psbt: Remove non_witness_utxo for segwit v1+

If all inputs are segwit v1+, the non_witness_utxos can be removed.

2 files changed, 32 insertions, 3 deletions

diff --git a/src/psbt.cpp b/src/psbt.cpp
index e4bac4aa63..98fde4cb04 100644
--- a/src/psbt.cpp
+++ b/src/psbt.cpp
@@ -184,7 +184,6 @@ void PSBTInput::Merge(const PSBTInput& input)
 {
     if (!non_witness_utxo && input.non_witness_utxo) non_witness_utxo = input.non_witness_utxo;
     if (witness_utxo.IsNull() && !input.witness_utxo.IsNull()) {
-        // TODO: For segwit v1, we will want to clear out the non-witness utxo when setting a witness one. For v0 and non-segwit, this is not safe
         witness_utxo = input.witness_utxo;
     }
 
@@ -367,10 +366,11 @@ bool SignPSBTInput(const SigningProvider& provider, PartiallySignedTransaction&
     input.FromSignatureData(sigdata);
 
     // If we have a witness signature, put a witness UTXO.
-    // TODO: For segwit v1, we should remove the non_witness_utxo
     if (sigdata.witness) {
         input.witness_utxo = utxo;
-        // input.non_witness_utxo = nullptr;
+        // We can remove the non_witness_utxo if and only if there are no non-segwit or segwit v0
+        // inputs in this transaction. Since this requires inspecting the entire transaction, this
+        // is something for the caller to deal with (i.e. FillPSBT).
     }
 
     // Fill in the missing info
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index d0b093bbb7..041481559b 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -2006,6 +2006,35 @@ TransactionError CWallet::FillPSBT(PartiallySignedTransaction& psbtx, bool& comp
         }
     }
 
+    // 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);
+            }
+        }
+
+        // Drop the non_witness_utxos that we can drop
+        for (unsigned int i : to_drop) {
+            psbtx.inputs.at(i).non_witness_utxo = nullptr;
+        }
+    }
+
     // Complete if every input is now signed
     complete = true;
     for (const auto& input : psbtx.inputs) {