From 5840ce473e430de4c4e3698734e9667cc476fee4 Mon Sep 17 00:00:00 2001
From: vmammal <valuedmammal@protonmail.com>
Date: Wed, 31 Jan 2024 16:18:09 -0500
Subject: [PATCH] fix(bdk): Remove extra taproot fields when finalizing Psbt

We currently allow removing `partial_sigs` from a finalized Psbt,
which is relevant to non-taproot inputs, however taproot related Psbt
fields were left in place despite the recommendation of BIP371 to remove
them once the `final_script_witness` is constructed. This can cause
confusion for parsers that encounter extra taproot metadata in an
already satisfied input.

Fix this by introducing a new member to SignOptions
`remove_taproot_extras`, which when true will remove extra taproot
related data from a Psbt upon successful finalization. This change
makes removal of all taproot extras the default but configurable.

test(wallet): Add test
`test_taproot_remove_tapfields_after_finalize_sign_option`
that checks various fields have been cleared for taproot
Psbt `Input`s and `Output`s according to BIP371.
---
 crates/bdk/src/wallet/mod.rs    | 15 +++++++++++++++
 crates/bdk/src/wallet/signer.rs | 11 +++++++++++
 crates/bdk/tests/wallet.rs      | 26 ++++++++++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/crates/bdk/src/wallet/mod.rs b/crates/bdk/src/wallet/mod.rs
index c6c922d5..406f60e0 100644
--- a/crates/bdk/src/wallet/mod.rs
+++ b/crates/bdk/src/wallet/mod.rs
@@ -1972,6 +1972,15 @@ impl<D> Wallet<D> {
                             if sign_options.remove_partial_sigs {
                                 psbt_input.partial_sigs.clear();
                             }
+                            if sign_options.remove_taproot_extras {
+                                // We just constructed the final witness, clear these fields.
+                                psbt_input.tap_key_sig = None;
+                                psbt_input.tap_script_sigs.clear();
+                                psbt_input.tap_scripts.clear();
+                                psbt_input.tap_key_origins.clear();
+                                psbt_input.tap_internal_key = None;
+                                psbt_input.tap_merkle_root = None;
+                            }
                         }
                         Err(_) => finished = false,
                     }
@@ -1980,6 +1989,12 @@ impl<D> Wallet<D> {
             }
         }
 
+        if finished && sign_options.remove_taproot_extras {
+            for output in &mut psbt.outputs {
+                output.tap_key_origins.clear();
+            }
+        }
+
         Ok(finished)
     }
 
diff --git a/crates/bdk/src/wallet/signer.rs b/crates/bdk/src/wallet/signer.rs
index 63784a3f..38e51dfb 100644
--- a/crates/bdk/src/wallet/signer.rs
+++ b/crates/bdk/src/wallet/signer.rs
@@ -782,6 +782,16 @@ pub struct SignOptions {
     /// Defaults to `true` which will remove partial signatures during finalization.
     pub remove_partial_sigs: bool,
 
+    /// Whether to remove taproot specific fields from the PSBT on finalization.
+    ///
+    /// For inputs this includes the taproot internal key, merkle root, and individual
+    /// scripts and signatures. For both inputs and outputs it includes key origin info.
+    ///
+    /// Defaults to `true` which will remove all of the above mentioned fields when finalizing.
+    ///
+    /// See [`BIP371`](https://github.com/bitcoin/bips/blob/master/bip-0371.mediawiki) for details.
+    pub remove_taproot_extras: bool,
+
     /// Whether to try finalizing the PSBT after the inputs are signed.
     ///
     /// Defaults to `true` which will try finalizing PSBT after inputs are signed.
@@ -827,6 +837,7 @@ impl Default for SignOptions {
             assume_height: None,
             allow_all_sighashes: false,
             remove_partial_sigs: true,
+            remove_taproot_extras: true,
             try_finalize: true,
             tap_leaves_options: TapLeavesOptions::default(),
             sign_with_tap_internal_key: true,
diff --git a/crates/bdk/tests/wallet.rs b/crates/bdk/tests/wallet.rs
index 271b8716..c2f7d40f 100644
--- a/crates/bdk/tests/wallet.rs
+++ b/crates/bdk/tests/wallet.rs
@@ -2813,6 +2813,32 @@ fn test_get_address_no_reuse_single_descriptor() {
     });
 }
 
+#[test]
+fn test_taproot_remove_tapfields_after_finalize_sign_option() {
+    let (mut wallet, _) = get_funded_wallet(get_test_tr_with_taptree());
+
+    let addr = wallet.get_address(New);
+    let mut builder = wallet.build_tx();
+    builder.drain_to(addr.script_pubkey()).drain_wallet();
+    let mut psbt = builder.finish().unwrap();
+    let finalized = wallet.sign(&mut psbt, SignOptions::default()).unwrap();
+    assert!(finalized);
+
+    // removes tap_* from inputs
+    for input in &psbt.inputs {
+        assert!(input.tap_key_sig.is_none());
+        assert!(input.tap_script_sigs.is_empty());
+        assert!(input.tap_scripts.is_empty());
+        assert!(input.tap_key_origins.is_empty());
+        assert!(input.tap_internal_key.is_none());
+        assert!(input.tap_merkle_root.is_none());
+    }
+    // removes key origins from outputs
+    for output in &psbt.outputs {
+        assert!(output.tap_key_origins.is_empty());
+    }
+}
+
 #[test]
 fn test_taproot_psbt_populate_tap_key_origins() {
     let (mut wallet, _) = get_funded_wallet(get_test_tr_single_sig_xprv());