Populate tap_key_origin in PSBT inputs and outputs

src/descriptor/mod.rs

@@ -19,7 +19,7 @@ use std::ops::Deref;
 
 use bitcoin::secp256k1;
 use bitcoin::util::bip32::{ChildNumber, DerivationPath, ExtendedPubKey, Fingerprint, KeySource};
-use bitcoin::util::psbt;
+use bitcoin::util::{psbt, taproot};
 use bitcoin::{Network, Script, TxOut};
 
 use miniscript::descriptor::{DescriptorType, InnerXKey};
@@ -61,6 +61,13 @@ pub type DerivedDescriptor<'s> = Descriptor<DerivedDescriptorKey<'s>>;
 /// [`psbt::Output`]: bitcoin::util::psbt::Output
 pub type HdKeyPaths = BTreeMap<secp256k1::PublicKey, KeySource>;
 
+/// Alias for the type of maps that represent taproot key origins in a [`psbt::Input`] or
+/// [`psbt::Output`]
+///
+/// [`psbt::Input`]: bitcoin::util::psbt::Input
+/// [`psbt::Output`]: bitcoin::util::psbt::Output
+pub type TapKeyOrigins = BTreeMap<bitcoin::XOnlyPublicKey, (Vec<taproot::TapLeafHash>, KeySource)>;
+
 /// Trait for types which can be converted into an [`ExtendedDescriptor`] and a [`KeyMap`] usable by a wallet in a specific [`Network`]
 pub trait IntoWalletDescriptor {
     /// Convert to wallet descriptor
@@ -302,7 +309,8 @@ where
 }
 
 pub(crate) trait DerivedDescriptorMeta {
-    fn get_hd_keypaths(&self, secp: &SecpCtx) -> Result<HdKeyPaths, DescriptorError>;
+    fn get_hd_keypaths(&self, secp: &SecpCtx) -> HdKeyPaths;
+    fn get_tap_key_origins(&self, secp: &SecpCtx) -> TapKeyOrigins;
 }
 
 pub(crate) trait DescriptorMeta {
@@ -497,7 +505,7 @@ impl DescriptorMeta for ExtendedDescriptor {
 }
 
 impl<'s> DerivedDescriptorMeta for DerivedDescriptor<'s> {
-    fn get_hd_keypaths(&self, secp: &SecpCtx) -> Result<HdKeyPaths, DescriptorError> {
+    fn get_hd_keypaths(&self, secp: &SecpCtx) -> HdKeyPaths {
         let mut answer = BTreeMap::new();
         self.for_each_key(|key| {
             if let DescriptorPublicKey::XPub(xpub) = key.as_key().deref() {
@@ -515,7 +523,64 @@ impl<'s> DerivedDescriptorMeta for DerivedDescriptor<'s> {
             true
         });
 
-        Ok(answer)
+        answer
+    }
+
+    fn get_tap_key_origins(&self, secp: &SecpCtx) -> TapKeyOrigins {
+        use miniscript::ToPublicKey;
+
+        let mut answer = BTreeMap::new();
+        let mut insert_path = |pk: &DerivedDescriptorKey<'_>, lh| {
+            let key_origin = match pk.deref() {
+                DescriptorPublicKey::XPub(xpub) => {
+                    Some((xpub.root_fingerprint(secp), xpub.full_path(&[])))
+                }
+                DescriptorPublicKey::SinglePub(_) => None,
+            };
+
+            // If this is the internal key, we only insert the key origin if it's not None.
+            // For keys found in the tap tree we always insert a key origin (because the signer
+            // looks for it to know which leaves to sign for), even though it may be None
+            match (lh, key_origin) {
+                (None, Some(ko)) => {
+                    answer
+                        .entry(pk.to_x_only_pubkey())
+                        .or_insert_with(|| (vec![], ko));
+                }
+                (Some(lh), origin) => {
+                    answer
+                        .entry(pk.to_x_only_pubkey())
+                        .or_insert_with(|| (vec![], origin.unwrap_or_default()))
+                        .0
+                        .push(lh);
+                }
+                _ => {}
+            }
+        };
+
+        if let Descriptor::Tr(tr) = &self {
+            // Internal key first, then iterate the scripts
+            insert_path(tr.internal_key(), None);
+
+            for (_, ms) in tr.iter_scripts() {
+                // Assume always the same leaf version
+                let leaf_hash = taproot::TapLeafHash::from_script(
+                    &ms.encode(),
+                    taproot::LeafVersion::TapScript,
+                );
+
+                for key in ms.iter_pk_pkh() {
+                    let key = match key {
+                        miniscript::miniscript::iter::PkPkh::PlainPubkey(pk) => pk,
+                        miniscript::miniscript::iter::PkPkh::HashedPubkey(pk) => pk,
+                    };
+
+                    insert_path(&key, Some(leaf_hash));
+                }
+            }
+        }
+
+        answer
     }
 }
 

src/wallet/mod.rs

@@ -1225,7 +1225,7 @@ where
 
         let derived_descriptor = descriptor.as_derived(index, &self.secp);
 
-        let hd_keypaths = derived_descriptor.get_hd_keypaths(&self.secp)?;
+        let hd_keypaths = derived_descriptor.get_hd_keypaths(&self.secp);
         let script = derived_descriptor.script_pubkey();
 
         for validator in &self.address_validators {
@@ -1436,6 +1436,7 @@ where
                     psbt_input: foreign_psbt_input,
                     outpoint,
                 } => {
+                    // TODO: do not require non_witness_utxo for taproot utxos
                     if !params.only_witness_utxo && foreign_psbt_input.non_witness_utxo.is_none() {
                         return Err(Error::Generic(format!(
                             "Missing non_witness_utxo on foreign utxo {}",
@@ -1461,7 +1462,15 @@ where
                 let (desc, _) = self._get_descriptor_for_keychain(keychain);
                 let derived_descriptor = desc.as_derived(child, &self.secp);
 
-                psbt_output.bip32_derivation = derived_descriptor.get_hd_keypaths(&self.secp)?;
+                if desc.is_taproot() {
+                    psbt_output
+                        .tap_key_origins
+                        .append(&mut derived_descriptor.get_tap_key_origins(&self.secp));
+                } else {
+                    psbt_output
+                        .bip32_derivation
+                        .append(&mut derived_descriptor.get_hd_keypaths(&self.secp));
+                }
                 if params.include_output_redeem_witness_script {
                     psbt_output.witness_script = derived_descriptor.psbt_witness_script();
                     psbt_output.redeem_script = derived_descriptor.psbt_redeem_script();
@@ -1494,17 +1503,21 @@ where
 
         let desc = self.get_descriptor_for_keychain(keychain);
         let derived_descriptor = desc.as_derived(child, &self.secp);
-        psbt_input.bip32_derivation = derived_descriptor.get_hd_keypaths(&self.secp)?;
+        if desc.is_taproot() {
+            psbt_input.tap_key_origins = derived_descriptor.get_tap_key_origins(&self.secp);
+        } else {
+            psbt_input.bip32_derivation = derived_descriptor.get_hd_keypaths(&self.secp);
+        }
 
         psbt_input.redeem_script = derived_descriptor.psbt_redeem_script();
         psbt_input.witness_script = derived_descriptor.psbt_witness_script();
 
         let prev_output = utxo.outpoint;
         if let Some(prev_tx) = self.database.borrow().get_raw_tx(&prev_output.txid)? {
-            if desc.is_witness() {
+            if desc.is_witness() || desc.is_taproot() {
                 psbt_input.witness_utxo = Some(prev_tx.output[prev_output.vout as usize].clone());
             }
-            if !desc.is_witness() || !only_witness_utxo {
+            if !desc.is_taproot() && (!desc.is_witness() || !only_witness_utxo) {
                 psbt_input.non_witness_utxo = Some(prev_tx);
             }
         }
@@ -1530,12 +1543,19 @@ where
                 {
                     debug!("Found descriptor {:?}/{}", keychain, child);
 
-                    // merge hd_keypaths
+                    // merge hd_keypaths or tap_key_origins
                     let desc = self.get_descriptor_for_keychain(keychain);
-                    let mut hd_keypaths = desc
+                    if desc.is_taproot() {
-                        .as_derived(child, &self.secp)
+                        let mut tap_key_origins = desc
-                        .get_hd_keypaths(&self.secp)?;
+                            .as_derived(child, &self.secp)
-                    psbt_input.bip32_derivation.append(&mut hd_keypaths);
+                            .get_tap_key_origins(&self.secp);
+                        psbt_input.tap_key_origins.append(&mut tap_key_origins);
+                    } else {
+                        let mut hd_keypaths = desc
+                            .as_derived(child, &self.secp)
+                            .get_hd_keypaths(&self.secp);
+                        psbt_input.bip32_derivation.append(&mut hd_keypaths);
+                    }
                 }
             }
         }
@@ -1790,6 +1810,10 @@ pub(crate) mod test {
         "wsh(and_v(v:pk(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW),after(100000)))"
     }
 
+    pub(crate) fn get_test_tr_single_sig() -> &'static str {
+        "tr(tprv8ZgxMBicQKsPdDArR4xSAECuVxeX1jwwSXR4ApKbkYgZiziDc4LdBy2WvJeGDfUSE4UT4hHhbgEwbdq8ajjUHiKDegkwrNU6V55CxcxonVN/*)"
+    }
+
     macro_rules! assert_fee_rate {
         ($tx:expr, $fees:expr, $fee_rate:expr $( ,@dust_change $( $dust_change:expr )* )* $( ,@add_signature $( $add_signature:expr )* )* ) => ({
             let mut tx = $tx.clone();
@@ -1819,6 +1843,17 @@ pub(crate) mod test {
         });
     }
 
+    macro_rules! from_str {
+        ($e:expr, $t:ty) => {{
+            use std::str::FromStr;
+            <$t>::from_str($e).unwrap()
+        }};
+
+        ($e:expr) => {
+            from_str!($e, _)
+        };
+    }
+
     #[test]
     #[should_panic(expected = "NoRecipients")]
     fn test_create_tx_empty_recipients() {
@@ -4095,4 +4130,39 @@ pub(crate) mod test {
             "when there's no internal descriptor it should just use external"
         );
     }
+
+    #[test]
+    fn test_taproot_psbt_populate_tap_key_origins() {
+        let (wallet, _, _) = get_funded_wallet(get_test_tr_single_sig());
+        let addr = wallet.get_address(AddressIndex::New).unwrap();
+
+        let mut builder = wallet.build_tx();
+        builder.add_recipient(addr.script_pubkey(), 25_000);
+        let (psbt, _) = builder.finish().unwrap();
+
+        assert_eq!(
+            psbt.inputs[0]
+                .tap_key_origins
+                .clone()
+                .into_iter()
+                .collect::<Vec<_>>(),
+            vec![(
+                from_str!("b96d3a3dc76a4fc74e976511b23aecb78e0754c23c0ed7a6513e18cbbc7178e9"),
+                (vec![], (from_str!("f6a5cb8b"), from_str!("m/0")))
+            )],
+            "Wrong input tap_key_origins"
+        );
+        assert_eq!(
+            psbt.outputs[0]
+                .tap_key_origins
+                .clone()
+                .into_iter()
+                .collect::<Vec<_>>(),
+            vec![(
+                from_str!("e9b03068cf4a2621d4f81e68f6c4216e6bd260fe6edf6acc55c8d8ae5aeff0a8"),
+                (vec![], (from_str!("f6a5cb8b"), from_str!("m/1")))
+            )],
+            "Wrong output tap_key_origins"
+        );
+    }
 }