11use  alloc:: string:: ToString ; 
2+ use  alloc:: vec:: Vec ; 
3+ use  std:: collections:: BTreeMap ; 
24
35use  bitcoin:: { 
46    psbt:: { GetKey ,  GetKeyError ,  KeyRequest } , 
@@ -24,9 +26,10 @@ impl GetKey for Signer {
2426        for  entry in  & self . 0  { 
2527            match  entry { 
2628                ( _,  DescriptorSecretKey :: Single ( prv) )  => { 
27-                     let  pk = prv. key . public_key ( secp) ; 
28-                     if  key_request == KeyRequest :: Pubkey ( pk)  { 
29-                         return  Ok ( Some ( prv. key ) ) ; 
29+                     let  map:  BTreeMap < _ ,  _ >  =
30+                         core:: iter:: once ( ( prv. key . public_key ( secp) ,  prv. key ) ) . collect ( ) ; 
31+                     if  let  Ok ( Some ( prv) )  = GetKey :: get_key ( & map,  key_request. clone ( ) ,  secp)  { 
32+                         return  Ok ( Some ( prv) ) ; 
3033                    } 
3134                } 
3235                ( _,  desc_sk)  => { 
@@ -44,7 +47,11 @@ impl GetKey for Signer {
4447                                    if  fingerprint == fp
4548                                        && derivation. to_string ( ) . starts_with ( & path. to_string ( ) ) 
4649                                    { 
47-                                         let  to_derive = & derivation[ k. xkey . depth  as  usize ..] ; 
50+                                         let  to_derive = derivation
51+                                             . into_iter ( ) 
52+                                             . skip ( path. len ( ) ) 
53+                                             . cloned ( ) 
54+                                             . collect :: < Vec < _ > > ( ) ; 
4855                                        let  derived = k. xkey . derive_priv ( secp,  & to_derive) ?; 
4956                                        return  Ok ( Some ( derived. to_priv ( ) ) ) ; 
5057                                    } 
@@ -61,6 +68,8 @@ impl GetKey for Signer {
6168
6269#[ cfg( test) ]  
6370mod  test { 
71+     use  crate :: bitcoin:: bip32:: ChildNumber ; 
72+     use  core:: str:: FromStr ; 
6473    use  std:: string:: String ; 
6574
6675    use  bitcoin:: bip32:: { DerivationPath ,  Xpriv } ; 
@@ -89,6 +98,26 @@ mod test {
8998        Ok ( ( ) ) 
9099    } 
91100
101+     #[ test]  
102+     fn  get_key_x_only_pubkey ( )  -> anyhow:: Result < ( ) >  { 
103+         let  secp = Secp256k1 :: new ( ) ; 
104+         let  wif = "cU6BxEezV8FnkEPBCaFtc4WNuUKmgFaAu6sJErB154GXgMUjhgWe" ; 
105+         let  prv = bitcoin:: PrivateKey :: from_wif ( wif) ?; 
106+         let  ( x_only_pk,  _parity)  = prv. inner . x_only_public_key ( & secp) ; 
107+ 
108+         let  s = format ! ( "wpkh({wif})" ) ; 
109+         let  ( _,  keymap)  = Descriptor :: parse_descriptor ( & secp,  & s) . unwrap ( ) ; 
110+ 
111+         let  signer = Signer ( keymap) ; 
112+         let  req = KeyRequest :: XOnlyPubkey ( x_only_pk) ; 
113+         let  res = signer. get_key ( req,  & secp) ; 
114+         assert ! ( matches!( 
115+             res, 
116+             Ok ( Some ( k) )  if  k. inner. x_only_public_key( & secp) . 0  == x_only_pk
117+         ) ) ; 
118+         Ok ( ( ) ) 
119+     } 
120+ 
92121    // Test `Signer` can fulfill a bip32 KeyRequest if we know the key origin 
93122    #[ test]  
94123    fn  get_key_bip32 ( )  -> anyhow:: Result < ( ) >  { 
@@ -141,4 +170,35 @@ mod test {
141170
142171        Ok ( ( ) ) 
143172    } 
173+ 
174+     #[ test]  
175+     fn  get_key_xpriv_with_key_origin ( )  -> anyhow:: Result < ( ) >  { 
176+         let  secp = Secp256k1 :: new ( ) ; 
177+         let  s = "wpkh([d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)" ; 
178+         let  ( _,  keymap)  = Descriptor :: parse_descriptor ( & secp,  s) ?; 
179+ 
180+         let  desc_sk = DescriptorSecretKey :: from_str ( "[d34db33f/84h/1h/0h]tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*" ) ?; 
181+         let  desc_xkey = match  desc_sk { 
182+             DescriptorSecretKey :: XPrv ( k)  => k, 
183+             _ => panic ! ( ) , 
184+         } ; 
185+ 
186+         let  ( fp,  _)  = desc_xkey. origin . clone ( ) . unwrap ( ) ; 
187+         let  path = DerivationPath :: from_str ( "84h/1h/0h/7" ) ?; 
188+         let  req = KeyRequest :: Bip32 ( ( fp,  path) ) ; 
189+ 
190+         let  exp_prv = desc_xkey
191+             . xkey 
192+             . derive_priv ( & secp,  & [ ChildNumber :: from ( 7 ) ] ) ?
193+             . to_priv ( ) ; 
194+ 
195+         let  res = Signer ( keymap) . get_key ( req,  & secp) ; 
196+ 
197+         assert ! ( matches!( 
198+             res, 
199+             Ok ( Some ( k) )  if  k == exp_prv, 
200+         ) ) ; 
201+ 
202+         Ok ( ( ) ) 
203+     } 
144204} 
0 commit comments