Merge rust-bitcoin/rust-miniscript#557: Add method to convert expr_raw_pkh into pkh

2f3e457fd21a022db2521280542fd74b624ccc9a Using BTreeMap in place of HashMap (Harshil Jani)
c17965a97c71032d5a0718d59d766b83c6b1762f Add method to convert expr_raw_pkh into pkh (Harshil Jani)

Pull request description:

  The PR adds up the method to convert the expr_raw_pkh into pkh by looking at the available keys and their hases being stored inside of the Hashmap of PSBT Input. If the key is found corresponding to given hash, It will update to pkh.

  To-Do :
  - [ ] Write a  test case.
  - [ ] Make it work for taproot.

ACKs for top commit:
  sanket1729:
    utACK 2f3e457fd21a022db2521280542fd74b624ccc9a. Thanks for the contribution

Tree-SHA512: c98d7dfe0055e01f7ff3b0c23a73faa79136aa1d446f2b1e035931c1d256712668d4a150674dd35448682dd4c91db6ea426557c48ec006f6335af796ddb8a4f4

Author: heap-coder

src/miniscript/astelem.rs

@@ -179,6 +179,75 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> Terminal<Pk, Ctx> {
         };
         Ok(frag)
     }
+
+    /// Substitutes raw public keys hashes with the public keys as provided by map.
+    pub fn substitute_raw_pkh(&self, pk_map: &BTreeMap<hash160::Hash, Pk>) -> Terminal<Pk, Ctx> {
+        match self {
+            Terminal::RawPkH(ref p) => match pk_map.get(p) {
+                Some(pk) => Terminal::PkH(pk.clone()).into(),
+                None => Terminal::RawPkH(*p).into(),
+            },
+            Terminal::PkK(..)
+            | Terminal::PkH(..)
+            | Terminal::Multi(..)
+            | Terminal::MultiA(..)
+            | Terminal::After(..)
+            | Terminal::Older(..)
+            | Terminal::Sha256(..)
+            | Terminal::Hash256(..)
+            | Terminal::Ripemd160(..)
+            | Terminal::Hash160(..)
+            | Terminal::True
+            | Terminal::False => self.clone().into(),
+            Terminal::Alt(ref sub) => Terminal::Alt(Arc::new(sub.substitute_raw_pkh(pk_map))),
+            Terminal::Swap(ref sub) => Terminal::Swap(Arc::new(sub.substitute_raw_pkh(pk_map))),
+            Terminal::Check(ref sub) => Terminal::Check(Arc::new(sub.substitute_raw_pkh(pk_map))),
+            Terminal::DupIf(ref sub) => Terminal::DupIf(Arc::new(sub.substitute_raw_pkh(pk_map))),
+            Terminal::Verify(ref sub) => Terminal::Verify(Arc::new(sub.substitute_raw_pkh(pk_map))),
+            Terminal::NonZero(ref sub) => {
+                Terminal::NonZero(Arc::new(sub.substitute_raw_pkh(pk_map)))
+            }
+            Terminal::ZeroNotEqual(ref sub) => {
+                Terminal::ZeroNotEqual(Arc::new(sub.substitute_raw_pkh(pk_map)))
+            }
+            Terminal::AndV(ref left, ref right) => Terminal::AndV(
+                Arc::new(left.substitute_raw_pkh(pk_map)),
+                Arc::new(right.substitute_raw_pkh(pk_map)),
+            ),
+            Terminal::AndB(ref left, ref right) => Terminal::AndB(
+                Arc::new(left.substitute_raw_pkh(pk_map)),
+                Arc::new(right.substitute_raw_pkh(pk_map)),
+            ),
+            Terminal::AndOr(ref a, ref b, ref c) => Terminal::AndOr(
+                Arc::new(a.substitute_raw_pkh(pk_map)),
+                Arc::new(b.substitute_raw_pkh(pk_map)),
+                Arc::new(c.substitute_raw_pkh(pk_map)),
+            ),
+            Terminal::OrB(ref left, ref right) => Terminal::OrB(
+                Arc::new(left.substitute_raw_pkh(pk_map)),
+                Arc::new(right.substitute_raw_pkh(pk_map)),
+            ),
+            Terminal::OrD(ref left, ref right) => Terminal::OrD(
+                Arc::new(left.substitute_raw_pkh(pk_map)),
+                Arc::new(right.substitute_raw_pkh(pk_map)),
+            ),
+            Terminal::OrC(ref left, ref right) => Terminal::OrC(
+                Arc::new(left.substitute_raw_pkh(pk_map)),
+                Arc::new(right.substitute_raw_pkh(pk_map)),
+            ),
+            Terminal::OrI(ref left, ref right) => Terminal::OrI(
+                Arc::new(left.substitute_raw_pkh(pk_map)),
+                Arc::new(right.substitute_raw_pkh(pk_map)),
+            ),
+            Terminal::Thresh(k, ref subs) => {
+                let subs: Vec<Arc<Miniscript<_, _>>> = subs
+                    .iter()
+                    .map(|s| Arc::new(s.substitute_raw_pkh(pk_map)))
+                    .collect();
+                Terminal::Thresh(*k, subs)
+            }
+        }
+    }
 }
 
 impl<Pk: MiniscriptKey, Ctx: ScriptContext> ForEachKey<Pk> for Terminal<Pk, Ctx> {

src/miniscript/mod.rs

@@ -16,6 +16,7 @@
 use core::marker::PhantomData;
 use core::{fmt, hash, str};
 
+use bitcoin::hashes::hash160;
 use bitcoin::script;
 use bitcoin::taproot::{LeafVersion, TapLeafHash};
 
@@ -334,6 +335,11 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> Miniscript<Pk, Ctx> {
         let inner = self.node.real_translate_pk(t)?;
         Miniscript::from_ast(inner).map_err(TranslateErr::OuterError)
     }
+
+    /// Substitutes raw public keys hashes with the public keys as provided by map.
+    pub fn substitute_raw_pkh(&self, pk_map: &BTreeMap<hash160::Hash, Pk>) -> Miniscript<Pk, Ctx> {
+        Miniscript::from_ast(self.node.substitute_raw_pkh(pk_map)).expect("type check failed")
+    }
 }
 
 impl_block_str!(

src/psbt/finalizer.rs

@@ -8,6 +8,7 @@
 //! `https://github.com/bitcoin/bips/blob/master/bip-0174.mediawiki`
 //!
 
+use bitcoin::hashes::hash160;
 use bitcoin::key::XOnlyPublicKey;
 use bitcoin::secp256k1::{self, Secp256k1};
 use bitcoin::sighash::Prevouts;
@@ -122,6 +123,22 @@ pub(super) fn prevouts(psbt: &Psbt) -> Result<Vec<&bitcoin::TxOut>, super::Error
 // we want to move the script is probably already created
 // and we want to satisfy it in any way possible.
 fn get_descriptor(psbt: &Psbt, index: usize) -> Result<Descriptor<PublicKey>, InputError> {
+    // Create a HashMap <Hash(Pk),Pk>
+    let mut hash_map: BTreeMap<hash160::Hash, PublicKey> = BTreeMap::new();
+    let psbt_inputs = &psbt.inputs;
+    for psbt_input in psbt_inputs {
+        // Use BIP32 Derviation to get set of all possible keys.
+        let public_keys = psbt_input.bip32_derivation.keys();
+        for key in public_keys {
+            // Convert bitcoin::secp256k1::PublicKey into just bitcoin::PublicKey
+            let bitcoin_key = bitcoin::PublicKey::new(*key);
+            // Convert PubKeyHash into Hash::hash160
+            let hash = bitcoin_key.pubkey_hash().to_raw_hash();
+            // Insert pair in HashMap
+            hash_map.insert(hash, bitcoin_key);
+        }
+    }
+
     // Figure out Scriptpubkey
     let script_pubkey = get_scriptpubkey(psbt, index)?;
     let inp = &psbt.inputs[index];
@@ -178,7 +195,7 @@ fn get_descriptor(psbt: &Psbt, index: usize) -> Result<Descriptor<PublicKey>, In
                 witness_script,
                 &ExtParams::allow_all(),
             )?;
-            Ok(Descriptor::new_wsh(ms)?)
+            Ok(Descriptor::new_wsh(ms.substitute_raw_pkh(&hash_map))?)
         } else {
             Err(InputError::MissingWitnessScript)
         }
@@ -205,7 +222,7 @@ fn get_descriptor(psbt: &Psbt, index: usize) -> Result<Descriptor<PublicKey>, In
                             witness_script,
                             &ExtParams::allow_all(),
                         )?;
-                        Ok(Descriptor::new_sh_wsh(ms)?)
+                        Ok(Descriptor::new_sh_wsh(ms.substitute_raw_pkh(&hash_map))?)
                     } else {
                         Err(InputError::MissingWitnessScript)
                     }
@@ -249,7 +266,7 @@ fn get_descriptor(psbt: &Psbt, index: usize) -> Result<Descriptor<PublicKey>, In
             &script_pubkey,
             &ExtParams::allow_all(),
         )?;
-        Ok(Descriptor::new_bare(ms)?)
+        Ok(Descriptor::new_bare(ms.substitute_raw_pkh(&hash_map))?)
     }
 }