From d73d3439f67fe528479b759013529d7049c8cba2 Mon Sep 17 00:00:00 2001
From: Craig Raw <craigraw@gmail.com>
Date: Thu, 11 Jul 2024 11:48:32 +0200
Subject: [PATCH] handle psbt keytypes with values greater than single byte
 compact integers

---
 .../sparrowwallet/drongo/psbt/PSBTEntry.java  | 37 +++++++++----------
 .../sparrowwallet/drongo/psbt/PSBTTest.java   |  7 ++++
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/src/main/java/com/sparrowwallet/drongo/psbt/PSBTEntry.java b/src/main/java/com/sparrowwallet/drongo/psbt/PSBTEntry.java
index 54be38f..1010baa 100644
--- a/src/main/java/com/sparrowwallet/drongo/psbt/PSBTEntry.java
+++ b/src/main/java/com/sparrowwallet/drongo/psbt/PSBTEntry.java
@@ -16,11 +16,11 @@ import java.util.Map;
 
 public class PSBTEntry {
     private final byte[] key;
-    private final byte keyType;
+    private final int keyType;
     private final byte[] keyData;
     private final byte[] data;
 
-    public PSBTEntry(byte[] key, byte keyType, byte[] keyData, byte[] data) {
+    public PSBTEntry(byte[] key, int keyType, byte[] keyData, byte[] data) {
         this.key = key;
         this.keyType = keyType;
         this.keyData = keyData;
@@ -39,12 +39,13 @@ public class PSBTEntry {
             byte[] key = new byte[keyLen];
             psbtByteBuffer.get(key);
 
-            byte keyType = key[0];
+            ByteBuffer keyBuf = ByteBuffer.wrap(key);
+            int keyType = readCompactInt(keyBuf);
 
             byte[] keyData = null;
-            if (key.length > 1) {
-                keyData = new byte[key.length - 1];
-                System.arraycopy(key, 1, keyData, 0, keyData.length);
+            if(keyBuf.hasRemaining()) {
+                keyData = new byte[keyBuf.remaining()];
+                keyBuf.get(keyData);
             }
 
             int dataLen = readCompactInt(psbtByteBuffer);
@@ -143,22 +144,20 @@ public class PSBTEntry {
         return baos.toByteArray();
     }
 
-    static PSBTEntry populateEntry(byte type, byte[] keydata, byte[] data) {
-        return new PSBTEntry(new byte[] {type}, type, keydata, data);
-    }
-
-    void serializeToStream(ByteArrayOutputStream baos) {
-        int keyLen = 1;
-        if(keyData != null) {
-            keyLen += keyData.length;
-        }
-
-        baos.writeBytes(writeCompactInt(keyLen));
-        baos.writeBytes(key);
+    static PSBTEntry populateEntry(int type, byte[] keyData, byte[] data) {
+        ByteArrayOutputStream baos = new ByteArrayOutputStream(1 + (keyData == null ? 0 : keyData.length));
+        baos.writeBytes(writeCompactInt(type));
         if(keyData != null) {
             baos.writeBytes(keyData);
         }
 
+        return new PSBTEntry(baos.toByteArray(), type, keyData, data);
+    }
+
+    void serializeToStream(ByteArrayOutputStream baos) {
+        baos.writeBytes(writeCompactInt(key.length));
+        baos.writeBytes(key);
+
         baos.writeBytes(writeCompactInt(data.length));
         baos.writeBytes(data);
     }
@@ -167,7 +166,7 @@ public class PSBTEntry {
         return key;
     }
 
-    public byte getKeyType() {
+    public int getKeyType() {
         return keyType;
     }
 
diff --git a/src/test/java/com/sparrowwallet/drongo/psbt/PSBTTest.java b/src/test/java/com/sparrowwallet/drongo/psbt/PSBTTest.java
index 1220e96..0fbfbd9 100644
--- a/src/test/java/com/sparrowwallet/drongo/psbt/PSBTTest.java
+++ b/src/test/java/com/sparrowwallet/drongo/psbt/PSBTTest.java
@@ -313,6 +313,13 @@ public class PSBTTest {
         Assertions.assertEquals(noWitnessDataPsbt, psbt4.toBase64String());
     }
 
+    @Test
+    public void largeKeyType() throws PSBTParseException {
+        String psbt = "cHNidP8BAD8CAAAAAf//////////////////////////////////////////AAAAAAD/////AQAAAAAAAAAAA2oBAAAAAAAF/v//AAAAAAAA";
+        PSBT psbt1 = PSBT.fromString(psbt);
+        Assertions.assertEquals("cHNidP8BAD8CAAAAAf//////////////////////////////////////////AAAAAAD/////AQAAAAAAAAAAA2oBAAAAAAAAAAA=", psbt1.toBase64String());
+    }
+
     @Test
     public void creatorBip() throws PSBTParseException {
         Transaction transaction = new Transaction();