Fix tap_key_origins when there are pkh(..) terms

Author: LLFourn

integration_test/src/test_desc.rs

@@ -109,10 +109,10 @@ pub fn test_desc_satisfy(cl: &Client, testdata: &TestData, desc: &str) -> Witnes
         script_pubkey: addr.script_pubkey(),
     });
     let mut input = psbt::Input::default();
+    input.update_with_descriptor_unchecked(&desc).unwrap();
     input.witness_utxo = Some(witness_utxo.clone());
     psbt.inputs.push(input);
     psbt.outputs.push(psbt::Output::default());
-    psbt.update_inp_with_descriptor(0, &desc, true, false).unwrap();
 
     // --------------------------------------------
     // Sign the transactions with all keys

src/psbt/mod.rs

@@ -33,11 +33,15 @@ use bitcoin::{EcdsaSigHashType, Script};
 use bitcoin::util::taproot::{self, ControlBlock, LeafVersion, TapLeafHash};
 use descriptor;
 use interpreter;
+use miniscript::iter::PkPkh;
 use miniscript::limits::SEQUENCE_LOCKTIME_DISABLE_FLAG;
 use miniscript::satisfy::{After, Older};
 use Preimage32;
 use Satisfier;
-use {Descriptor, DescriptorPublicKey, DescriptorTrait, MiniscriptKey, ToPublicKey, TranslatePk2};
+use {
+    Descriptor, DescriptorPublicKey, DescriptorTrait, MiniscriptKey, ToPublicKey, TranslatePk,
+    TranslatePk2,
+};
 
 mod finalizer;
 
@@ -872,6 +876,12 @@ pub trait PsbtInputExt {
     /// Note that his method doesn't check that the `witness_utxo` or `non_witness_utxo` is
     /// consistent with the descriptor. To do that see [`update_input_with_descriptor`].
     ///
+    /// ## Return value
+    ///
+    /// For convenience, this returns the concrete descriptor that is computed internally to fill
+    /// out the PSBT input fields. If you are checking the validity of `witness_utxo` or
+    /// `non_witness_utxo` yourself you should check the `script_pubkey` against it.
+    ///
     /// [`update_input_with_descriptor`]: PsbtExt::update_input_with_descriptor
     fn update_with_descriptor_unchecked(
         &mut self,
@@ -889,48 +899,75 @@ impl PsbtInputExt for psbt::Input {
 
         let derived = if let Descriptor::Tr(_) = &descriptor {
             // have to use a RefCell because we can't pass FnMut to translate_pk2
-            let tap_key_origins = RefCell::new(BTreeMap::new());
-            let derived = descriptor.translate_pk2(|xpk| {
-                let derived = xpk.derive_public_key(&secp)?;
-                tap_key_origins.borrow_mut().insert(
-                    derived.to_x_only_pubkey(),
+            let mut hash_lookup = BTreeMap::new();
+            let derived = descriptor.translate_pk(
+                |xpk| xpk.derive_public_key(&secp),
+                |xpk| {
+                    let xonly = xpk.derive_public_key(&secp)?.to_x_only_pubkey();
+                    let hash = xonly.to_pubkeyhash();
+                    hash_lookup.insert(hash, xonly);
+                    Ok(hash)
+                },
+            )?;
+
+            // NOTE: they will both always be Tr
+            if let (Descriptor::Tr(tr_derived), Descriptor::Tr(tr_xpk)) = (&derived, descriptor) {
+                let spend_info = tr_derived.spend_info();
+                let ik_derived = spend_info.internal_key();
+                let ik_xpk = tr_xpk.internal_key();
+                self.tap_internal_key = Some(ik_derived);
+                self.tap_merkle_root = spend_info.merkle_root();
+                self.tap_key_origins.insert(
+                    ik_derived,
                     (
-                        vec![/* we'll populate this in the next loop */],
-                        (xpk.master_fingerprint(), xpk.full_derivation_path()),
+                        vec![],
+                        (ik_xpk.master_fingerprint(), ik_xpk.full_derivation_path()),
                     ),
                 );
-                Ok(derived)
-            })?;
 
-            let mut tap_key_origins = tap_key_origins.into_inner();
-
-            // NOTE: the derived descriptor will always be Tr
-            if let Descriptor::Tr(tr) = &derived {
-                let spend_info = tr.spend_info();
-                for (_depth, ms) in tr.iter_scripts() {
-                    let leaf_script = (ms.encode(), LeafVersion::TapScript);
+                for ((_depth_der, ms_derived), (_depth, ms)) in
+                    tr_derived.iter_scripts().zip(tr_xpk.iter_scripts())
+                {
+                    debug_assert_eq!(_depth_der, _depth);
+                    let leaf_script = (ms_derived.encode(), LeafVersion::TapScript);
                     let tapleaf_hash = TapLeafHash::from_script(&leaf_script.0, leaf_script.1);
                     let control_block = spend_info
                         .control_block(&leaf_script)
                         .expect("Control block must exist in script map for every known leaf");
                     self.tap_scripts.insert(control_block, leaf_script);
 
-                    for pk in ms.iter_pk() {
-                        if let Some((tapleaf_hashes, _)) =
-                            tap_key_origins.get_mut(&pk.to_x_only_pubkey())
-                        {
-                            // To avoid duplication when the same key is in the same leaf more than
-                            // once. (even though I think this is usually not allowed)
-                            if tapleaf_hashes.last() != Some(&tapleaf_hash) {
-                                tapleaf_hashes.push(tapleaf_hash)
+                    for (pk_pkh_derived, pk_pkh_xpk) in
+                        ms_derived.iter_pk_pkh().zip(ms.iter_pk_pkh())
+                    {
+                        let (xonly, xpk) = match (pk_pkh_derived, pk_pkh_xpk) {
+                            (PkPkh::PlainPubkey(pk), PkPkh::PlainPubkey(xpk)) => {
+                                (pk.to_x_only_pubkey(), xpk)
                             }
-                        }
+                            (PkPkh::HashedPubkey(hash), PkPkh::HashedPubkey(xpk)) => (
+                                hash_lookup
+                                    .get(&hash)
+                                    .expect("translate_pk inserted an entry for every hash")
+                                    .clone(),
+                                xpk,
+                            ),
+                            _ => unreachable!("the iterators work in the same order"),
+                        };
+
+                        self.tap_key_origins
+                            .entry(xonly)
+                            .and_modify(|(tapleaf_hashes, _)| {
+                                if tapleaf_hashes.last() != Some(&tapleaf_hash) {
+                                    tapleaf_hashes.push(tapleaf_hash);
+                                }
+                            })
+                            .or_insert_with(|| {
+                                (
+                                    vec![tapleaf_hash],
+                                    (xpk.master_fingerprint(), xpk.full_derivation_path()),
+                                )
+                            });
                     }
                 }
-
-                self.tap_internal_key = Some(spend_info.internal_key());
-                self.tap_merkle_root = spend_info.merkle_root();
-                self.tap_key_origins = tap_key_origins;
             }
 
             derived
@@ -944,6 +981,7 @@ impl PsbtInputExt for psbt::Input {
                 );
                 Ok(derived)
             })?;
+
             self.bip32_derivation = bip32_derivation.into_inner();
 
             match &derived {