From d25c62bf4572dc76f00b4bf89f48cff1a01db04a Mon Sep 17 00:00:00 2001 From: Nadav Ivgi Date: Fri, 13 Sep 2024 11:28:17 +0300 Subject: [PATCH 1/4] Fix GetKey for sets to properly compare the fingerprint --- bitcoin/src/psbt/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bitcoin/src/psbt/mod.rs b/bitcoin/src/psbt/mod.rs index da9cd3efd..09799f95e 100644 --- a/bitcoin/src/psbt/mod.rs +++ b/bitcoin/src/psbt/mod.rs @@ -807,7 +807,7 @@ impl GetKey for $set { KeyRequest::Pubkey(_) => Err(GetKeyError::NotSupported), KeyRequest::Bip32((fingerprint, path)) => { for xpriv in self.iter() { - if xpriv.parent_fingerprint == fingerprint { + if xpriv.fingerprint(secp) == fingerprint { let k = xpriv.derive_priv(secp, &path); return Ok(Some(k.to_priv())); } From d15c57bd1f4a3928ae5bfa23118c6d28f21b6b9d Mon Sep 17 00:00:00 2001 From: Nadav Ivgi Date: Fri, 13 Sep 2024 11:35:32 +0300 Subject: [PATCH 2/4] Refactor GetKey for sets to internally use Xpriv::get_key() --- bitcoin/src/psbt/mod.rs | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/bitcoin/src/psbt/mod.rs b/bitcoin/src/psbt/mod.rs index 09799f95e..af39d22fb 100644 --- a/bitcoin/src/psbt/mod.rs +++ b/bitcoin/src/psbt/mod.rs @@ -803,18 +803,11 @@ impl GetKey for $set { key_request: KeyRequest, secp: &Secp256k1 ) -> Result, Self::Error> { - match key_request { - KeyRequest::Pubkey(_) => Err(GetKeyError::NotSupported), - KeyRequest::Bip32((fingerprint, path)) => { - for xpriv in self.iter() { - if xpriv.fingerprint(secp) == fingerprint { - let k = xpriv.derive_priv(secp, &path); - return Ok(Some(k.to_priv())); - } - } - Ok(None) - } - } + // OK to stop at the first error because Xpriv::get_key() can only fail + // if this isn't a KeyRequest::Bip32, which would fail for all Xprivs. + self.iter() + .find_map(|xpriv| xpriv.get_key(key_request.clone(), secp).transpose()) + .transpose() } }}} impl_get_key_for_set!(BTreeSet); From 055aa9d4dcc500183be087189363ac3270df9be9 Mon Sep 17 00:00:00 2001 From: Nadav Ivgi Date: Fri, 13 Sep 2024 11:38:38 +0300 Subject: [PATCH 3/4] Refactor GetKey to take the KeyRequest by reference To avoid cloning when looking it up in sets. --- bitcoin/src/psbt/mod.rs | 20 ++++++++++---------- bitcoin/tests/psbt-sign-taproot.rs | 4 ++-- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/bitcoin/src/psbt/mod.rs b/bitcoin/src/psbt/mod.rs index af39d22fb..472847655 100644 --- a/bitcoin/src/psbt/mod.rs +++ b/bitcoin/src/psbt/mod.rs @@ -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( &self, - key_request: KeyRequest, + key_request: &KeyRequest, secp: &Secp256k1, ) -> Result, Self::Error>; } @@ -755,13 +755,13 @@ impl GetKey for Xpriv { fn get_key( &self, - key_request: KeyRequest, + key_request: &KeyRequest, secp: &Secp256k1, ) -> Result, 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 { @@ -800,13 +800,13 @@ impl GetKey for $set { fn get_key( &self, - key_request: KeyRequest, + key_request: &KeyRequest, secp: &Secp256k1 ) -> Result, Self::Error> { // OK to stop at the first error because Xpriv::get_key() can only fail // if this isn't a KeyRequest::Bip32, which would fail for all Xprivs. self.iter() - .find_map(|xpriv| xpriv.get_key(key_request.clone(), secp).transpose()) + .find_map(|xpriv| xpriv.get_key(key_request, secp).transpose()) .transpose() } }}} @@ -823,7 +823,7 @@ impl GetKey for $map { fn get_key( &self, - key_request: KeyRequest, + key_request: &KeyRequest, _: &Secp256k1, ) -> Result, Self::Error> { match key_request { @@ -2126,7 +2126,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) } diff --git a/bitcoin/tests/psbt-sign-taproot.rs b/bitcoin/tests/psbt-sign-taproot.rs index e79a1fe1a..6f4563988 100644 --- a/bitcoin/tests/psbt-sign-taproot.rs +++ b/bitcoin/tests/psbt-sign-taproot.rs @@ -27,12 +27,12 @@ fn psbt_sign_taproot() { type Error = SignError; fn get_key( &self, - key_request: KeyRequest, + key_request: &KeyRequest, _secp: &Secp256k1, ) -> Result, Self::Error> { match key_request { KeyRequest::Bip32((mfp, _)) => - if mfp == self.mfp { + if *mfp == self.mfp { Ok(Some(self.sk)) } else { Err(SignError::KeyNotFound) From b593c886e34e89b7c628e9b34de84a0e20e805af Mon Sep 17 00:00:00 2001 From: Nadav Ivgi Date: Fri, 13 Sep 2024 11:46:07 +0300 Subject: [PATCH 4/4] Support GetKey where the Xpriv is a direct child of the looked up KeySource --- bitcoin/src/psbt/mod.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/bitcoin/src/psbt/mod.rs b/bitcoin/src/psbt/mod.rs index 472847655..dd27611af 100644 --- a/bitcoin/src/psbt/mod.rs +++ b/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}; @@ -764,6 +764,13 @@ impl GetKey for Xpriv { 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 { None };