feat(descriptor): backport GetKey impl from #851

Cargo-minimal.lock

@@ -53,9 +53,9 @@ checksum = "d965446196e3b7decd44aa7ee49e31d630118f90ef12f97900f262eb915c951d"
 
 [[package]]
 name = "bitcoin"
-version = "0.32.0"
+version = "0.32.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7170e7750a20974246f17ece04311b4205a6155f1db564c5b224af817663c3ea"
+checksum = "ad8929a18b8e33ea6b3c09297b687baaa71fb1b97353243a3f1029fad5c59c5b"
 dependencies = [
  "base58ck",
  "base64 0.21.7",

Cargo-recent.lock

@@ -53,9 +53,9 @@ checksum = "d965446196e3b7decd44aa7ee49e31d630118f90ef12f97900f262eb915c951d"
 
 [[package]]
 name = "bitcoin"
-version = "0.32.0"
+version = "0.32.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7170e7750a20974246f17ece04311b4205a6155f1db564c5b224af817663c3ea"
+checksum = "ad8929a18b8e33ea6b3c09297b687baaa71fb1b97353243a3f1029fad5c59c5b"
 dependencies = [
  "base58ck",
  "base64 0.21.7",

Cargo.toml

@@ -23,7 +23,7 @@ base64 = ["bitcoin/base64"]
 
 [dependencies]
 bech32 = { version = "0.11.0", default-features = false }
-bitcoin = { version = "0.32.0", default-features = false }
+bitcoin = { version = "0.32.6", default-features = false }
 
 # Do NOT use this as a feature! Use the `serde` feature instead.
 actual-serde = { package = "serde", version = "1.0.103", optional = true }

src/descriptor/key_map.rs

@@ -0,0 +1,361 @@
+use bitcoin::psbt::{GetKey, GetKeyError, KeyRequest};
+use bitcoin::secp256k1::{Secp256k1, Signing};
+use bitcoin::PrivateKey;
+
+use crate::descriptor::{DescriptorSecretKey, KeyMap};
+use crate::BTreeMap;
+
+/// A wrapper around KeyMap that implements GetKey for PSBT signing.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct KeyMapWrapper {
+    map: KeyMap,
+}
+
+impl From<KeyMap> for KeyMapWrapper {
+    fn from(map: KeyMap) -> Self { KeyMapWrapper { map } }
+}
+
+impl GetKey for KeyMapWrapper {
+    type Error = GetKeyError;
+
+    fn get_key<C: Signing>(
+        &self,
+        key_request: KeyRequest,
+        secp: &Secp256k1<C>,
+    ) -> Result<Option<bitcoin::PrivateKey>, Self::Error> {
+        Ok(self
+            .map
+            .iter()
+            .find_map(|(_desc_pk, desc_sk)| -> Option<PrivateKey> {
+                match desc_sk.get_key(key_request.clone(), secp) {
+                    Ok(Some(pk)) => Some(pk),
+                    // When looking up keys in a map, we eat errors on individual keys, on
+                    // the assumption that some other key in the map might not error.
+                    Ok(None) | Err(_) => None,
+                }
+            }))
+    }
+}
+
+impl GetKey for DescriptorSecretKey {
+    type Error = GetKeyError;
+
+    fn get_key<C: Signing>(
+        &self,
+        key_request: KeyRequest,
+        secp: &Secp256k1<C>,
+    ) -> Result<Option<PrivateKey>, Self::Error> {
+        match (self, key_request) {
+            (DescriptorSecretKey::Single(single_priv), key_request) => {
+                let sk = single_priv.key;
+                let pk = sk.public_key(secp);
+                let pubkey_map = BTreeMap::from([(pk, sk)]);
+                pubkey_map.get_key(key_request, secp)
+            }
+            (DescriptorSecretKey::XPrv(descriptor_xkey), KeyRequest::Pubkey(public_key)) => {
+                let xpriv = descriptor_xkey
+                    .xkey
+                    .derive_priv(secp, &descriptor_xkey.derivation_path)
+                    .map_err(GetKeyError::Bip32)?;
+                let pk = xpriv.private_key.public_key(secp);
+
+                if public_key.inner.eq(&pk) {
+                    Ok(Some(xpriv.to_priv()))
+                } else {
+                    Ok(None)
+                }
+            }
+            (
+                DescriptorSecretKey::XPrv(descriptor_xkey),
+                ref key_request @ KeyRequest::Bip32(ref key_source),
+            ) => {
+                if let Some(key) = descriptor_xkey.xkey.get_key(key_request.clone(), secp)? {
+                    return Ok(Some(key));
+                }
+
+                if descriptor_xkey.matches(key_source, secp).is_some() {
+                    let (_, derivation_path) = key_source;
+                    return Ok(Some(
+                        descriptor_xkey
+                            .xkey
+                            .derive_priv(secp, &derivation_path)
+                            .map_err(GetKeyError::Bip32)?
+                            .to_priv(),
+                    ));
+                }
+
+                Ok(None)
+            }
+            (DescriptorSecretKey::XPrv(_), KeyRequest::XOnlyPubkey(_)) => {
+                Err(GetKeyError::NotSupported)
+            }
+            (
+                desc_multi_sk @ DescriptorSecretKey::MultiXPrv(_descriptor_multi_xkey),
+                key_request,
+            ) => {
+                for desc_sk in &desc_multi_sk.clone().into_single_keys() {
+                    // If any key is an error, then all of them will, so here we propagate errors with ?.
+                    if let Some(pk) = desc_sk.get_key(key_request.clone(), secp)? {
+                        return Ok(Some(pk));
+                    }
+                }
+                Ok(None)
+            }
+            _ => Ok(None),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use core::str::FromStr;
+
+    use bitcoin::bip32::{ChildNumber, DerivationPath, IntoDerivationPath, Xpriv};
+
+    use super::*;
+    use crate::Descriptor;
+
+    #[test]
+    fn get_key_single_key() {
+        let secp = Secp256k1::new();
+
+        let descriptor_sk_s =
+            "[90b6a706/44'/0'/0'/0/0]cMk8gWmj1KpjdYnAWwsEDekodMYhbyYBhG8gMtCCxucJ98JzcNij";
+
+        let single = match descriptor_sk_s.parse::<DescriptorSecretKey>().unwrap() {
+            DescriptorSecretKey::Single(single) => single,
+            _ => panic!("unexpected DescriptorSecretKey variant"),
+        };
+
+        let want_sk = single.key;
+        let descriptor_s = format!("wpkh({})", descriptor_sk_s);
+        let (_, keymap) = Descriptor::parse_descriptor(&secp, &descriptor_s).unwrap();
+        let keymap_wrapper = KeyMapWrapper::from(keymap);
+
+        let pk = want_sk.public_key(&secp);
+        let request = KeyRequest::Pubkey(pk);
+        let got_sk = keymap_wrapper
+            .get_key(request, &secp)
+            .expect("get_key call errored")
+            .expect("failed to find the key");
+        assert_eq!(got_sk, want_sk)
+    }
+
+    #[test]
+    fn get_key_xpriv_single_key_xpriv() {
+        let secp = Secp256k1::new();
+
+        let s = "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi";
+
+        let xpriv = s.parse::<Xpriv>().unwrap();
+        let xpriv_fingerprint = xpriv.fingerprint(&secp);
+
+        // Sanity check.
+        {
+            let descriptor_sk_s = format!("[{}]{}", xpriv_fingerprint, xpriv);
+            let descriptor_sk = descriptor_sk_s.parse::<DescriptorSecretKey>().unwrap();
+            let got = match descriptor_sk {
+                DescriptorSecretKey::XPrv(x) => x.xkey,
+                _ => panic!("unexpected DescriptorSecretKey variant"),
+            };
+            assert_eq!(got, xpriv);
+        }
+
+        let want_sk = xpriv.to_priv();
+        let descriptor_s = format!("wpkh([{}]{})", xpriv_fingerprint, xpriv);
+        let (_, keymap) = Descriptor::parse_descriptor(&secp, &descriptor_s).unwrap();
+        let keymap_wrapper = KeyMapWrapper::from(keymap);
+
+        let pk = want_sk.public_key(&secp);
+        let request = KeyRequest::Pubkey(pk);
+        let got_sk = keymap_wrapper
+            .get_key(request, &secp)
+            .expect("get_key call errored")
+            .expect("failed to find the key");
+        assert_eq!(got_sk, want_sk)
+    }
+
+    #[test]
+    fn get_key_xpriv_child_depth_one() {
+        let secp = Secp256k1::new();
+
+        let s = "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi";
+        let master = s.parse::<Xpriv>().unwrap();
+        let master_fingerprint = master.fingerprint(&secp);
+
+        let child_number = ChildNumber::from_hardened_idx(44).unwrap();
+        let child = master.derive_priv(&secp, &[child_number]).unwrap();
+
+        // Sanity check.
+        {
+            let descriptor_sk_s = format!("[{}/44']{}", master_fingerprint, child);
+            let descriptor_sk = descriptor_sk_s.parse::<DescriptorSecretKey>().unwrap();
+            let got = match descriptor_sk {
+                DescriptorSecretKey::XPrv(ref x) => x.xkey,
+                _ => panic!("unexpected DescriptorSecretKey variant"),
+            };
+            assert_eq!(got, child);
+        }
+
+        let want_sk = child.to_priv();
+        let descriptor_s = format!("wpkh({}/44')", s);
+        let (_, keymap) = Descriptor::parse_descriptor(&secp, &descriptor_s).unwrap();
+        let keymap_wrapper = KeyMapWrapper::from(keymap);
+
+        let pk = want_sk.public_key(&secp);
+        let request = KeyRequest::Pubkey(pk);
+        let got_sk = keymap_wrapper
+            .get_key(request, &secp)
+            .expect("get_key call errored")
+            .expect("failed to find the key");
+        assert_eq!(got_sk, want_sk)
+    }
+
+    #[test]
+    fn get_key_xpriv_with_path() {
+        let secp = Secp256k1::new();
+
+        let s = "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi";
+        let master = s.parse::<Xpriv>().unwrap();
+        let master_fingerprint = master.fingerprint(&secp);
+
+        let first_external_child = "44'/0'/0'/0/0";
+        let derivation_path = first_external_child.into_derivation_path().unwrap();
+
+        let child = master.derive_priv(&secp, &derivation_path).unwrap();
+
+        // Sanity check.
+        {
+            let descriptor_sk_s =
+                format!("[{}/{}]{}", master_fingerprint, first_external_child, child);
+            let descriptor_sk = descriptor_sk_s.parse::<DescriptorSecretKey>().unwrap();
+            let got = match descriptor_sk {
+                DescriptorSecretKey::XPrv(ref x) => x.xkey,
+                _ => panic!("unexpected DescriptorSecretKey variant"),
+            };
+            assert_eq!(got, child);
+        }
+
+        let want_sk = child.to_priv();
+        let descriptor_s = format!("wpkh({}/44'/0'/0'/0/*)", s);
+        let (_, keymap) = Descriptor::parse_descriptor(&secp, &descriptor_s).unwrap();
+        let keymap_wrapper = KeyMapWrapper::from(keymap);
+
+        let key_source = (master_fingerprint, derivation_path);
+        let request = KeyRequest::Bip32(key_source);
+        let got_sk = keymap_wrapper
+            .get_key(request, &secp)
+            .expect("get_key call errored")
+            .expect("failed to find the key");
+
+        assert_eq!(got_sk, want_sk)
+    }
+
+    #[test]
+    fn get_key_xpriv_with_key_origin() {
+        let secp = Secp256k1::new();
+
+        let descriptor_str = "wpkh([d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)";
+        let (_descriptor_pk, keymap) = Descriptor::parse_descriptor(&secp, descriptor_str).unwrap();
+        let keymap_wrapper = KeyMapWrapper::from(keymap);
+
+        let descriptor_sk = DescriptorSecretKey::from_str("[d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*").unwrap();
+        let xpriv = match descriptor_sk {
+            DescriptorSecretKey::XPrv(descriptor_xkey) => descriptor_xkey,
+            _ => unreachable!(),
+        };
+
+        let path = DerivationPath::from_str("84'/1'/0'/0").unwrap();
+        let expected_pk = xpriv.xkey.derive_priv(&secp, &path).unwrap().to_priv();
+
+        let (fp, _) = xpriv.origin.unwrap();
+        let key_request = KeyRequest::Bip32((fp, path));
+
+        let pk = keymap_wrapper
+            .get_key(key_request, &secp)
+            .expect("get_key should not fail")
+            .expect("get_key should return a `PrivateKey`");
+
+        assert_eq!(pk, expected_pk);
+    }
+
+    #[test]
+    fn get_key_keymap_no_match() {
+        let secp = Secp256k1::new();
+
+        // Create a keymap with one key
+        let descriptor_s = "wpkh(cMk8gWmj1KpjdYnAWwsEDekodMYhbyYBhG8gMtCCxucJ98JzcNij)";
+        let (_, keymap) = Descriptor::parse_descriptor(&secp, descriptor_s).unwrap();
+        let keymap_wrapper = KeyMapWrapper::from(keymap);
+
+        // Request a different public key that doesn't exist in the keymap
+        let different_sk =
+            PrivateKey::from_str("cNJFgo1driFnPcBdBX8BrJrpxchBWXwXCvNH5SoSkdcF6JXXwHMm").unwrap();
+        let different_pk = different_sk.public_key(&secp);
+        let request = KeyRequest::Pubkey(different_pk);
+
+        let result = keymap_wrapper.get_key(request, &secp).unwrap();
+        assert!(result.is_none(), "Should return None when no matching key is found");
+    }
+
+    #[test]
+    fn get_key_descriptor_secret_key_xonly_not_supported() {
+        let secp = Secp256k1::new();
+
+        let descriptor_sk = DescriptorSecretKey::from_str("xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi").unwrap();
+
+        // Create an x-only public key request
+        let sk =
+            PrivateKey::from_str("cMk8gWmj1KpjdYnAWwsEDekodMYhbyYBhG8gMtCCxucJ98JzcNij").unwrap();
+        let xonly_pk = sk.public_key(&secp).inner.x_only_public_key().0;
+        let request = KeyRequest::XOnlyPubkey(xonly_pk);
+
+        let result = descriptor_sk.get_key(request.clone(), &secp);
+        assert!(matches!(result, Err(GetKeyError::NotSupported)));
+
+        // Also test with KeyMap
+        let descriptor_s = "wpkh(xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi)";
+        let (_, keymap) = Descriptor::parse_descriptor(&secp, descriptor_s).unwrap();
+        let keymap_wrapper = KeyMapWrapper::from(keymap);
+
+        // While requesting an x-only key from an individual xpriv, that's an error.
+        // But from a keymap, which might have both x-only keys and regular xprivs,
+        // we treat errors as "key not found".
+        let result = keymap_wrapper.get_key(request, &secp);
+        assert!(matches!(result, Ok(None)));
+    }
+
+    #[test]
+    fn get_key_descriptor_secret_key_xonly_multipath() {
+        let secp = Secp256k1::new();
+
+        let descriptor_sk = DescriptorSecretKey::from_str("[d34db33f/84h/0h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/<0;1>").unwrap();
+
+        // Request with a different fingerprint
+        let different_fingerprint = bitcoin::bip32::Fingerprint::from([0x12, 0x34, 0x56, 0x78]);
+        let path = DerivationPath::from_str("84'/1'/0'/0").unwrap();
+        let request = KeyRequest::Bip32((different_fingerprint, path));
+
+        let result = descriptor_sk.get_key(request.clone(), &secp).unwrap();
+        assert!(result.is_none(), "Should return None when fingerprint doesn't match");
+
+        // Create an x-only public key request -- now we get "not supported".
+        let sk =
+            PrivateKey::from_str("cMk8gWmj1KpjdYnAWwsEDekodMYhbyYBhG8gMtCCxucJ98JzcNij").unwrap();
+        let xonly_pk = sk.public_key(&secp).inner.x_only_public_key().0;
+        let request_x = KeyRequest::XOnlyPubkey(xonly_pk);
+
+        let result = descriptor_sk.get_key(request_x.clone(), &secp);
+        assert!(matches!(result, Err(GetKeyError::NotSupported)));
+
+        // Also test with KeyMap; as in the previous test, the error turns to None.
+        let descriptor_s = "wpkh([d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)";
+        let (_, keymap) = Descriptor::parse_descriptor(&secp, descriptor_s).unwrap();
+        let keymap_wrapper = KeyMapWrapper::from(keymap);
+
+        let result = keymap_wrapper.get_key(request, &secp).unwrap();
+        assert!(result.is_none(), "Should return None when fingerprint doesn't match");
+        let result = keymap_wrapper.get_key(request_x, &secp).unwrap();
+        assert!(result.is_none(), "Should return None even on error");
+    }
+}