From 4bff77856409dafc4d8caca6f6e82b7084cdb7da Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Tue, 25 Jun 2024 13:54:58 +0200
Subject: [PATCH] psbt: decode keytype as compact size

Fixes #2199.

Previous to this fix the keytype was only interpreted as a single byte,
even though BIP-0174 states it is to be parsed as a CompactSize/VarInt.
---
 btcutil/psbt/psbt.go      |  6 ++++++
 btcutil/psbt/psbt_test.go |  5 +++++
 btcutil/psbt/utils.go     | 25 +++++++++++++++++--------
 3 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/btcutil/psbt/psbt.go b/btcutil/psbt/psbt.go
index 964061bdc5..a72b7fc551 100644
--- a/btcutil/psbt/psbt.go
+++ b/btcutil/psbt/psbt.go
@@ -37,6 +37,12 @@ const MaxPsbtValueLength = 4000000
 // deserialize from the wire. Anything more will return ErrInvalidKeyData.
 const MaxPsbtKeyLength = 10000
 
+// MaxPsbtKeyValue is the maximum value of a key type in a PSBT. This maximum
+// isn't specified by the BIP but used by bitcoind in various places to limit
+// the number of items processed. So we use it to validate the key type in order
+// to have a consistent behavior.
+const MaxPsbtKeyValue = 0x02000000
+
 var (
 
 	// ErrInvalidPsbtFormat is a generic error for any situation in which a
diff --git a/btcutil/psbt/psbt_test.go b/btcutil/psbt/psbt_test.go
index 2309b07e40..a39f52923a 100644
--- a/btcutil/psbt/psbt_test.go
+++ b/btcutil/psbt/psbt_test.go
@@ -123,6 +123,8 @@ var invalidPsbtHex = map[int]string{
 	18: "70736274ff0100550200000001279a2323a5dfb51fc45f220fa58b0fc13e1e3342792a85d7e36cd6333b5cbc390000000000ffffffff01a05aea0b000000001976a914ffe9c0061097cc3b636f2cb0460fa4fc427d2b4588ac0000000000010120955eea0b0000000017a9146345200f68d189e1adc0df1c4d16ea8f14c0dbeb87220203b1341ccba7683b6af4f1238cd6e97e7167d569fac47f1e48d47541844355bd4646304302200424b58effaaa694e1559ea5c93bbfd4a89064224055cdf070b6771469442d07021f5c8eb0fea6516d60b8acb33ad64ede60e8785bfb3aa94b99bdf86151db9a9a01220203b1341ccba7683b6af4f1238cd6e97e7167d569fac47f1e48d47541844355bd4646304302200424b58effaaa694e1559ea5c93bbfd4a89064224055cdf070b6771469442d07021f5c8eb0fea6516d60b8acb33ad64ede60e8785bfb3aa94b99bdf86151db9a9a010104220020771fd18ad459666dd49f3d564e3dbc42f4c84774e360ada16816a8ed488d5681010547522103b1341ccba7683b6af4f1238cd6e97e7167d569fac47f1e48d47541844355bd462103de55d1e1dac805e3f8a58c1fbf9b94c02f3dbaafe127fefca4995f26f82083bd52ae220603b1341ccba7683b6af4f1238cd6e97e7167d569fac47f1e48d47541844355bd4610b4a6ba67000000800000008004000080220603de55d1e1dac805e3f8a58c1fbf9b94c02f3dbaafe127fefca4995f26f82083bd10b4a6ba670000008000000080050000800000",
 	// Invalid duplicate BIP32 derivation (different derivs, same key)
 	19: "70736274ff0100550200000001279a2323a5dfb51fc45f220fa58b0fc13e1e3342792a85d7e36cd6333b5cbc390000000000ffffffff01a05aea0b000000001976a914ffe9c0061097cc3b636f2cb0460fa4fc427d2b4588ac0000000000010120955eea0b0000000017a9146345200f68d189e1adc0df1c4d16ea8f14c0dbeb87220203b1341ccba7683b6af4f1238cd6e97e7167d569fac47f1e48d47541844355bd4646304302200424b58effaaa694e1559ea5c93bbfd4a89064224055cdf070b6771469442d07021f5c8eb0fea6516d60b8acb33ad64ede60e8785bfb3aa94b99bdf86151db9a9a010104220020771fd18ad459666dd49f3d564e3dbc42f4c84774e360ada16816a8ed488d5681010547522103b1341ccba7683b6af4f1238cd6e97e7167d569fac47f1e48d47541844355bd462103de55d1e1dac805e3f8a58c1fbf9b94c02f3dbaafe127fefca4995f26f82083bd52ae220603b1341ccba7683b6af4f1238cd6e97e7167d569fac47f1e48d47541844355bd4610b4a6ba67000000800000008004000080220603b1341ccba7683b6af4f1238cd6e97e7167d569fac47f1e48d47541844355bd4610b4a6ba670000008000000080050000800000",
+	// Invalid var int for key type
+	20: "70736274ff01001c000000000002000000000000000000000000736210ff01000001010010ff70ff01001c00000000000000000000000000000000000000000000",
 }
 
 // All following PSBTs are Taproot specific invalid packets taken from
@@ -682,6 +684,9 @@ func TestFinalize2of3(t *testing.T) {
 		t.Fatalf("Error decoding hex: %v", err)
 	}
 	p, err := NewFromRawBytes(bytes.NewReader(b), false)
+	if err != nil {
+		t.Fatalf("Error parsing PSBT: %v", err)
+	}
 	if p.IsComplete() {
 		t.Fatalf("Psbt is complete")
 	}
diff --git a/btcutil/psbt/utils.go b/btcutil/psbt/utils.go
index 0a9002798e..c47f6afd4d 100644
--- a/btcutil/psbt/utils.go
+++ b/btcutil/psbt/utils.go
@@ -250,23 +250,32 @@ func getKey(r io.Reader) (int, []byte, error) {
 
 	// Next, we ready out the designated number of bytes, which may include
 	// a type, key, and optional data.
-	keyTypeAndData := make([]byte, count)
-	if _, err := io.ReadFull(r, keyTypeAndData[:]); err != nil {
-		return -1, nil, err
+	keyTypeReader := io.LimitReader(r, int64(count))
+	keyType, err := wire.ReadVarInt(keyTypeReader, 0)
+	if err != nil {
+		return -1, nil, ErrInvalidPsbtFormat
 	}
 
-	keyType := int(string(keyTypeAndData)[0])
+	// The maximum value of a compact size int is capped in bitcoind, do the
+	// same here to mimic the behavior.
+	if keyType > MaxPsbtKeyValue {
+		return -1, nil, ErrInvalidPsbtFormat
+	}
+
+	keyData, err := io.ReadAll(keyTypeReader)
+	if err != nil {
+		return -1, nil, ErrInvalidPsbtFormat
+	}
 
 	// Note that the second return value will usually be empty, since most
 	// keys contain no more than the key type byte.
-	if len(keyTypeAndData) == 1 {
-		return keyType, nil, nil
+	if len(keyData) == 0 {
+		return int(keyType), nil, nil
 	}
 
 	// Otherwise, we return the key, along with any data that it may
 	// contain.
-	return keyType, keyTypeAndData[1:], nil
-
+	return int(keyType), keyData, nil
 }
 
 // readTxOut is a limited version of wire.ReadTxOut, because the latter is not