From 17a9068599c3a550a5c263512dcf82fb2e2894b3 Mon Sep 17 00:00:00 2001 From: "Tobin C. Harding" Date: Tue, 29 Oct 2024 12:19:35 +1100 Subject: [PATCH] Implement GetKey for KeyMap Create a `KeyMay` type to replace the current `BTreeMap` alias and implement `bitcoin::psbt::GetKey` for it. Close: #709 --- src/descriptor/key_map.rs | 277 ++++++++++++++++++++++++++++++++++++++ src/descriptor/mod.rs | 24 ++-- src/plan.rs | 2 +- 3 files changed, 288 insertions(+), 15 deletions(-) create mode 100644 src/descriptor/key_map.rs diff --git a/src/descriptor/key_map.rs b/src/descriptor/key_map.rs new file mode 100644 index 000000000..3191280a0 --- /dev/null +++ b/src/descriptor/key_map.rs @@ -0,0 +1,277 @@ +// SPDX-License-Identifier: CC0-1.0 + +//! A map of public key to secret key. + +use bitcoin::psbt::{GetKey, GetKeyError, KeyRequest}; +use bitcoin::secp256k1::{Secp256k1, Signing}; + +use super::{DescriptorPublicKey, DescriptorSecretKey, SinglePubKey}; +use crate::prelude::BTreeMap; + +/// Alias type for a map of public key to secret key +/// +/// This map is returned whenever a descriptor that contains secrets is parsed using +/// [`Descriptor::parse_descriptor`], since the descriptor will always only contain +/// public keys. This map allows looking up the corresponding secret key given a +/// public key from the descriptor. +#[derive(Debug, Eq, PartialEq, Clone)] +pub struct KeyMap { + /// The map. + pub map: BTreeMap +} + +impl KeyMap { + /// Creates a new empty `KeyMap`. + pub fn new() -> Self { Self{ map: BTreeMap::new()} } + + /// Inserts `key` and `value` into the map. + /// + /// See [`BTreeMap::insert`] for more information. + pub fn insert(&mut self, key: DescriptorPublicKey, value: DescriptorSecretKey) -> Option { + self.map.insert(key, value) + } + + /// Gets `key` if it exists. + /// + /// See [`BTreeMap::insert`] for more information. + pub fn get(&self, key: &DescriptorPublicKey) -> Option<&DescriptorSecretKey> { + self.map.get(key) + } + + /// Returns the number of items in this map. + /// + /// See [`BTreeMap::insert`] for more information. + pub fn len(&self) -> usize { self.map.len() } +} + +impl GetKey for KeyMap { + type Error = GetKeyError; + + fn get_key( + &self, + key_request: &KeyRequest, + secp: &Secp256k1 + ) -> Result, Self::Error> { + Ok(self.map.iter().find_map(|(k, v)| { + match k { + DescriptorPublicKey::Single(ref pk) => { + match key_request { + KeyRequest::Pubkey(ref request) => { + match pk.key { + SinglePubKey::FullKey(ref pk) => { + if pk == request { + match v { + DescriptorSecretKey::Single(ref sk) => Some(sk.key), + _ => unreachable!("Single maps to Single"), + } + } else { + None + } + } + SinglePubKey::XOnly(_) => None, + } + }, + _ => None, + } + }, + // Performance: Might be faster to check the origin and then if it matches return + // the key directly instead of calling `get_key` on the xpriv. + DescriptorPublicKey::XPub(ref xpub) => { + let pk = xpub.xkey.public_key; + match key_request { + KeyRequest::Pubkey(ref request) => { + if pk == request.inner { + match v { + DescriptorSecretKey::XPrv(xpriv) => { + let xkey = xpriv.xkey; + let child = xkey.derive_priv(secp, &xpriv.derivation_path); + Some(bitcoin::PrivateKey::new(child.private_key, xkey.network)) + } + _ => unreachable!("XPrv maps to XPrv"), + } + } else { + None + } + }, + KeyRequest::Bip32(..) => { + match v { + DescriptorSecretKey::XPrv(xpriv) => { + if let Ok(Some(sk)) = xpriv.xkey.get_key(key_request, secp) { + Some(sk) + } else { + None + } + } + _ => unreachable!("XPrv maps to XPrv"), + } + }, + _ => unreachable!("rust-bitcoin v0.32") + } + } + DescriptorPublicKey::MultiXPub(ref xpub) => { + let pk = xpub.xkey.public_key; + match key_request { + KeyRequest::Pubkey(ref request) => { + if pk == request.inner { + match v { + DescriptorSecretKey::MultiXPrv(xpriv) => Some(xpriv.xkey.to_private_key()), + _ => unreachable!("MultiXPrv maps to MultiXPrv"), + } + } else { + None + } + }, + KeyRequest::Bip32(..) => { + match v { + DescriptorSecretKey::MultiXPrv(xpriv) => { + if let Ok(Some(sk)) = xpriv.xkey.get_key(key_request, secp) { + Some(sk) + } else { + None + } + } + _ => unreachable!("MultiXPrv maps to MultiXPrv"), + } + }, + _ => unreachable!("rust-bitcoin v0.32") + } + }, + } + })) + } +} + +#[cfg(test)] +mod tests { + // use bitcoin::NetworkKind; + use bitcoin::bip32::{ChildNumber, 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::().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 pk = want_sk.public_key(&secp); + let request = KeyRequest::Pubkey(pk); + let got_sk = keymap.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::().unwrap(); + let xpriv_fingerprint = xpriv.fingerprint(&secp); + + // Sanity check. + { + let descriptor_sk_s = format!("[{}]{}", xpriv_fingerprint, xpriv); + let descriptor_sk = descriptor_sk_s.parse::().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 pk = want_sk.public_key(&secp); + let request = KeyRequest::Pubkey(pk); + let got_sk = keymap.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::().unwrap(); + let master_fingerprint = master.fingerprint(&secp); + + let child_number = ChildNumber::from_hardened_idx(44).unwrap(); + let child = master.derive_priv(&secp, &[child_number]); + + // Sanity check. + { + let descriptor_sk_s = format!("[{}/44']{}", master_fingerprint, child); + let descriptor_sk = descriptor_sk_s.parse::().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 pk = want_sk.public_key(&secp); + let request = KeyRequest::Pubkey(pk); + let got_sk = keymap.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::().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); + + // Sanity check. + { + let descriptor_sk_s = format!("[{}/{}]{}", master_fingerprint, first_external_child, child); + let descriptor_sk = descriptor_sk_s.parse::().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 key_source = (master_fingerprint, derivation_path); + let request = KeyRequest::Bip32(key_source); + let got_sk = keymap.get_key(&request, &secp) + .expect("get_key call errored") + .expect("failed to find the key"); + + assert_eq!(got_sk, want_sk) + } +} diff --git a/src/descriptor/mod.rs b/src/descriptor/mod.rs index b172aa69e..da9aeda30 100644 --- a/src/descriptor/mod.rs +++ b/src/descriptor/mod.rs @@ -46,21 +46,17 @@ pub use self::tr::{TapTree, Tr}; pub mod checksum; mod key; - -pub use self::key::{ - ConversionError, DefiniteDescriptorKey, DerivPaths, DescriptorKeyParseError, - DescriptorMultiXKey, DescriptorPublicKey, DescriptorSecretKey, DescriptorXKey, InnerXKey, - SinglePriv, SinglePub, SinglePubKey, Wildcard, +mod key_map; + +pub use self::{ + key::{ + ConversionError, DefiniteDescriptorKey, DerivPaths, DescriptorKeyParseError, + DescriptorMultiXKey, DescriptorPublicKey, DescriptorSecretKey, DescriptorXKey, InnerXKey, + SinglePriv, SinglePub, SinglePubKey, Wildcard, + }, + key_map::KeyMap, }; -/// Alias type for a map of public key to secret key -/// -/// This map is returned whenever a descriptor that contains secrets is parsed using -/// [`Descriptor::parse_descriptor`], since the descriptor will always only contain -/// public keys. This map allows looking up the corresponding secret key given a -/// public key from the descriptor. -pub type KeyMap = BTreeMap; - /// Script descriptor #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum Descriptor { @@ -728,7 +724,7 @@ impl Descriptor { Ok(public_key) } - let mut keymap_pk = KeyMapWrapper(BTreeMap::new(), secp); + let mut keymap_pk = KeyMapWrapper(KeyMap::new(), secp); struct KeyMapWrapper<'a, C: secp256k1::Signing>(KeyMap, &'a secp256k1::Secp256k1); diff --git a/src/plan.rs b/src/plan.rs index bcdfb311b..a9f5e305e 100644 --- a/src/plan.rs +++ b/src/plan.rs @@ -661,7 +661,7 @@ pub trait IntoAssets { } impl IntoAssets for KeyMap { - fn into_assets(self) -> Assets { Assets::from_iter(self.into_iter().map(|(k, _)| k)) } + fn into_assets(self) -> Assets { Assets::from_iter(self.map.into_iter().map(|(k, _)| k)) } } impl IntoAssets for DescriptorPublicKey {