Fix descriptor derivation from psbt input

afilini · afilini · commit b6a3322571b6 · 2022-04-29T12:59:09.000+02:00
diff --git a/src/descriptor/mod.rs b/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::{secp256k1, XOnlyPublicKey, PublicKey};
 use bitcoin::util::bip32::{ChildNumber, DerivationPath, ExtendedPubKey, Fingerprint, KeySource};
 use bitcoin::util::{psbt, taproot};
 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,
@@ -327,6 +327,7 @@ pub(crate) trait DescriptorMeta {
         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,
@@ -395,78 +396,101 @@ impl DescriptorMeta for ExtendedDescriptor {
         Ok(answer)
     }
 
-    fn derive_from_hd_keypaths<'s>(
-        &self,
-        hd_keypaths: &HdKeyPaths,
-        secp: &'s SecpCtx,
-    ) -> Option<DerivedDescriptor<'s>> {
-        let index: HashMap<_, _> = hd_keypaths.values().map(|(a, b)| (a, b)).collect();
+    fn derive_from_psbt_key_origins<'s>(&self, key_origins: BTreeMap<Fingerprint, (&DerivationPath, SinglePubKey)>, secp: &'s SecpCtx) -> Option<DerivedDescriptor<'s>> {
+        // 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;
 
-        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;
+            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;
+
+        // 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)| {
-                        xpub.matches(&(**fingerprint, (*path).clone()), secp)
+                    .and_then(|(fingerprint, (path, expected))| {
+                        xpub.matches(&(*fingerprint, (*path).clone()), secp).zip(Some((path, expected)))
                     })
-                    .map(|prefix| {
-                        index
-                            .get(&xpub.root_fingerprint(secp))
-                            .unwrap()
-                            .into_iter()
+                    .and_then(|(prefix, (full_path, expected))| {
+                        let derive_path = full_path.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 an hd_keypaths map, by dropping the leaf hash vector and
-        // converting the public key type
-        let hd_keypaths = tap_key_origins
+        // "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, (_, keysource))| {
-                let mut data: Vec<u8> = vec![0x02];
-                data.extend(pk.serialize().iter());
-                let pk = secp256k1::PublicKey::from_slice(&data).expect("Invalid XOnlyPublicKey");
-                (pk, keysource.clone())
+            .map(|(pk, (_, (fingerprint, path)))| {
+                (*fingerprint, (path, SinglePubKey::XOnly(*pk)))
             })
             .collect();
-        self.derive_from_hd_keypaths(&hd_keypaths, secp)
+        self.derive_from_psbt_key_origins(key_origins, secp)
     }
 
     fn derive_from_psbt_input<'s>(
diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs
@@ -26,7 +26,7 @@ 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, SchnorrSighashType, Network, OutPoint, Script, Transaction, TxOut, Txid, Witness,
 };
 
 use miniscript::descriptor::DescriptorTrait;
@@ -1009,23 +1009,24 @@ 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
-        // has the `non_witness_utxo`
+        // If we aren't allowed to use `witness_utxo`, ensure that every input (except p2tr and finalized ones)
+        // have 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));
@@ -1818,7 +1819,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)})"
     }
 
     macro_rules! assert_fee_rate {
@@ -4163,6 +4164,33 @@ pub(crate) mod test {
         );
         assert_eq!(psbt.inputs[0].tap_scripts.len(), 2, "Empty tap_scripts");
     }
+
+    #[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");
+        wallet.sign(&mut psbt, Default::default()).unwrap();
+
+        dbg!(&psbt);
+    }
+
+    #[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");
+    }
 }
 
 /// Deterministically generate a unique name given the descriptors defining the wallet
diff --git a/src/wallet/signer.rs b/src/wallet/signer.rs
@@ -94,7 +94,7 @@ use bitcoin::util::{ecdsa, psbt, schnorr, sighash, taproot};
 use bitcoin::{secp256k1, XOnlyPublicKey};
 use bitcoin::{EcdsaSighashType, PrivateKey, PublicKey, SchnorrSighashType, Script};
 
-use miniscript::descriptor::{DescriptorSecretKey, DescriptorSinglePriv, DescriptorXKey, KeyMap};
+use miniscript::descriptor::{DescriptorSecretKey, DescriptorSinglePriv, DescriptorXKey, SinglePubKey, KeyMap};
 use miniscript::{Legacy, MiniscriptKey, Segwitv0, Tap};
 
 use super::utils::SecpCtx;
@@ -258,11 +258,14 @@ impl InputSigner for 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))| {
-                if self.matches(&(fingerprint, path.clone()), secp).is_some() {
+            .map(|(pk, keysource)| (SinglePubKey::FullKey(PublicKey::new(*pk)), keysource))
+            .chain(tap_key_origins.into_iter())
+            .filter_map(|(pk, keysource @ (_, path))| {
+                if self.matches(&keysource, secp).is_some() {
                     Some((pk, path))
                 } else {
                     None
@@ -285,7 +288,13 @@ impl InputSigner for 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
@@ -440,7 +449,7 @@ fn produce_schnorr_sig(
             .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);
     }
 }
 
