Additional checks for nonWitnessUtxo

Author: NicolasDorier

NBitcoin.Tests/PSBTTests.cs

@@ -133,6 +133,13 @@ public void CantCalculateUnsafeFee()
 			psbt.AddTransactions(funding);
 			Assert.True(psbt.TryGetFee(out _));
 			Assert.True(psbt.TryGetVirtualSize(out _));
+
+			psbt.Inputs[0].NonWitnessUtxo = tx;
+			Assert.False(psbt.TryGetFee(out _));
+			Assert.NotEmpty(psbt.Inputs[0].CheckSanity());
+			psbt.Inputs[0].NonWitnessUtxo = funding;
+			Assert.True(psbt.TryGetFee(out _));
+			Assert.Empty(psbt.Inputs[0].CheckSanity());
 		}
 
 		[Fact]
@@ -265,12 +272,8 @@ public void CanUpdate()
 			var psbtWithCoins = PSBT.FromTransaction(tx, Network.Main)
 				.AddCoins(coins);
 
-			Assert.Null(psbtWithCoins.Inputs[0].WitnessUtxo);
-			Assert.NotNull(psbtWithCoins.Inputs[1].WitnessUtxo);
-			Assert.Null(psbtWithCoins.Inputs[2].WitnessUtxo);
-			Assert.NotNull(psbtWithCoins.Inputs[3].WitnessUtxo);
-			Assert.NotNull(psbtWithCoins.Inputs[4].WitnessUtxo);
-			Assert.NotNull(psbtWithCoins.Inputs[5].WitnessUtxo);
+			foreach (var input in psbtWithCoins.Inputs)
+				Assert.NotNull(input.WitnessUtxo);
 
 			// Check if it holds scripts as expected.
 			Assert.Null(psbtWithCoins.Inputs[0].RedeemScript); // p2pkh

NBitcoin/BIP174/PSBT0.cs

@@ -20,6 +20,7 @@ internal PSBT0(Transaction transaction, Network network) : this(CreateMap(transa
 		if (transaction == null)
 			throw new ArgumentNullException(nameof(transaction));
 		tx = transaction.Clone();
+		tx.PrecomputeHash(true, true);
 	}
 
 	private static Maps CreateMap(Transaction transaction, Network network)
@@ -81,6 +82,7 @@ internal PSBT0(Maps maps, Network network) : base(maps, network, PSBTVersion.PSB
 		if (!maps.Global.TryRemove<byte[]>([PSBTConstants.PSBT_GLOBAL_UNSIGNED_TX], out var txBytes))
 			throw new FormatException("Invalid PSBT. No global TX");
 		tx = Transaction.Load(txBytes, Network);
+		tx.PrecomputeHash(true, true);
 		if (tx.Inputs.Any(txin => txin.ScriptSig != Script.Empty || txin.WitScript != WitScript.Empty))
 			throw new FormatException("Malformed global tx. It should not contain any scriptsig or witness by itself");
 		if (tx.Inputs.Count + tx.Outputs.Count + 1 != maps.Count)
@@ -110,11 +112,11 @@ internal override void FillMap(Map map)
 		map.Add([PSBTConstants.PSBT_GLOBAL_UNSIGNED_TX], this.GetGlobalTransaction(true).ToBytes());
 	}
 
-	Transaction tx;
+	readonly Transaction tx;
 
 	internal override Transaction GetGlobalTransaction(bool @unsafe) => @unsafe ? tx : tx.Clone();
 
-	class PSBT0Input : PSBTInput
+	internal class PSBT0Input : PSBTInput
 	{
 		internal PSBT0Input(Map map, PSBT0 parent, uint index) : base(map, parent, index)
 		{
@@ -132,7 +134,11 @@ internal PSBT0Input(Map map, PSBT0 parent, uint index) : base(map, parent, index
 		public override Sequence Sequence
 		{
 			get => txIn.Sequence;
-			set => txIn.Sequence = value;
+			set
+			{
+				txIn.Sequence = value;
+				((PSBT0)Parent).tx.PrecomputeHash(true, true);
+			}
 		}
 	}
 }

NBitcoin/BIP174/PSBTInput.cs

@@ -185,6 +185,11 @@ protected virtual void Load(Map map)
 		internal Transaction GetTransaction() => Parent.GetGlobalTransaction(true);
 
 		private Transaction? non_witness_utxo;
+		// Record the hash of the last time we checked the non_witness_utxo correspond to the outpoint of the input
+		// If NonWitnessUTXO change, we unset this.
+		// If the Outpoint change, then non_witness_utxo_check != prevout.Hash
+		// This is optimization to not have to calculate the hash of non_witness_utxo every time.
+		internal uint256? non_witness_utxo_check;
 		private TxOut? witness_utxo;
 		private PartialSigKVMap partial_sigs = new PartialSigKVMap(PubKeyComparer.Instance);
 
@@ -199,6 +204,7 @@ public Transaction? NonWitnessUtxo
 			set
 			{
 				non_witness_utxo = value;
+				non_witness_utxo_check = null;
 			}
 		}
 
@@ -535,6 +541,8 @@ public IList<PSBTError> CheckSanity()
 					errors.Add(new PSBTError(Index, "non_witness_utxo does not match the transaction id referenced by the global transaction sign"));
 					validOutpoint = false;
 				}
+				else
+					non_witness_utxo_check = prevOutTxId;
 				if (PrevOut.N >= NonWitnessUtxo.Outputs.Count)
 				{
 					errors.Add(new PSBTError(Index, "Global transaction referencing an out of bound output in non_witness_utxo"));
@@ -797,6 +805,13 @@ private string GetName(uint sighashType)
 			{
 				if (PrevOut.N >= NonWitnessUtxo.Outputs.Count)
 					return null;
+				if (non_witness_utxo_check != PrevOut.Hash)
+				{
+					if (NonWitnessUtxo.GetHash() != PrevOut.Hash)
+						return null;
+					non_witness_utxo_check = PrevOut.Hash;
+				}
+				non_witness_utxo_check = PrevOut.Hash;
 				return NonWitnessUtxo.Outputs[PrevOut.N];
 			}
 			if (WitnessUtxo != null)

NBitcoin/BIP174/PartiallySignedTransaction.cs

@@ -313,6 +313,8 @@ public PSBT AddTransactions(params Transaction[] parentTransactions)
 						continue;
 					var output = tx.Outputs[input.PrevOut.N];
 					input.NonWitnessUtxo = tx;
+					if (input is PSBT0.PSBT0Input input0)
+						input0.non_witness_utxo_check = input.PrevOut.Hash;
 					if (Network.Consensus.NeverNeedPreviousTxForSigning ||
 					input.GetCoin()?.IsMalleable is false)
 						input.WitnessUtxo = output;
@@ -635,11 +637,12 @@ public PSBT SignWithKeys(params ISecret[] keys)
 
 		public PSBT SignWithKeys(params Key[] keys)
 		{
-			AssertSanity();
+			var errors = CheckSanity();
+			var hasError = new HashSet<uint>(errors.Select(e => e.InputIndex));
 			var signingOptions = GetSigningOptions(null);
 			foreach (var key in keys)
 			{
-				foreach (var input in this.Inputs)
+				foreach (var input in this.Inputs.Where(i => !hasError.Contains(i.Index)))
 				{
 					input.Sign(key, signingOptions);
 				}