Merge bitcoindevkit/bdk#645: Allow signing only specific leaf hashes

a713a5a0629c9a643708a4b0d8d6ac3a87a13195 Better customize signing in taproot transactions (Daniela Brozzoni)

Pull request description:

  Fixes #616

  ### Checklists

  #### All Submissions:

  * [x] I've signed all my commits
  * [x] I followed the [contribution guidelines](https://github.com/bitcoindevkit/bdk/blob/master/CONTRIBUTING.md)
  * [x] I ran `cargo fmt` and `cargo clippy` before committing

  #### New Features:

  * [x] I've added tests for the new feature
  * [x] I've added docs for the new feature
  * [x] I've updated `CHANGELOG.md`

ACKs for top commit:
  afilini:
    ACK a713a5a0629c9a643708a4b0d8d6ac3a87a13195

Tree-SHA512: 1100d43cb394b429450fc34f49dd815a024701987c0e6dd163865bd5c4c6f7102127b1ea6e10ced5fdb319874be97baeeb0deea66b4138410871a1d68b4def10

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `Excess` enum to handle remaining amount after coin selection.
 - Move change creation from `Wallet::create_tx` to `CoinSelectionAlgorithm::coin_select`.
 - Change the interface of `SqliteDatabase::new` to accept any type that implement AsRef<Path>
+- Add the ability to specify which leaves to sign in a taproot transaction through `TapLeavesOptions` in `SignOptions`
+- Add the ability to specify whether a taproot transaction should be signed using the internal key or not, using `sign_with_tap_internal_key` in `SignOptions`
 
 ## [v0.20.0] - [v0.19.0]
 

src/wallet/mod.rs

@@ -1086,7 +1086,7 @@ where
             .iter()
             .chain(self.change_signers.signers().iter())
         {
-            signer.sign_transaction(psbt, &self.secp)?;
+            signer.sign_transaction(psbt, &sign_options, &self.secp)?;
         }
 
         // attempt to finalize
@@ -1975,6 +1975,10 @@ pub(crate) mod test {
         "tr(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55,{pk(cPZzKuNmpuUjD1e8jUU4PVzy2b5LngbSip8mBsxf4e7rSFZVb4Uh),pk(8aee2b8120a5f157f1223f72b5e62b825831a27a9fdf427db7cc697494d4a642)})"
     }
 
+    pub(crate) fn get_test_tr_with_taptree_both_priv() -> &'static str {
+        "tr(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55,{pk(cPZzKuNmpuUjD1e8jUU4PVzy2b5LngbSip8mBsxf4e7rSFZVb4Uh),pk(cNaQCDwmmh4dS9LzCgVtyy1e1xjCJ21GUDHe9K98nzb689JvinGV)})"
+    }
+
     pub(crate) fn get_test_tr_repeated_key() -> &'static str {
         "tr(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55,{and_v(v:pk(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW),after(100)),and_v(v:pk(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW),after(200))})"
     }
@@ -1984,7 +1988,7 @@ pub(crate) mod test {
     }
 
     pub(crate) fn get_test_tr_with_taptree_xprv() -> &'static str {
-        "tr(b511bd5771e47ee27558b1765e87b541668304ec567721c7b880edc0a010da55,{pk(tprv8ZgxMBicQKsPdDArR4xSAECuVxeX1jwwSXR4ApKbkYgZiziDc4LdBy2WvJeGDfUSE4UT4hHhbgEwbdq8ajjUHiKDegkwrNU6V55CxcxonVN/*),pk(8aee2b8120a5f157f1223f72b5e62b825831a27a9fdf427db7cc697494d4a642)})"
+        "tr(cNJmN3fH9DDbDt131fQNkVakkpzawJBSeybCUNmP1BovpmGQ45xG,{pk(tprv8ZgxMBicQKsPdDArR4xSAECuVxeX1jwwSXR4ApKbkYgZiziDc4LdBy2WvJeGDfUSE4UT4hHhbgEwbdq8ajjUHiKDegkwrNU6V55CxcxonVN/*),pk(8aee2b8120a5f157f1223f72b5e62b825831a27a9fdf427db7cc697494d4a642)})"
     }
 
     macro_rules! assert_fee_rate {
@@ -4830,6 +4834,31 @@ pub(crate) mod test {
         test_spend_from_wallet(wallet);
     }
 
+    #[test]
+    fn test_taproot_no_key_spend() {
+        let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
+        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,
+                    SignOptions {
+                        sign_with_tap_internal_key: false,
+                        ..Default::default()
+                    },
+                )
+                .unwrap(),
+            "Unable to finalize tx"
+        );
+
+        assert!(psbt.inputs.iter().all(|i| i.tap_key_sig.is_none()));
+    }
+
     #[test]
     fn test_taproot_script_spend() {
         let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree());
@@ -4839,6 +4868,142 @@ pub(crate) mod test {
         test_spend_from_wallet(wallet);
     }
 
+    #[test]
+    fn test_taproot_script_spend_sign_all_leaves() {
+        use crate::signer::TapLeavesOptions;
+        let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
+        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,
+                    SignOptions {
+                        tap_leaves_options: TapLeavesOptions::All,
+                        ..Default::default()
+                    },
+                )
+                .unwrap(),
+            "Unable to finalize tx"
+        );
+
+        assert!(psbt
+            .inputs
+            .iter()
+            .all(|i| i.tap_script_sigs.len() == i.tap_scripts.len()));
+    }
+
+    #[test]
+    fn test_taproot_script_spend_sign_include_some_leaves() {
+        use crate::signer::TapLeavesOptions;
+        use crate::wallet::taproot::TapLeafHash;
+
+        let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
+        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();
+        let mut script_leaves: Vec<_> = psbt.inputs[0]
+            .tap_scripts
+            .clone()
+            .values()
+            .map(|(script, version)| TapLeafHash::from_script(script, *version))
+            .collect();
+        let included_script_leaves = vec![script_leaves.pop().unwrap()];
+        let excluded_script_leaves = script_leaves;
+
+        assert!(
+            wallet
+                .sign(
+                    &mut psbt,
+                    SignOptions {
+                        tap_leaves_options: TapLeavesOptions::Include(
+                            included_script_leaves.clone()
+                        ),
+                        ..Default::default()
+                    },
+                )
+                .unwrap(),
+            "Unable to finalize tx"
+        );
+
+        assert!(psbt.inputs[0]
+            .tap_script_sigs
+            .iter()
+            .all(|s| included_script_leaves.contains(&s.0 .1)
+                && !excluded_script_leaves.contains(&s.0 .1)));
+    }
+
+    #[test]
+    fn test_taproot_script_spend_sign_exclude_some_leaves() {
+        use crate::signer::TapLeavesOptions;
+        use crate::wallet::taproot::TapLeafHash;
+
+        let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
+        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();
+        let mut script_leaves: Vec<_> = psbt.inputs[0]
+            .tap_scripts
+            .clone()
+            .values()
+            .map(|(script, version)| TapLeafHash::from_script(script, *version))
+            .collect();
+        let included_script_leaves = vec![script_leaves.pop().unwrap()];
+        let excluded_script_leaves = script_leaves;
+
+        assert!(
+            wallet
+                .sign(
+                    &mut psbt,
+                    SignOptions {
+                        tap_leaves_options: TapLeavesOptions::Exclude(
+                            excluded_script_leaves.clone()
+                        ),
+                        ..Default::default()
+                    },
+                )
+                .unwrap(),
+            "Unable to finalize tx"
+        );
+
+        assert!(psbt.inputs[0]
+            .tap_script_sigs
+            .iter()
+            .all(|s| included_script_leaves.contains(&s.0 .1)
+                && !excluded_script_leaves.contains(&s.0 .1)));
+    }
+
+    #[test]
+    fn test_taproot_script_spend_sign_no_leaves() {
+        use crate::signer::TapLeavesOptions;
+        let (wallet, _, _) = get_funded_wallet(get_test_tr_with_taptree_both_priv());
+        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();
+
+        wallet
+            .sign(
+                &mut psbt,
+                SignOptions {
+                    tap_leaves_options: TapLeavesOptions::None,
+                    ..Default::default()
+                },
+            )
+            .unwrap();
+
+        assert!(psbt.inputs.iter().all(|i| i.tap_script_sigs.is_empty()));
+    }
+
     #[test]
     fn test_taproot_sign_derive_index_from_psbt() {
         let (wallet, _, _) = get_funded_wallet(get_test_tr_single_sig_xprv());

src/wallet/signer.rs

@@ -58,6 +58,7 @@
 //!         &self,
 //!         psbt: &mut psbt::PartiallySignedTransaction,
 //!         input_index: usize,
+//!         _sign_options: &SignOptions,
 //!         _secp: &Secp256k1<All>,
 //!     ) -> Result<(), SignerError> {
 //!         self.device.hsm_sign_input(psbt, input_index)?;
@@ -241,6 +242,7 @@ pub trait InputSigner: SignerCommon {
         &self,
         psbt: &mut psbt::PartiallySignedTransaction,
         input_index: usize,
+        sign_options: &SignOptions,
         secp: &SecpCtx,
     ) -> Result<(), SignerError>;
 }
@@ -254,6 +256,7 @@ pub trait TransactionSigner: SignerCommon {
     fn sign_transaction(
         &self,
         psbt: &mut psbt::PartiallySignedTransaction,
+        sign_options: &SignOptions,
         secp: &SecpCtx,
     ) -> Result<(), SignerError>;
 }
@@ -262,10 +265,11 @@ impl<T: InputSigner> TransactionSigner for T {
     fn sign_transaction(
         &self,
         psbt: &mut psbt::PartiallySignedTransaction,
+        sign_options: &SignOptions,
         secp: &SecpCtx,
     ) -> Result<(), SignerError> {
         for input_index in 0..psbt.inputs.len() {
-            self.sign_input(psbt, input_index, secp)?;
+            self.sign_input(psbt, input_index, sign_options, secp)?;
         }
 
         Ok(())
@@ -287,6 +291,7 @@ impl InputSigner for SignerWrapper<DescriptorXKey<ExtendedPrivKey>> {
         &self,
         psbt: &mut psbt::PartiallySignedTransaction,
         input_index: usize,
+        sign_options: &SignOptions,
         secp: &SecpCtx,
     ) -> Result<(), SignerError> {
         if input_index >= psbt.inputs.len() {
@@ -346,7 +351,7 @@ impl InputSigner for SignerWrapper<DescriptorXKey<ExtendedPrivKey>> {
                 inner: derived_key.private_key,
             };
 
-            SignerWrapper::new(priv_key, self.ctx).sign_input(psbt, input_index, secp)
+            SignerWrapper::new(priv_key, self.ctx).sign_input(psbt, input_index, sign_options, secp)
         }
     }
 }
@@ -369,6 +374,7 @@ impl InputSigner for SignerWrapper<PrivateKey> {
         &self,
         psbt: &mut psbt::PartiallySignedTransaction,
         input_index: usize,
+        sign_options: &SignOptions,
         secp: &SecpCtx,
     ) -> Result<(), SignerError> {
         if input_index >= psbt.inputs.len() || input_index >= psbt.unsigned_tx.input.len() {
@@ -385,7 +391,10 @@ impl InputSigner for SignerWrapper<PrivateKey> {
         let x_only_pubkey = XOnlyPublicKey::from(pubkey.inner);
 
         if let SignerContext::Tap { is_internal_key } = self.ctx {
-            if is_internal_key && psbt.inputs[input_index].tap_key_sig.is_none() {
+            if is_internal_key
+                && psbt.inputs[input_index].tap_key_sig.is_none()
+                && sign_options.sign_with_tap_internal_key
+            {
                 let (hash, hash_ty) = Tap::sighash(psbt, input_index, None)?;
                 sign_psbt_schnorr(
                     &self.inner,
@@ -404,9 +413,18 @@ impl InputSigner for SignerWrapper<PrivateKey> {
                 let leaf_hashes = leaf_hashes
                     .iter()
                     .filter(|lh| {
-                        !psbt.inputs[input_index]
+                        // Removing the leaves we shouldn't sign for
-                            .tap_script_sigs
+                        let should_sign = match &sign_options.tap_leaves_options {
-                            .contains_key(&(x_only_pubkey, **lh))
+                            TapLeavesOptions::All => true,
+                            TapLeavesOptions::Include(v) => v.contains(lh),
+                            TapLeavesOptions::Exclude(v) => !v.contains(lh),
+                            TapLeavesOptions::None => false,
+                        };
+                        // Filtering out the leaves without our key
+                        should_sign
+                            && !psbt.inputs[input_index]
+                                .tap_script_sigs
+                                .contains_key(&(x_only_pubkey, **lh))
                     })
                     .cloned()
                     .collect::<Vec<_>>();
@@ -677,6 +695,38 @@ pub struct SignOptions {
     ///
     /// Defaults to `true` which will try fianlizing psbt after inputs are signed.
     pub try_finalize: bool,
+
+    /// Specifies which Taproot script-spend leaves we should sign for. This option is
+    /// ignored if we're signing a non-taproot PSBT.
+    ///
+    /// Defaults to All, i.e., the wallet will sign all the leaves it has a key for.
+    pub tap_leaves_options: TapLeavesOptions,
+
+    /// Whether we should try to sign a taproot transaction with the taproot internal key
+    /// or not. This option is ignored if we're signing a non-taproot PSBT.
+    ///
+    /// Defaults to `true`, i.e., we always try to sign with the taproot internal key.
+    pub sign_with_tap_internal_key: bool,
+}
+
+/// Customize which taproot script-path leaves the signer should sign.
+#[derive(Debug, Clone, PartialEq)]
+pub enum TapLeavesOptions {
+    /// The signer will sign all the leaves it has a key for.
+    All,
+    /// The signer won't sign leaves other than the ones specified. Note that it could still ignore
+    /// some of the specified leaves, if it doesn't have the right key to sign them.
+    Include(Vec<taproot::TapLeafHash>),
+    /// The signer won't sign the specified leaves.
+    Exclude(Vec<taproot::TapLeafHash>),
+    /// The signer won't sign any leaf.
+    None,
+}
+
+impl Default for TapLeavesOptions {
+    fn default() -> Self {
+        TapLeavesOptions::All
+    }
 }
 
 #[allow(clippy::derivable_impls)]
@@ -688,6 +738,8 @@ impl Default for SignOptions {
             allow_all_sighashes: false,
             remove_partial_sigs: true,
             try_finalize: true,
+            tap_leaves_options: TapLeavesOptions::default(),
+            sign_with_tap_internal_key: true,
         }
     }
 }
@@ -1018,6 +1070,7 @@ mod signers_container_tests {
         fn sign_transaction(
             &self,
             _psbt: &mut psbt::PartiallySignedTransaction,
+            _sign_options: &SignOptions,
             _secp: &SecpCtx,
         ) -> Result<(), SignerError> {
             Ok(())