From c7b6a5aace94a4046855fc9e8e967b8b26591f6b Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Mon, 20 Jul 2020 15:02:05 +0200
Subject: [PATCH] psbt: also check witness UTXO if both are set

A wallet that has patched the CVE-2020-14199 vulnerability will always
include a non-witness UTXO, even for witness inputs. In the signer, we
detect that the input we spend is a witness input and copy over the
TxOut to the witness UTXO field. Therefore it is possible that both UTXO
fields are set at the same time. We need to adjust the sanity checks
when adding a partial signature to account for that.
---
 psbt/updater.go | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/psbt/updater.go b/psbt/updater.go
index b8c2350..01fb2f3 100644
--- a/psbt/updater.go
+++ b/psbt/updater.go
@@ -102,6 +102,11 @@ func (p *Updater) addPartialSignature(inIndex int, sig []byte,
 		}
 	}
 
+	// Attaching signature without utxo field is not allowed.
+	if pInput.WitnessUtxo == nil && pInput.NonWitnessUtxo == nil {
+		return ErrInvalidPsbtFormat
+	}
+
 	// Next, we perform a series of additional sanity checks.
 	if pInput.NonWitnessUtxo != nil {
 		if len(p.Upsbt.UnsignedTx.TxIn) < inIndex+1 {
@@ -136,7 +141,16 @@ func (p *Updater) addPartialSignature(inIndex int, sig []byte,
 			}
 		}
 
-	} else if pInput.WitnessUtxo != nil {
+	}
+
+	// It could be that we set both the non-witness and witness UTXO fields
+	// in case it's from a wallet that patched the CVE-2020-14199
+	// vulnerability. We detect whether the input being spent is actually a
+	// witness input and then copy it over to the witness UTXO field in the
+	// signer. Run the witness checks as well, even if we might already have
+	// checked the script hash. But that should be a negligible performance
+	// penalty.
+	if pInput.WitnessUtxo != nil {
 		scriptPubKey := pInput.WitnessUtxo.PkScript
 
 		var script []byte
@@ -196,10 +210,6 @@ func (p *Updater) addPartialSignature(inIndex int, sig []byte,
 				return ErrInvalidSignatureForInput
 			}
 		}
-	} else {
-
-		// Attaching signature without utxo field is not allowed.
-		return ErrInvalidPsbtFormat
 	}
 
 	p.Upsbt.Inputs[inIndex].PartialSigs = append(