feat: update signer impl for XOnlyPubkey

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::<Vec<_>>();
                                         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(())
+    }
 }