Use tap_key_origins in PSBTs to derive descriptors

src/descriptor/mod.rs

@@ -14,15 +14,15 @@
 //! This module contains generic utilities to work with descriptors, plus some re-exported types
 //! from [`miniscript`].
 
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{BTreeMap, HashSet};
 use std::ops::Deref;
 
-use bitcoin::secp256k1;
 use bitcoin::util::bip32::{ChildNumber, DerivationPath, ExtendedPubKey, Fingerprint, KeySource};
 use bitcoin::util::{psbt, taproot};
+use bitcoin::{secp256k1, PublicKey, XOnlyPublicKey};
 use bitcoin::{Network, Script, TxOut};
 
-use miniscript::descriptor::{DescriptorType, InnerXKey};
+use miniscript::descriptor::{DescriptorType, InnerXKey, SinglePubKey};
 pub use miniscript::{
     descriptor::DescriptorXKey, descriptor::KeyMap, descriptor::Wildcard, Descriptor,
     DescriptorPublicKey, Legacy, Miniscript, ScriptContext, Segwitv0,
@@ -322,6 +322,16 @@ pub(crate) trait DescriptorMeta {
         hd_keypaths: &HdKeyPaths,
         secp: &'s SecpCtx,
     ) -> Option<DerivedDescriptor<'s>>;
+    fn derive_from_tap_key_origins<'s>(
+        &self,
+        tap_key_origins: &TapKeyOrigins,
+        secp: &'s SecpCtx,
+    ) -> Option<DerivedDescriptor<'s>>;
+    fn derive_from_psbt_key_origins<'s>(
+        &self,
+        key_origins: BTreeMap<Fingerprint, (&DerivationPath, SinglePubKey)>,
+        secp: &'s SecpCtx,
+    ) -> Option<DerivedDescriptor<'s>>;
     fn derive_from_psbt_input<'s>(
         &self,
         psbt_input: &psbt::Input,
@@ -401,61 +411,124 @@ impl DescriptorMeta for ExtendedDescriptor {
         Ok(answer)
     }
 
-    fn derive_from_hd_keypaths<'s>(
+    fn derive_from_psbt_key_origins<'s>(
         &self,
-        hd_keypaths: &HdKeyPaths,
+        key_origins: BTreeMap<Fingerprint, (&DerivationPath, SinglePubKey)>,
         secp: &'s SecpCtx,
     ) -> Option<DerivedDescriptor<'s>> {
-        let index: HashMap<_, _> = hd_keypaths.values().map(|(a, b)| (a, b)).collect();
+        // Ensure that deriving `xpub` with `path` yields `expected`
+        let verify_key = |xpub: &DescriptorXKey<ExtendedPubKey>,
+                          path: &DerivationPath,
+                          expected: &SinglePubKey| {
+            let derived = xpub
+                .xkey
+                .derive_pub(secp, path)
+                .expect("The path should never contain hardened derivation steps")
+                .public_key;
+
+            match expected {
+                SinglePubKey::FullKey(pk) if &PublicKey::new(derived) == pk => true,
+                SinglePubKey::XOnly(pk) if &XOnlyPublicKey::from(derived) == pk => true,
+                _ => false,
+            }
+        };
 
         let mut path_found = None;
-        self.for_each_key(|key| {
-            if path_found.is_some() {
-                // already found a matching path, we are done
-                return true;
-            }
 
+        // using `for_any_key` should make this stop as soon as we return `true`
+        self.for_any_key(|key| {
             if let DescriptorPublicKey::XPub(xpub) = key.as_key().deref() {
-                // Check if the key matches one entry in our `index`. If it does, `matches()` will
+                // Check if the key matches one entry in our `key_origins`. If it does, `matches()` will
                 // return the "prefix" that matched, so we remove that prefix from the full path
-                // found in `index` and save it in `derive_path`. We expect this to be a derivation
+                // found in `key_origins` and save it in `derive_path`. We expect this to be a derivation
                 // path of length 1 if the key is `wildcard` and an empty path otherwise.
                 let root_fingerprint = xpub.root_fingerprint(secp);
-                let derivation_path: Option<Vec<ChildNumber>> = index
+                let derive_path = key_origins
                     .get_key_value(&root_fingerprint)
-                    .and_then(|(fingerprint, path)| {
+                    .and_then(|(fingerprint, (path, expected))| {
-                        xpub.matches(&(**fingerprint, (*path).clone()), secp)
+                        xpub.matches(&(*fingerprint, (*path).clone()), secp)
+                            .zip(Some((path, expected)))
                     })
-                    .map(|prefix| {
+                    .and_then(|(prefix, (full_path, expected))| {
-                        index
+                        let derive_path = full_path
-                            .get(&xpub.root_fingerprint(secp))
-                            .unwrap()
                             .into_iter()
                             .skip(prefix.into_iter().count())
                             .cloned()
-                            .collect()
+                            .collect::<DerivationPath>();
+
+                        // `derive_path` only contains the replacement index for the wildcard, if present, or
+                        // an empty path for fixed descriptors. To verify the key we also need the normal steps
+                        // that come before the wildcard, so we take them directly from `xpub` and then append
+                        // the final index
+                        if verify_key(
+                            xpub,
+                            &xpub.derivation_path.extend(derive_path.clone()),
+                            expected,
+                        ) {
+                            Some(derive_path)
+                        } else {
+                            log::debug!(
+                                "Key `{}` derived with {} yields an unexpected key",
+                                root_fingerprint,
+                                derive_path
+                            );
+                            None
+                        }
                     });
 
-                match derivation_path {
+                match derive_path {
                     Some(path) if xpub.wildcard != Wildcard::None && path.len() == 1 => {
                         // Ignore hardened wildcards
                         if let ChildNumber::Normal { index } = path[0] {
-                            path_found = Some(index)
+                            path_found = Some(index);
+                            return true;
                         }
                     }
                     Some(path) if xpub.wildcard == Wildcard::None && path.is_empty() => {
-                        path_found = Some(0)
+                        path_found = Some(0);
+                        return true;
                     }
                     _ => {}
                 }
             }
 
-            true
+            false
         });
 
         path_found.map(|path| self.as_derived(path, secp))
     }
 
+    fn derive_from_hd_keypaths<'s>(
+        &self,
+        hd_keypaths: &HdKeyPaths,
+        secp: &'s SecpCtx,
+    ) -> Option<DerivedDescriptor<'s>> {
+        // "Convert" an hd_keypaths map to the format required by `derive_from_psbt_key_origins`
+        let key_origins = hd_keypaths
+            .iter()
+            .map(|(pk, (fingerprint, path))| {
+                (
+                    *fingerprint,
+                    (path, SinglePubKey::FullKey(PublicKey::new(*pk))),
+                )
+            })
+            .collect();
+        self.derive_from_psbt_key_origins(key_origins, secp)
+    }
+
+    fn derive_from_tap_key_origins<'s>(
+        &self,
+        tap_key_origins: &TapKeyOrigins,
+        secp: &'s SecpCtx,
+    ) -> Option<DerivedDescriptor<'s>> {
+        // "Convert" a tap_key_origins map to the format required by `derive_from_psbt_key_origins`
+        let key_origins = tap_key_origins
+            .iter()
+            .map(|(pk, (_, (fingerprint, path)))| (*fingerprint, (path, SinglePubKey::XOnly(*pk))))
+            .collect();
+        self.derive_from_psbt_key_origins(key_origins, secp)
+    }
+
     fn derive_from_psbt_input<'s>(
         &self,
         psbt_input: &psbt::Input,
@@ -465,6 +538,9 @@ impl DescriptorMeta for ExtendedDescriptor {
         if let Some(derived) = self.derive_from_hd_keypaths(&psbt_input.bip32_derivation, secp) {
             return Some(derived);
         }
+        if let Some(derived) = self.derive_from_tap_key_origins(&psbt_input.tap_key_origins, secp) {
+            return Some(derived);
+        }
         if self.is_deriveable() {
             // We can't try to bruteforce the derivation index, exit here
             return None;

src/wallet/mod.rs

@@ -26,7 +26,8 @@ use bitcoin::secp256k1::Secp256k1;
 use bitcoin::consensus::encode::serialize;
 use bitcoin::util::{psbt, taproot};
 use bitcoin::{
-    Address, EcdsaSighashType, Network, OutPoint, Script, Transaction, TxOut, Txid, Witness,
+    Address, EcdsaSighashType, Network, OutPoint, SchnorrSighashType, Script, Transaction, TxOut,
+    Txid, Witness,
 };
 
 use miniscript::descriptor::DescriptorTrait;
@@ -1013,23 +1014,27 @@ where
         // this helps us doing our job later
         self.add_input_hd_keypaths(psbt)?;
 
-        // If we aren't allowed to use `witness_utxo`, ensure that every input but finalized one
+        // If we aren't allowed to use `witness_utxo`, ensure that every input (except p2tr and finalized ones)
         // has the `non_witness_utxo`
         if !sign_options.trust_witness_utxo
             && psbt
                 .inputs
                 .iter()
                 .filter(|i| i.final_script_witness.is_none() && i.final_script_sig.is_none())
+                .filter(|i| i.tap_internal_key.is_none() && i.tap_merkle_root.is_none())
                 .any(|i| i.non_witness_utxo.is_none())
         {
             return Err(Error::Signer(signer::SignerError::MissingNonWitnessUtxo));
         }
 
         // If the user hasn't explicitly opted-in, refuse to sign the transaction unless every input
-        // is using `SIGHASH_ALL`
+        // is using `SIGHASH_ALL` or `SIGHASH_DEFAULT` for taproot
         if !sign_options.allow_all_sighashes
             && !psbt.inputs.iter().all(|i| {
-                i.sighash_type.is_none() || i.sighash_type == Some(EcdsaSighashType::All.into())
+                i.sighash_type.is_none()
+                    || i.sighash_type == Some(EcdsaSighashType::All.into())
+                    || i.sighash_type == Some(SchnorrSighashType::All.into())
+                    || i.sighash_type == Some(SchnorrSighashType::Default.into())
             })
         {
             return Err(Error::Signer(signer::SignerError::NonStandardSighash));
@@ -1854,7 +1859,7 @@ pub(crate) mod test {
     }
 
     pub(crate) fn get_test_tr_with_taptree() -> &'static str {
-        "tr(cPZzKuNmpuUjD1e8jUU4PVzy2b5LngbSip8mBsxf4e7rSFZVb4Uh,{pk(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55),pk(8aee2b8120a5f157f1223f72b5e62b825831a27a9fdf427db7cc697494d4a642)})"
+        "tr(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55,{pk(cPZzKuNmpuUjD1e8jUU4PVzy2b5LngbSip8mBsxf4e7rSFZVb4Uh),pk(8aee2b8120a5f157f1223f72b5e62b825831a27a9fdf427db7cc697494d4a642)})"
     }
 
     pub(crate) fn get_test_tr_repeated_key() -> &'static str {
@@ -4218,7 +4223,7 @@ pub(crate) mod test {
         let (wallet, _, _) = get_funded_wallet(get_test_tr_repeated_key());
         let addr = wallet.get_address(AddressIndex::New).unwrap();
 
-        let path = vec![("rn4nre9c".to_string(), vec![0])]
+        let path = vec![("u6ugnnck".to_string(), vec![0])]
             .into_iter()
             .collect();
 
@@ -4290,7 +4295,7 @@ pub(crate) mod test {
             psbt.inputs[0].tap_merkle_root,
             Some(
                 FromHex::from_hex(
-                    "d9ca24475ed2c4081ae181d8faa7461649961237bee7bc692f1de448d2d62031"
+                    "61f81509635053e52d9d1217545916167394490da2287aca4693606e43851986"
                 )
                 .unwrap()
             ),
@@ -4298,8 +4303,8 @@ pub(crate) mod test {
         assert_eq!(
             psbt.inputs[0].tap_scripts.clone().into_iter().collect::<Vec<_>>(),
             vec![
-                (taproot::ControlBlock::from_slice(&Vec::<u8>::from_hex("c151494dc22e24a32fe9dcfbd7e85faf345fa1df296fb49d156e859ef3452012958b0afded0ee3149dbf6710d349dc30e55ae30c319c30efbc79efe19cf70f46a8").unwrap()).unwrap(), (from_str!("208aee2b8120a5f157f1223f72b5e62b825831a27a9fdf427db7cc697494d4a642ac"), taproot::LeafVersion::TapScript)),
+                (taproot::ControlBlock::from_slice(&Vec::<u8>::from_hex("c0b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55b7ef769a745e625ed4b9a4982a4dc08274c59187e73e6f07171108f455081cb2").unwrap()).unwrap(), (from_str!("208aee2b8120a5f157f1223f72b5e62b825831a27a9fdf427db7cc697494d4a642ac"), taproot::LeafVersion::TapScript)),
-                (taproot::ControlBlock::from_slice(&Vec::<u8>::from_hex("c151494dc22e24a32fe9dcfbd7e85faf345fa1df296fb49d156e859ef345201295b9a515f7be31a70186e3c5937ee4a70cc4b4e1efe876c1d38e408222ffc64834").unwrap()).unwrap(), (from_str!("20b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55ac"), taproot::LeafVersion::TapScript)),
+                (taproot::ControlBlock::from_slice(&Vec::<u8>::from_hex("c0b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55b9a515f7be31a70186e3c5937ee4a70cc4b4e1efe876c1d38e408222ffc64834").unwrap()).unwrap(), (from_str!("2051494dc22e24a32fe9dcfbd7e85faf345fa1df296fb49d156e859ef345201295ac"), taproot::LeafVersion::TapScript)),
             ],
         );
         assert_eq!(
@@ -4356,4 +4361,34 @@ pub(crate) mod test {
             "foreign_utxo should be in there"
         );
     }
+
+    #[test]
+    fn test_taproot_key_spend() {
+        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 (mut psbt, _) = builder.finish().unwrap();
+
+        assert!(
+            wallet.sign(&mut psbt, Default::default()).unwrap(),
+            "Unable to finalize taproot key spend"
+        );
+    }
+
+    #[test]
+    fn test_taproot_script_spend() {
+        let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree());
+        let addr = wallet.get_address(AddressIndex::New).unwrap();
+
+        let mut builder = wallet.build_tx();
+        builder.add_recipient(addr.script_pubkey(), 25_000);
+        let (mut psbt, _) = builder.finish().unwrap();
+
+        assert!(
+            wallet.sign(&mut psbt, Default::default()).unwrap(),
+            "Unable to finalize taproot script spend"
+        );
+    }
 }

src/wallet/signer.rs

@@ -96,7 +96,7 @@ use bitcoin::{EcdsaSighashType, PrivateKey, PublicKey, SchnorrSighashType, Scrip
 
 use miniscript::descriptor::{
     Descriptor, DescriptorPublicKey, DescriptorSecretKey, DescriptorSinglePriv, DescriptorXKey,
-    KeyMap,
+    KeyMap, SinglePubKey,
 };
 use miniscript::{Legacy, MiniscriptKey, Segwitv0, Tap};
 
@@ -299,19 +299,24 @@ impl InputSigner for SignerWrapper<DescriptorXKey<ExtendedPrivKey>> {
             return Ok(());
         }
 
+        let tap_key_origins = psbt.inputs[input_index]
+            .tap_key_origins
+            .iter()
+            .map(|(pk, (_, keysource))| (SinglePubKey::XOnly(*pk), keysource))
+            .collect::<Vec<_>>();
         let (public_key, full_path) = match psbt.inputs[input_index]
             .bip32_derivation
             .iter()
-            .filter_map(|(pk, &(fingerprint, ref path))| {
+            .map(|(pk, keysource)| (SinglePubKey::FullKey(PublicKey::new(*pk)), keysource))
-                if self.matches(&(fingerprint, path.clone()), secp).is_some() {
+            .chain(tap_key_origins)
-                    Some((pk, path))
+            .find_map(|(pk, keysource)| {
+                if self.matches(keysource, secp).is_some() {
+                    Some((pk, keysource.1.clone()))
                 } else {
                     None
                 }
-            })
+            }) {
-            .next()
+            Some((pk, full_path)) => (pk, full_path),
-        {
-            Some((pk, full_path)) => (pk, full_path.clone()),
             None => return Ok(()),
         };
 
@@ -326,7 +331,13 @@ impl InputSigner for SignerWrapper<DescriptorXKey<ExtendedPrivKey>> {
             None => self.xkey.derive_priv(secp, &full_path).unwrap(),
         };
 
-        if &secp256k1::PublicKey::from_secret_key(secp, &derived_key.private_key) != public_key {
+        let computed_pk = secp256k1::PublicKey::from_secret_key(secp, &derived_key.private_key);
+        let valid_key = match public_key {
+            SinglePubKey::FullKey(pk) if pk.inner == computed_pk => true,
+            SinglePubKey::XOnly(x_only) if XOnlyPublicKey::from(computed_pk) == x_only => true,
+            _ => false,
+        };
+        if !valid_key {
             Err(SignerError::InvalidKey)
         } else {
             // HD wallets imply compressed keys
@@ -488,7 +499,7 @@ fn sign_psbt_schnorr(
             .tap_script_sigs
             .insert((pubkey, lh), final_signature);
     } else {
-        psbt_input.tap_key_sig = Some(final_signature)
+        psbt_input.tap_key_sig = Some(final_signature);
     }
 }