Skip to content

Commit 4f304d1

Browse files
oleonardolimaoleonardolima
committed
wip(feat): impl GetKey for KeyMap 
1 parent 1005951 commit 4f304d1

File tree

1 file changed

+334
-1
lines changed

1 file changed

+334
-1
lines changed

‎src/descriptor/key_map.rs

Lines changed: 334 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,11 @@
44
55
use core::iter;
66

7-
use bitcoin::secp256k1::{Secp256k1, Signing};
7+
use bitcoin::{
8+
psbt::{GetKey, GetKeyError, KeyRequest},
9+
secp256k1::{Secp256k1, Signing},
10+
PrivateKey,
11+
};
812

913
#[cfg(doc)]
1014
use super::Descriptor;
@@ -87,3 +91,332 @@ impl iter::Extend<(DescriptorPublicKey, DescriptorSecretKey)> for KeyMap {
8791
self.map.extend(iter)
8892
}
8993
}
94+
95+
impl GetKey for KeyMap {
96+
type Error = GetKeyError;
97+
98+
fn get_key<C: Signing>(
99+
&self,
100+
key_request: KeyRequest,
101+
secp: &Secp256k1<C>,
102+
) -> Result<Option<bitcoin::PrivateKey>, Self::Error> {
103+
Ok(self
104+
.map
105+
.iter()
106+
.find_map(|(desc_pk, desc_sk)| -> Option<PrivateKey> {
107+
match (desc_sk, key_request.clone()) {
108+
(DescriptorSecretKey::Single(single_priv), key_request) => {
109+
let sk = single_priv.key;
110+
let pk = sk.public_key(secp);
111+
let pubkey_map = BTreeMap::from([(pk, sk)]);
112+
match pubkey_map.get_key(key_request, secp) {
113+
Ok(maybe_pk) => maybe_pk,
114+
Err(_) => None,
115+
}
116+
}
117+
(
118+
DescriptorSecretKey::XPrv(descriptor_xkey),
119+
KeyRequest::Pubkey(public_key),
120+
) => {
121+
let pk = match desc_pk {
122+
DescriptorPublicKey::XPub(descriptor_xkey) => {
123+
descriptor_xkey.xkey.public_key
124+
}
125+
_ => unreachable!(),
126+
};
127+
128+
if public_key.inner.eq(&pk) {
129+
if let Ok(xpriv) = descriptor_xkey
130+
.xkey
131+
.derive_priv(secp, &descriptor_xkey.derivation_path)
132+
{
133+
Some(xpriv.to_priv())
134+
} else {
135+
None
136+
}
137+
} else {
138+
None
139+
}
140+
}
141+
(
142+
DescriptorSecretKey::XPrv(descriptor_xkey),
143+
ref key_request @ KeyRequest::Bip32(ref key_source),
144+
) => {
145+
if let Ok(Some(key)) =
146+
descriptor_xkey.xkey.get_key(key_request.clone(), secp)
147+
{
148+
Some(key);
149+
}
150+
151+
if let Some(_derivation_path) = descriptor_xkey.matches(key_source, secp) {
152+
let (_fp, derivation_path) = key_source;
153+
154+
if let Ok(xpriv) =
155+
descriptor_xkey.xkey.derive_priv(secp, &derivation_path)
156+
{
157+
Some(xpriv.to_priv())
158+
} else {
159+
None
160+
}
161+
162+
// if let Some((_fp, origin_derivation_path)) = &descriptor_xkey.origin {
163+
// let derivation_path: DerivationPath =
164+
// derivation_path[origin_derivation_path.len()..].into();
165+
166+
// if let Ok(xpriv) =
167+
// descriptor_xkey.xkey.derive_priv(secp, &derivation_path)
168+
// {
169+
// Some(xpriv.to_priv())
170+
// } else {
171+
// None
172+
// }
173+
174+
// // return Some(
175+
// // descriptor_xkey
176+
// // .xkey
177+
// // .derive_priv(secp, &derivation_path)
178+
// // .map_err(GetKeyError::Bip32)?
179+
// // .to_priv(),
180+
// // );
181+
// } else {
182+
// if let Ok(xpriv) =
183+
// descriptor_xkey.xkey.derive_priv(secp, &derivation_path)
184+
// {
185+
// Some(xpriv.to_priv())
186+
// } else {
187+
// None
188+
// }
189+
190+
// // return Some(
191+
// // descriptor_xkey
192+
// // .xkey
193+
// // .derive_priv(secp, derivation_path)
194+
// // .map_err(GetKeyError::Bip32)?
195+
// // .to_priv(),
196+
// // );
197+
// }
198+
} else {
199+
None
200+
}
201+
}
202+
(DescriptorSecretKey::XPrv(_), KeyRequest::XOnlyPubkey(_)) => None,
203+
(
204+
DescriptorSecretKey::MultiXPrv(descriptor_multi_xkey),
205+
KeyRequest::Pubkey(_public_key),
206+
) => {
207+
// let pk = match desc_pk {
208+
// DescriptorPublicKey::MultiXPub(descriptor_multi_xkey) => {
209+
// descriptor_multi_xkey.xkey.public_key
210+
// }
211+
// _ => unreachable!(),
212+
// };
213+
214+
Some(descriptor_multi_xkey.xkey.to_priv())
215+
// if public_key.inner.eq(&pk) {
216+
// if let Ok(xpriv) = descriptor_multi_xkey
217+
// .xkey
218+
// .derive_priv(secp, &descriptor_multi_xkey.derivation_path)
219+
// {
220+
// Some(xpriv.to_priv())
221+
// } else {
222+
// None
223+
// }
224+
// } else {
225+
// None
226+
// }
227+
}
228+
(
229+
DescriptorSecretKey::MultiXPrv(descriptor_multi_xkey),
230+
KeyRequest::Bip32(_),
231+
) => {
232+
if let Ok(Some(key)) = descriptor_multi_xkey
233+
.xkey
234+
.get_key(key_request.clone(), secp)
235+
{
236+
Some(key)
237+
} else {
238+
None
239+
}
240+
}
241+
(DescriptorSecretKey::MultiXPrv(_), KeyRequest::XOnlyPubkey(_)) => None,
242+
_ => unreachable!(),
243+
}
244+
}))
245+
}
246+
}
247+
248+
#[cfg(test)]
249+
mod tests {
250+
use core::str::FromStr;
251+
252+
// use bitcoin::NetworkKind;
253+
use bitcoin::bip32::{ChildNumber, DerivationPath, IntoDerivationPath, Xpriv};
254+
255+
use super::*;
256+
use crate::Descriptor;
257+
258+
#[test]
259+
fn get_key_single_key() {
260+
let secp = Secp256k1::new();
261+
262+
let descriptor_sk_s =
263+
"[90b6a706/44'/0'/0'/0/0]cMk8gWmj1KpjdYnAWwsEDekodMYhbyYBhG8gMtCCxucJ98JzcNij";
264+
265+
let single = match descriptor_sk_s.parse::<DescriptorSecretKey>().unwrap() {
266+
DescriptorSecretKey::Single(single) => single,
267+
_ => panic!("unexpected DescriptorSecretKey variant"),
268+
};
269+
270+
let want_sk = single.key;
271+
let descriptor_s = format!("wpkh({})", descriptor_sk_s);
272+
let (_, keymap) = Descriptor::parse_descriptor(&secp, &descriptor_s).unwrap();
273+
274+
let pk = want_sk.public_key(&secp);
275+
let request = KeyRequest::Pubkey(pk);
276+
let got_sk = keymap
277+
.get_key(request, &secp)
278+
.expect("get_key call errored")
279+
.expect("failed to find the key");
280+
assert_eq!(got_sk, want_sk)
281+
}
282+
283+
#[test]
284+
fn get_key_xpriv_single_key_xpriv() {
285+
let secp = Secp256k1::new();
286+
287+
let s = "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi";
288+
289+
let xpriv = s.parse::<Xpriv>().unwrap();
290+
let xpriv_fingerprint = xpriv.fingerprint(&secp);
291+
292+
// Sanity check.
293+
{
294+
let descriptor_sk_s = format!("[{}]{}", xpriv_fingerprint, xpriv);
295+
let descriptor_sk = descriptor_sk_s.parse::<DescriptorSecretKey>().unwrap();
296+
let got = match descriptor_sk {
297+
DescriptorSecretKey::XPrv(x) => x.xkey,
298+
_ => panic!("unexpected DescriptorSecretKey variant"),
299+
};
300+
assert_eq!(got, xpriv);
301+
}
302+
303+
let want_sk = xpriv.to_priv();
304+
let descriptor_s = format!("wpkh([{}]{})", xpriv_fingerprint, xpriv);
305+
let (_, keymap) = Descriptor::parse_descriptor(&secp, &descriptor_s).unwrap();
306+
307+
let pk = want_sk.public_key(&secp);
308+
let request = KeyRequest::Pubkey(pk);
309+
let got_sk = keymap
310+
.get_key(request, &secp)
311+
.expect("get_key call errored")
312+
.expect("failed to find the key");
313+
assert_eq!(got_sk, want_sk)
314+
}
315+
316+
#[test]
317+
fn get_key_xpriv_child_depth_one() {
318+
let secp = Secp256k1::new();
319+
320+
let s = "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi";
321+
let master = s.parse::<Xpriv>().unwrap();
322+
let master_fingerprint = master.fingerprint(&secp);
323+
324+
let child_number = ChildNumber::from_hardened_idx(44).unwrap();
325+
let child = master.derive_priv(&secp, &[child_number]).unwrap();
326+
327+
// Sanity check.
328+
{
329+
let descriptor_sk_s = format!("[{}/44']{}", master_fingerprint, child);
330+
let descriptor_sk = descriptor_sk_s.parse::<DescriptorSecretKey>().unwrap();
331+
let got = match descriptor_sk {
332+
DescriptorSecretKey::XPrv(ref x) => x.xkey,
333+
_ => panic!("unexpected DescriptorSecretKey variant"),
334+
};
335+
assert_eq!(got, child);
336+
}
337+
338+
let want_sk = child.to_priv();
339+
let descriptor_s = format!("wpkh({}/44')", s);
340+
let (_, keymap) = Descriptor::parse_descriptor(&secp, &descriptor_s).unwrap();
341+
342+
let pk = want_sk.public_key(&secp);
343+
let request = KeyRequest::Pubkey(pk);
344+
let got_sk = keymap
345+
.get_key(request, &secp)
346+
.expect("get_key call errored")
347+
.expect("failed to find the key");
348+
assert_eq!(got_sk, want_sk)
349+
}
350+
351+
#[test]
352+
fn get_key_xpriv_with_path() {
353+
let secp = Secp256k1::new();
354+
355+
let s = "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi";
356+
let master = s.parse::<Xpriv>().unwrap();
357+
let master_fingerprint = master.fingerprint(&secp);
358+
359+
let first_external_child = "44'/0'/0'/0/0";
360+
let derivation_path = first_external_child.into_derivation_path().unwrap();
361+
362+
let child = master.derive_priv(&secp, &derivation_path).unwrap();
363+
364+
// Sanity check.
365+
{
366+
let descriptor_sk_s =
367+
format!("[{}/{}]{}", master_fingerprint, first_external_child, child);
368+
let descriptor_sk = descriptor_sk_s.parse::<DescriptorSecretKey>().unwrap();
369+
let got = match descriptor_sk {
370+
DescriptorSecretKey::XPrv(ref x) => x.xkey,
371+
_ => panic!("unexpected DescriptorSecretKey variant"),
372+
};
373+
assert_eq!(got, child);
374+
}
375+
376+
let want_sk = child.to_priv();
377+
let descriptor_s = format!("wpkh({}/44'/0'/0'/0/*)", s);
378+
let (_, keymap) = Descriptor::parse_descriptor(&secp, &descriptor_s).unwrap();
379+
380+
let key_source = (master_fingerprint, derivation_path);
381+
let request = KeyRequest::Bip32(key_source);
382+
let got_sk = keymap
383+
.get_key(request, &secp)
384+
.expect("get_key call errored")
385+
.expect("failed to find the key");
386+
387+
assert_eq!(got_sk, want_sk)
388+
}
389+
390+
#[test]
391+
fn get_key_xpriv_with_key_origin() {
392+
let secp = Secp256k1::new();
393+
394+
let descriptor_str = "wpkh([d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)";
395+
let (_descriptor_pk, keymap) = Descriptor::parse_descriptor(&secp, descriptor_str).unwrap();
396+
397+
let descriptor_sk = DescriptorSecretKey::from_str("[d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*").unwrap();
398+
let xpriv = match descriptor_sk {
399+
DescriptorSecretKey::XPrv(descriptor_xkey) => descriptor_xkey,
400+
_ => unreachable!(),
401+
};
402+
403+
let path = DerivationPath::from_str("84'/1'/0'/0").unwrap();
404+
// let child_number = ChildNumber::from_normal_idx(0).unwrap();
405+
// let child = master.derive_priv(&secp, &[child_number]).unwrap();
406+
let expected_pk = xpriv
407+
.xkey
408+
.derive_priv(&secp, &path)
409+
.unwrap()
410+
.to_priv();
411+
412+
let (fp, _) = xpriv.origin.unwrap();
413+
let key_request = KeyRequest::Bip32((fp, path));
414+
415+
let pk = keymap
416+
.get_key(key_request, &secp)
417+
.expect("get_key should not fail")
418+
.expect("get_key should return a `PrivateKey`");
419+
420+
assert_eq!(pk, expected_pk);
421+
}
422+
}

0 commit comments

Comments
 (0)