Merge rust-bitcoin/rust-bitcoin#3356: Fix GetKey for sets (plus some related changes)

b593c886e3 Support GetKey where the Xpriv is a direct child of the looked up KeySource (Nadav Ivgi)
055aa9d4dc Refactor GetKey to take the KeyRequest by reference (Nadav Ivgi)
d15c57bd1f Refactor GetKey for sets to internally use Xpriv::get_key() (Nadav Ivgi)
d25c62bf45 Fix GetKey for sets to properly compare the fingerprint (Nadav Ivgi)

Pull request description:

  - The first commit is the simplest fix for a bug where the fingerprint wasn't compared correctly.

  - The second & third commits are optional refactoring to reuse `Xpriv::get_key` for `$set<Xpriv>::get_key`, so the Xpriv matching logic only has to be maintained in one place.

  - The forth commit adds support for signing with `Xpriv`s that are direct children of the `KeySource` -- possibly what the original (buggy) code author had in mind?

  Of course, feel free to take just the first commit if the others seem unnecessary. The last one is kind of meh, not sure if really useful.

  Note that multi-`Xpriv` signing does not actually work until #2850 is addressed too.

ACKs for top commit:
  apoelstra:
    ACK b593c886e3 successfully ran local tests
  tcharding:
    ACK b593c886e3

Tree-SHA512: a6f84fe12b68ddb71fc722649b1596d1561ab08a07d47f2345bd8932e155571f2b1e23c32a9c03d588fb697540340cb8ca38275dfe11f4910fd48265a993762a

bitcoin/src/psbt/mod.rs

@@ -20,7 +20,7 @@ use std::collections::{HashMap, HashSet};
 use internals::write_err;
 use secp256k1::{Keypair, Message, Secp256k1, Signing, Verification};
 
-use crate::bip32::{self, KeySource, Xpriv, Xpub};
+use crate::bip32::{self, DerivationPath, KeySource, Xpriv, Xpub};
 use crate::crypto::key::{PrivateKey, PublicKey};
 use crate::crypto::{ecdsa, taproot};
 use crate::key::{TapTweak, XOnlyPublicKey};
@@ -365,9 +365,9 @@ impl Psbt {
         let mut used = vec![]; // List of pubkeys used to sign the input.
 
         for (pk, key_source) in input.bip32_derivation.iter() {
-            let sk = if let Ok(Some(sk)) = k.get_key(KeyRequest::Bip32(key_source.clone()), secp) {
+            let sk = if let Ok(Some(sk)) = k.get_key(&KeyRequest::Bip32(key_source.clone()), secp) {
                 sk
-            } else if let Ok(Some(sk)) = k.get_key(KeyRequest::Pubkey(PublicKey::new(*pk)), secp) {
+            } else if let Ok(Some(sk)) = k.get_key(&KeyRequest::Pubkey(PublicKey::new(*pk)), secp) {
                 sk
             } else {
                 continue;
@@ -419,7 +419,7 @@ impl Psbt {
 
         for (&xonly, (leaf_hashes, key_source)) in input.tap_key_origins.iter() {
             let sk = if let Ok(Some(secret_key)) =
-                k.get_key(KeyRequest::Bip32(key_source.clone()), secp)
+                k.get_key(&KeyRequest::Bip32(key_source.clone()), secp)
             {
                 secret_key
             } else {
@@ -745,7 +745,7 @@ pub trait GetKey {
     /// - `Err` if an error was encountered while looking for the key.
     fn get_key<C: Signing>(
         &self,
-        key_request: KeyRequest,
+        key_request: &KeyRequest,
         secp: &Secp256k1<C>,
     ) -> Result<Option<PrivateKey>, Self::Error>;
 }
@@ -755,13 +755,20 @@ impl GetKey for Xpriv {
 
     fn get_key<C: Signing>(
         &self,
-        key_request: KeyRequest,
+        key_request: &KeyRequest,
         secp: &Secp256k1<C>,
     ) -> Result<Option<PrivateKey>, Self::Error> {
         match key_request {
             KeyRequest::Pubkey(_) => Err(GetKeyError::NotSupported),
             KeyRequest::Bip32((fingerprint, path)) => {
-                let key = if self.fingerprint(secp) == fingerprint {
+                let key = if self.fingerprint(secp) == *fingerprint {
+                    let k = self.derive_priv(secp, &path);
+                    Some(k.to_priv())
+                } else if self.parent_fingerprint == *fingerprint
+                    && !path.is_empty()
+                    && path[0] == self.child_number
+                {
+                    let path = DerivationPath::from_iter(path.into_iter().skip(1).copied());
                     let k = self.derive_priv(secp, &path);
                     Some(k.to_priv())
                 } else {
@@ -800,21 +807,14 @@ impl GetKey for $set<Xpriv> {
 
     fn get_key<C: Signing>(
         &self,
-        key_request: KeyRequest,
+        key_request: &KeyRequest,
         secp: &Secp256k1<C>
     ) -> Result<Option<PrivateKey>, Self::Error> {
-        match key_request {
+        // OK to stop at the first error because Xpriv::get_key() can only fail
-            KeyRequest::Pubkey(_) => Err(GetKeyError::NotSupported),
+        // if this isn't a KeyRequest::Bip32, which would fail for all Xprivs.
-            KeyRequest::Bip32((fingerprint, path)) => {
+        self.iter()
-                for xpriv in self.iter() {
+            .find_map(|xpriv| xpriv.get_key(key_request, secp).transpose())
-                    if xpriv.parent_fingerprint == fingerprint {
+            .transpose()
-                        let k = xpriv.derive_priv(secp, &path);
-                        return Ok(Some(k.to_priv()));
-                    }
-                }
-                Ok(None)
-            }
-        }
     }
 }}}
 impl_get_key_for_set!(BTreeSet);
@@ -830,7 +830,7 @@ impl GetKey for $map<PublicKey, PrivateKey> {
 
     fn get_key<C: Signing>(
         &self,
-        key_request: KeyRequest,
+        key_request: &KeyRequest,
         _: &Secp256k1<C>,
     ) -> Result<Option<PrivateKey>, Self::Error> {
         match key_request {
@@ -2133,7 +2133,7 @@ mod tests {
         let mut key_map = BTreeMap::new();
         key_map.insert(pk, priv_key);
 
-        let got = key_map.get_key(KeyRequest::Pubkey(pk), &secp).expect("failed to get key");
+        let got = key_map.get_key(&KeyRequest::Pubkey(pk), &secp).expect("failed to get key");
         assert_eq!(got.unwrap(), priv_key)
     }
 

bitcoin/tests/psbt-sign-taproot.rs

@@ -27,12 +27,12 @@ fn psbt_sign_taproot() {
         type Error = SignError;
         fn get_key<C: Signing>(
             &self,
-            key_request: KeyRequest,
+            key_request: &KeyRequest,
             _secp: &Secp256k1<C>,
         ) -> Result<Option<PrivateKey>, Self::Error> {
             match key_request {
                 KeyRequest::Bip32((mfp, _)) =>
-                    if mfp == self.mfp {
+                    if *mfp == self.mfp {
                         Ok(Some(self.sk))
                     } else {
                         Err(SignError::KeyNotFound)