From c2e4ba8cbd77f1a41d269b621b0001d042c1ea57 Mon Sep 17 00:00:00 2001
From: KaFai Choi <kafai@river.com>
Date: Sun, 29 May 2022 10:53:37 +0700
Subject: [PATCH] add remove_partial_sigs to SignOptions

---
 CHANGELOG.md         |  1 +
 src/wallet/mod.rs    | 33 +++++++++++++++++++++++++++++++++
 src/wallet/signer.rs |  5 +++++
 3 files changed, 39 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7eaf9bff..e89ad8dd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - Signing Taproot PSBTs (key spend and script spend)
   - Support for `tr()` descriptors in the `descriptor!()` macro
 - Add support for Bitcoin Core 23.0 when using the `rpc` blockchain
+- Add `remove_partial_sigs` to `SignOptions`
 
 ## [v0.18.0] - [v0.17.0]
 
diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs
index 6d2ff385..5adcf831 100644
--- a/src/wallet/mod.rs
+++ b/src/wallet/mod.rs
@@ -1161,6 +1161,9 @@ where
                             let psbt_input = &mut psbt.inputs[n];
                             psbt_input.final_script_sig = Some(tmp_input.script_sig);
                             psbt_input.final_script_witness = Some(tmp_input.witness);
+                            if sign_options.remove_partial_sigs {
+                                psbt_input.partial_sigs.clear();
+                            }
                         }
                         Err(e) => {
                             debug!("satisfy error {:?} for input {}", e, n);
@@ -3878,6 +3881,36 @@ pub(crate) mod test {
         )
     }
 
+    #[test]
+    fn test_remove_partial_sigs_after_finalize_sign_option() {
+        let (wallet, _, _) = get_funded_wallet("wpkh(tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*)");
+
+        for remove_partial_sigs in &[true, false] {
+            let addr = wallet.get_address(New).unwrap();
+            let mut builder = wallet.build_tx();
+            builder.drain_to(addr.script_pubkey()).drain_wallet();
+            let mut psbt = builder.finish().unwrap().0;
+
+            assert!(wallet
+                .sign(
+                    &mut psbt,
+                    SignOptions {
+                        remove_partial_sigs: *remove_partial_sigs,
+                        ..Default::default()
+                    },
+                )
+                .unwrap());
+
+            psbt.inputs.iter().for_each(|input| {
+                if *remove_partial_sigs {
+                    assert!(input.partial_sigs.is_empty())
+                } else {
+                    assert!(!input.partial_sigs.is_empty())
+                }
+            });
+        }
+    }
+
     #[test]
     fn test_sign_nonstandard_sighash() {
         let sighash = EcdsaSighashType::NonePlusAnyoneCanPay;
diff --git a/src/wallet/signer.rs b/src/wallet/signer.rs
index 96e3034e..2267aae3 100644
--- a/src/wallet/signer.rs
+++ b/src/wallet/signer.rs
@@ -667,6 +667,10 @@ pub struct SignOptions {
     ///
     /// Defaults to `false` which will only allow signing using `SIGHASH_ALL`.
     pub allow_all_sighashes: bool,
+    /// Whether to remove partial_sigs from psbt inputs while finalizing psbt.
+    ///
+    /// Defaults to `true` which will remove partial_sigs after finalizing.
+    pub remove_partial_sigs: bool,
 }
 
 #[allow(clippy::derivable_impls)]
@@ -676,6 +680,7 @@ impl Default for SignOptions {
             trust_witness_utxo: false,
             assume_height: None,
             allow_all_sighashes: false,
+            remove_partial_sigs: true,
         }
     }
 }