From a14daadcf1c001b5118eb4b30715defdb43f0a9f Mon Sep 17 00:00:00 2001
From: Ava Chow <github@achow101.com>
Date: Tue, 21 Jan 2025 18:02:15 -0500
Subject: [PATCH] psbt: use sighash type field to determine whether to remove
 non-witness utxos

Since the sighash type field is written for atypical sighash types, we
can look at that field to figure out whether the psbt contains
unnecessary transactions.
---
 src/psbt.cpp               | 56 ++++++++++++++++++++------------------
 src/psbt.h                 |  2 +-
 src/rpc/rawtransaction.cpp |  2 +-
 src/wallet/wallet.cpp      |  2 +-
 4 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/src/psbt.cpp b/src/psbt.cpp
index c0926cfc88b33..414215d9003ae 100644
--- a/src/psbt.cpp
+++ b/src/psbt.cpp
@@ -467,37 +467,41 @@ PSBTError SignPSBTInput(const SigningProvider& provider, PartiallySignedTransact
     return sig_complete ? PSBTError::OK : PSBTError::INCOMPLETE;
 }
 
-void RemoveUnnecessaryTransactions(PartiallySignedTransaction& psbtx, std::optional<int> sighash_type)
+void RemoveUnnecessaryTransactions(PartiallySignedTransaction& psbtx)
 {
-    if (!sighash_type) sighash_type = SIGHASH_DEFAULT;
-    // Only drop non_witness_utxos if sighash_type != SIGHASH_ANYONECANPAY
-    if ((*sighash_type & 0x80) != SIGHASH_ANYONECANPAY) {
-        // Figure out if any non_witness_utxos should be dropped
-        std::vector<unsigned int> to_drop;
-        for (unsigned int i = 0; i < psbtx.inputs.size(); ++i) {
-            const auto& input = psbtx.inputs.at(i);
-            int wit_ver;
-            std::vector<unsigned char> wit_prog;
-            if (input.witness_utxo.IsNull() || !input.witness_utxo.scriptPubKey.IsWitnessProgram(wit_ver, wit_prog)) {
-                // There's a non-segwit input or Segwit v0, so we cannot drop any witness_utxos
-                to_drop.clear();
-                break;
-            }
-            if (wit_ver == 0) {
-                // Segwit v0, so we cannot drop any non_witness_utxos
-                to_drop.clear();
-                break;
-            }
-            if (input.non_witness_utxo) {
-                to_drop.push_back(i);
-            }
+    // Figure out if any non_witness_utxos should be dropped
+    std::vector<unsigned int> to_drop;
+    for (unsigned int i = 0; i < psbtx.inputs.size(); ++i) {
+        const auto& input = psbtx.inputs.at(i);
+        int wit_ver;
+        std::vector<unsigned char> wit_prog;
+        if (input.witness_utxo.IsNull() || !input.witness_utxo.scriptPubKey.IsWitnessProgram(wit_ver, wit_prog)) {
+            // There's a non-segwit input, so we cannot drop any non_witness_utxos
+            to_drop.clear();
+            break;
+        }
+        if (wit_ver == 0) {
+            // Segwit v0, so we cannot drop any non_witness_utxos
+            to_drop.clear();
+            break;
+        }
+        // non_witness_utxos cannot be dropped if the sighash type includes SIGHASH_ANYONECANPAY
+        // Since callers should have called SignPSBTInput which updates the sighash type in the PSBT, we only
+        // need to look at that field. If it is not present, then we can assume SIGHASH_DEFAULT or SIGHASH_ALL.
+        if (input.sighash_type != std::nullopt && (*input.sighash_type & 0x80) == SIGHASH_ANYONECANPAY) {
+            to_drop.clear();
+            break;
         }
 
-        // Drop the non_witness_utxos that we can drop
-        for (unsigned int i : to_drop) {
-            psbtx.inputs.at(i).non_witness_utxo = nullptr;
+        if (input.non_witness_utxo) {
+            to_drop.push_back(i);
         }
     }
+
+    // Drop the non_witness_utxos that we can drop
+    for (unsigned int i : to_drop) {
+        psbtx.inputs.at(i).non_witness_utxo = nullptr;
+    }
 }
 
 bool FinalizePSBT(PartiallySignedTransaction& psbtx)
diff --git a/src/psbt.h b/src/psbt.h
index f84d70a53bf18..fecec1f0dd85b 100644
--- a/src/psbt.h
+++ b/src/psbt.h
@@ -1242,7 +1242,7 @@ bool PSBTInputSignedAndVerified(const PartiallySignedTransaction psbt, unsigned
 [[nodiscard]] PSBTError SignPSBTInput(const SigningProvider& provider, PartiallySignedTransaction& psbt, int index, const PrecomputedTransactionData* txdata, std::optional<int> sighash = std::nullopt, SignatureData* out_sigdata = nullptr, bool finalize = true);
 
 /**  Reduces the size of the PSBT by dropping unnecessary `non_witness_utxos` (i.e. complete previous transactions) from a psbt when all inputs are segwit v1. */
-void RemoveUnnecessaryTransactions(PartiallySignedTransaction& psbtx, std::optional<int> sighash_type);
+void RemoveUnnecessaryTransactions(PartiallySignedTransaction& psbtx);
 
 /** Counts the unsigned inputs of a PSBT. */
 size_t CountPSBTUnsignedInputs(const PartiallySignedTransaction& psbt);
diff --git a/src/rpc/rawtransaction.cpp b/src/rpc/rawtransaction.cpp
index 0c9d3540b0542..ae3f791e6579a 100644
--- a/src/rpc/rawtransaction.cpp
+++ b/src/rpc/rawtransaction.cpp
@@ -244,7 +244,7 @@ PartiallySignedTransaction ProcessPSBT(const std::string& psbt_string, const std
         UpdatePSBTOutput(provider, psbtx, i);
     }
 
-    RemoveUnnecessaryTransactions(psbtx, /*sighash_type=*/std::nullopt);
+    RemoveUnnecessaryTransactions(psbtx);
 
     return psbtx;
 }
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index cc095e198ef03..38b5f105bb788 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -2240,7 +2240,7 @@ std::optional<PSBTError> CWallet::FillPSBT(PartiallySignedTransaction& psbtx, bo
         }
     }
 
-    RemoveUnnecessaryTransactions(psbtx, sighash_type);
+    RemoveUnnecessaryTransactions(psbtx);
 
     // Complete if every input is now signed
     complete = true;