diff --git a/src/signer.rs b/src/signer.rs index 43fe05f..a8b9a8d 100644 --- a/src/signer.rs +++ b/src/signer.rs @@ -1,4 +1,6 @@ use alloc::string::ToString; +use alloc::vec::Vec; +use std::collections::BTreeMap; use bitcoin::{ psbt::{GetKey, GetKeyError, KeyRequest}, @@ -24,9 +26,10 @@ impl GetKey for Signer { for entry in &self.0 { match entry { (_, DescriptorSecretKey::Single(prv)) => { - let pk = prv.key.public_key(secp); - if key_request == KeyRequest::Pubkey(pk) { - return Ok(Some(prv.key)); + let map: BTreeMap<_, _> = + core::iter::once((prv.key.public_key(secp), prv.key)).collect(); + if let Ok(Some(prv)) = GetKey::get_key(&map, key_request.clone(), secp) { + return Ok(Some(prv)); } } (_, desc_sk) => { @@ -44,7 +47,11 @@ impl GetKey for Signer { if fingerprint == fp && derivation.to_string().starts_with(&path.to_string()) { - let to_derive = &derivation[k.xkey.depth as usize..]; + let to_derive = derivation + .into_iter() + .skip(path.len()) + .cloned() + .collect::>(); let derived = k.xkey.derive_priv(secp, &to_derive)?; return Ok(Some(derived.to_priv())); } @@ -61,6 +68,8 @@ impl GetKey for Signer { #[cfg(test)] mod test { + use crate::bitcoin::bip32::ChildNumber; + use core::str::FromStr; use std::string::String; use bitcoin::bip32::{DerivationPath, Xpriv}; @@ -89,6 +98,26 @@ mod test { Ok(()) } + #[test] + fn get_key_x_only_pubkey() -> anyhow::Result<()> { + let secp = Secp256k1::new(); + let wif = "cU6BxEezV8FnkEPBCaFtc4WNuUKmgFaAu6sJErB154GXgMUjhgWe"; + let prv = bitcoin::PrivateKey::from_wif(wif)?; + let (x_only_pk, _parity) = prv.inner.x_only_public_key(&secp); + + let s = format!("wpkh({wif})"); + let (_, keymap) = Descriptor::parse_descriptor(&secp, &s).unwrap(); + + let signer = Signer(keymap); + let req = KeyRequest::XOnlyPubkey(x_only_pk); + let res = signer.get_key(req, &secp); + assert!(matches!( + res, + Ok(Some(k)) if k.inner.x_only_public_key(&secp).0 == x_only_pk + )); + Ok(()) + } + // Test `Signer` can fulfill a bip32 KeyRequest if we know the key origin #[test] fn get_key_bip32() -> anyhow::Result<()> { @@ -141,4 +170,35 @@ mod test { Ok(()) } + + #[test] + fn get_key_xpriv_with_key_origin() -> anyhow::Result<()> { + let secp = Secp256k1::new(); + let s = "wpkh([d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)"; + let (_, keymap) = Descriptor::parse_descriptor(&secp, s)?; + + let desc_sk = DescriptorSecretKey::from_str("[d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*")?; + let desc_xkey = match desc_sk { + DescriptorSecretKey::XPrv(k) => k, + _ => panic!(), + }; + + let (fp, _) = desc_xkey.origin.clone().unwrap(); + let path = DerivationPath::from_str("84h/1h/0h/7")?; + let req = KeyRequest::Bip32((fp, path)); + + let exp_prv = desc_xkey + .xkey + .derive_priv(&secp, &[ChildNumber::from(7)])? + .to_priv(); + + let res = Signer(keymap).get_key(req, &secp); + + assert!(matches!( + res, + Ok(Some(k)) if k == exp_prv, + )); + + Ok(()) + } }