Add XOnlyPublicKey support for PSBT key retrieval and improve Taproot signing

erickcestari · erickcestari · commit 069d2fd07e7d · 2025-03-21T11:28:54.000-03:00
This commit enhances PSBT signing functionality by: 1. Added new KeyRequest::XOnlyPubkey variant to support direct retrieval using XOnly public keys 2. Implemented GetKey for HashMap<XOnlyPublicKey, PrivateKey> for more efficient Taproot key management 3. Modified HashMap<PublicKey, PrivateKey> implementation to handle XOnlyPublicKey requests by checking both even and odd parity variants These changes allow for more flexible key management in Taproot transactions. Specifically, wallet implementations can now store keys indexed by either PublicKey or XOnlyPublicKey and successfully sign PSBTs with Taproot inputs. Added tests for both implementations to verify correct behavior. Added test for odd parity key retrieval. Closes rust-bitcoin#4150
diff --git a/bitcoin/src/crypto/key.rs b/bitcoin/src/crypto/key.rs
@@ -533,6 +533,20 @@ impl PrivateKey {
             inner: secp256k1::SecretKey::from_byte_array(key)?,
         })
     }
+
+    /// Returns a new private key with the negated secret value.
+    ///
+    /// The resulting key corresponds to the same x-only public key (identical x-coordinate)
+    /// but with the opposite y-coordinate parity. This is useful for ensuring compatibility 
+    /// with specific public key formats and BIP-340 requirements.
+    #[inline]
+    pub fn negate(&self) -> Self {
+        PrivateKey {
+            compressed: self.compressed,
+            network: self.network,
+            inner: self.inner.negate(),
+        }
+    }
 }
 
 impl fmt::Display for PrivateKey {
diff --git a/bitcoin/src/psbt/mod.rs b/bitcoin/src/psbt/mod.rs
@@ -427,6 +427,8 @@ impl Psbt {
                 k.get_key(&KeyRequest::Bip32(key_source.clone()), secp)
             {
                 secret_key
+            } else if let Ok(Some(sk)) = k.get_key(&KeyRequest::XOnlyPubkey(xonly), secp) {
+                sk
             } else {
                 continue;
             };
@@ -737,6 +739,8 @@ pub enum KeyRequest {
     Pubkey(PublicKey),
     /// Request a private key using BIP-32 fingerprint and derivation path.
     Bip32(KeySource),
+    /// Request a private key using the associated x-only public key.
+    XOnlyPubkey(XOnlyPublicKey),
 }
 
 /// Trait to get a private key from a key request, key is then used to sign an input.
@@ -768,6 +772,7 @@ impl GetKey for Xpriv {
     ) -> Result<Option<PrivateKey>, Self::Error> {
         match key_request {
             KeyRequest::Pubkey(_) => Err(GetKeyError::NotSupported),
+            KeyRequest::XOnlyPubkey(_) => Err(GetKeyError::NotSupported),
             KeyRequest::Bip32((fingerprint, path)) => {
                 let key = if self.fingerprint(secp) == *fingerprint {
                     let k = self.derive_xpriv(secp, &path);
@@ -831,7 +836,7 @@ impl_get_key_for_set!(BTreeSet);
 impl_get_key_for_set!(HashSet);
 
 #[rustfmt::skip]
-macro_rules! impl_get_key_for_map {
+macro_rules! impl_get_key_for_pubkey_map {
     ($map:ident) => {
 
 impl GetKey for $map<PublicKey, PrivateKey> {
@@ -844,13 +849,67 @@ impl GetKey for $map<PublicKey, PrivateKey> {
     ) -> Result<Option<PrivateKey>, Self::Error> {
         match key_request {
             KeyRequest::Pubkey(pk) => Ok(self.get(&pk).cloned()),
+            KeyRequest::XOnlyPubkey(xonly) => {
+                let pubkey_even = PublicKey::new(xonly.public_key(secp256k1::Parity::Even));
+                let key = self.get(&pubkey_even).cloned();
+                
+                if key.is_some() {
+                    return Ok(key);
+                }
+                
+                let pubkey_odd = PublicKey::new(xonly.public_key(secp256k1::Parity::Odd));
+                if let Some(priv_key) = self.get(&pubkey_odd).copied() {
+                    let negated_priv_key  = priv_key.negate();
+                    return Ok(Some(negated_priv_key));
+                }
+
+                Ok(None)
+            },
+            KeyRequest::Bip32(_) => Err(GetKeyError::NotSupported),
+        }
+    }
+}}}
+impl_get_key_for_pubkey_map!(BTreeMap);
+#[cfg(feature = "std")]
+impl_get_key_for_pubkey_map!(HashMap);
+
+#[rustfmt::skip]
+macro_rules! impl_get_key_for_xonly_map {
+    ($map:ident) => {
+
+impl GetKey for $map<XOnlyPublicKey, PrivateKey> {
+    type Error = GetKeyError;
+
+    fn get_key<C: Signing>(
+        &self,
+        key_request: &KeyRequest,
+        secp: &Secp256k1<C>,
+    ) -> Result<Option<PrivateKey>, Self::Error> {
+        match key_request {
+            KeyRequest::XOnlyPubkey(xonly) => Ok(self.get(xonly).cloned()),
+            KeyRequest::Pubkey(pk) => {
+                let (xonly, parity) = pk.inner.x_only_public_key();
+                
+                if let Some(mut priv_key) = self.get(&XOnlyPublicKey::from(xonly)).cloned() {
+                    let computed_pk = priv_key.public_key(&secp);
+                    let (_, computed_parity) = computed_pk.inner.x_only_public_key();
+                    
+                    if computed_parity != parity {
+                        priv_key = priv_key.negate();
+                    }
+                    
+                    return Ok(Some(priv_key));
+                }
+                
+                Ok(None)
+            },
             KeyRequest::Bip32(_) => Err(GetKeyError::NotSupported),
         }
     }
 }}}
-impl_get_key_for_map!(BTreeMap);
+impl_get_key_for_xonly_map!(BTreeMap);
 #[cfg(feature = "std")]
-impl_get_key_for_map!(HashMap);
+impl_get_key_for_xonly_map!(HashMap);
 
 /// Errors when getting a key.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -1228,7 +1287,14 @@ mod tests {
     use hashes::{hash160, ripemd160, sha256};
     use hex::{test_hex_unwrap as hex, FromHex};
     #[cfg(feature = "rand-std")]
-    use secp256k1::{All, SecretKey};
+    use {
+        crate::address::script_pubkey::ScriptBufExt as _,
+        crate::bip32::{DerivationPath, Fingerprint},
+        crate::locktime,
+        crate::witness_version::WitnessVersion,
+        crate::WitnessProgram,
+        secp256k1::{All, SecretKey},
+    };
 
     use super::*;
     use crate::address::script_pubkey::ScriptExt as _;
@@ -2169,6 +2235,42 @@ mod tests {
         assert_eq!(got.unwrap(), priv_key)
     }
 
+    #[test]
+    #[cfg(feature = "rand-std")]
+    fn pubkey_map_get_key_negates_odd_parity_keys() {
+        use crate::psbt::{GetKey, KeyRequest};
+
+        let (mut priv_key, mut pk, secp) = gen_keys();
+        let (xonly, parity) = pk.inner.x_only_public_key();
+
+        let mut pubkey_map: HashMap<PublicKey, PrivateKey> = HashMap::new();
+
+        if parity == secp256k1::Parity::Even {
+            priv_key = PrivateKey {
+                compressed: priv_key.compressed,
+                network: priv_key.network,
+                inner: priv_key.inner.negate(),
+            };
+            pk = priv_key.public_key(&secp);
+        }
+
+        pubkey_map.insert(pk, priv_key);
+
+        let req_result = pubkey_map.get_key(&KeyRequest::XOnlyPubkey(xonly), &secp).unwrap();
+
+        let retrieved_key = req_result.unwrap();
+
+        let retrieved_pub_key = retrieved_key.public_key(&secp);
+        let (retrieved_xonly, retrieved_parity) = retrieved_pub_key.inner.x_only_public_key();
+
+        assert_eq!(xonly, retrieved_xonly);
+        assert_eq!(
+            retrieved_parity,
+            secp256k1::Parity::Even,
+            "Key should be normalized to have even parity, even when original had odd parity"
+        );
+    }
+
     #[test]
     fn fee() {
         let output_0_val = Amount::from_sat_u32(99_999_699);
@@ -2273,12 +2375,73 @@ mod tests {
 
     #[test]
     #[cfg(feature = "rand-std")]
-    fn sign_psbt() {
-        use crate::address::script_pubkey::ScriptBufExt as _;
-        use crate::bip32::{DerivationPath, Fingerprint};
-        use crate::witness_version::WitnessVersion;
-        use crate::WitnessProgram;
+    fn hashmap_can_sign_taproot() {
+        let (priv_key, pk, secp) = gen_keys();
+        let internal_key: XOnlyPublicKey = pk.inner.into();
+
+        let tx = Transaction {
+            version: transaction::Version::TWO,
+            lock_time: locktime::absolute::LockTime::ZERO,
+            input: vec![TxIn::EMPTY_COINBASE],
+            output: vec![TxOut { value: Amount::ZERO, script_pubkey: ScriptBuf::new() }],
+        };
 
+        let mut psbt = Psbt::from_unsigned_tx(tx).unwrap();
+        psbt.inputs[0].tap_internal_key = Some(internal_key);
+        psbt.inputs[0].witness_utxo = Some(transaction::TxOut {
+            value: Amount::from_sat_u32(10),
+            script_pubkey: ScriptBuf::new_p2tr(&secp, internal_key, None),
+        });
+
+        let mut key_map: HashMap<PublicKey, PrivateKey> = HashMap::new();
+        key_map.insert(pk, priv_key);
+
+        let key_source = (Fingerprint::default(), DerivationPath::default());
+        let mut tap_key_origins = std::collections::BTreeMap::new();
+        tap_key_origins.insert(internal_key, (vec![], key_source));
+        psbt.inputs[0].tap_key_origins = tap_key_origins;
+
+        let signing_keys = psbt.sign(&key_map, &secp).unwrap();
+        assert_eq!(signing_keys.len(), 1);
+        assert_eq!(signing_keys[&0], SigningKeys::Schnorr(vec![internal_key]));
+    }
+
+    #[test]
+    #[cfg(feature = "rand-std")]
+    fn xonly_hashmap_can_sign_taproot() {
+        let (priv_key, pk, secp) = gen_keys();
+        let internal_key: XOnlyPublicKey = pk.inner.into();
+
+        let tx = Transaction {
+            version: transaction::Version::TWO,
+            lock_time: locktime::absolute::LockTime::ZERO,
+            input: vec![TxIn::EMPTY_COINBASE],
+            output: vec![TxOut { value: Amount::ZERO, script_pubkey: ScriptBuf::new() }],
+        };
+
+        let mut psbt = Psbt::from_unsigned_tx(tx).unwrap();
+        psbt.inputs[0].tap_internal_key = Some(internal_key);
+        psbt.inputs[0].witness_utxo = Some(transaction::TxOut {
+            value: Amount::from_sat_u32(10),
+            script_pubkey: ScriptBuf::new_p2tr(&secp, internal_key, None),
+        });
+
+        let mut xonly_key_map: HashMap<XOnlyPublicKey, PrivateKey> = HashMap::new();
+        xonly_key_map.insert(internal_key, priv_key);
+
+        let key_source = (Fingerprint::default(), DerivationPath::default());
+        let mut tap_key_origins = std::collections::BTreeMap::new();
+        tap_key_origins.insert(internal_key, (vec![], key_source));
+        psbt.inputs[0].tap_key_origins = tap_key_origins;
+
+        let signing_keys = psbt.sign(&xonly_key_map, &secp).unwrap();
+        assert_eq!(signing_keys.len(), 1);
+        assert_eq!(signing_keys[&0], SigningKeys::Schnorr(vec![internal_key]));
+    }
+
+    #[test]
+    #[cfg(feature = "rand-std")]
+    fn sign_psbt() {
         let unsigned_tx = Transaction {
             version: transaction::Version::TWO,
             lock_time: absolute::LockTime::ZERO,
