From dcc4d930687c8817fc4c2e4e8782885ef2ca66a3 Mon Sep 17 00:00:00 2001
From: Otto Allmendinger <otto@bitgo.com>
Date: Fri, 27 Sep 2024 11:14:30 +0200
Subject: [PATCH] feat(utxo-lib): add toPsbtBuffer

First checks if the input is already a buffer that starts with the magic PSBT
byte sequence. If not, it checks if the input is a base64- or hex-encoded string
that starts with the PSBT header.

This function is useful when reading a file or request that could be in any of
the above formats.

Issue: BTC-1351
---
 modules/utxo-lib/src/bitgo/PsbtUtil.ts        | 42 +++++++++++++++++++
 modules/utxo-lib/src/bitgo/transaction.ts     | 20 +++++++++
 .../utxo-lib/test/bitgo/psbt/toPsbtBuffer.ts  | 27 ++++++++++++
 3 files changed, 89 insertions(+)
 create mode 100644 modules/utxo-lib/test/bitgo/psbt/toPsbtBuffer.ts

diff --git a/modules/utxo-lib/src/bitgo/PsbtUtil.ts b/modules/utxo-lib/src/bitgo/PsbtUtil.ts
index 61bf78913d..e28cea1bde 100644
--- a/modules/utxo-lib/src/bitgo/PsbtUtil.ts
+++ b/modules/utxo-lib/src/bitgo/PsbtUtil.ts
@@ -106,6 +106,48 @@ export function isPsbt(data: Buffer | string): boolean {
   return 5 <= data.length && data.readUInt32BE(0) === 0x70736274 && data.readUInt8(4) === 0xff;
 }
 
+/**
+ * First checks if the input is already a buffer that starts with the magic PSBT byte sequence.
+ * If not, it checks if the input is a base64- or hex-encoded string that starts with PSBT header.
+ *
+ * This function is useful when reading a file that could be in any of the above formats or when
+ * dealing with a request that could contain a hex or base64 encoded PSBT.
+ *
+ * @param data
+ * @return buffer that starts with the magic PSBT byte sequence
+ * @throws Error when conversion is not possible
+ */
+export function toPsbtBuffer(data: Buffer | string): Buffer {
+  if (Buffer.isBuffer(data)) {
+    // we are dealing with a buffer that looks like a psbt already
+    if (isPsbt(data)) {
+      return data;
+    }
+
+    // we could be dealing with a buffer that could be a hex or base64 encoded psbt
+    data = data.toString('ascii');
+  }
+
+  if (typeof data === 'string') {
+    const encodings = ['hex', 'base64'] as const;
+    for (const encoding of encodings) {
+      let buffer: Buffer;
+      try {
+        buffer = Buffer.from(data, encoding);
+      } catch (e) {
+        continue;
+      }
+      if (isPsbt(buffer)) {
+        return buffer;
+      }
+    }
+
+    throw new Error(`data is not in any of the following formats: ${encodings.join(', ')}`);
+  }
+
+  throw new Error('data must be a buffer or a string');
+}
+
 /**
  * This function allows signing or validating a psbt with non-segwit inputs those do not contain nonWitnessUtxo.
  */
diff --git a/modules/utxo-lib/src/bitgo/transaction.ts b/modules/utxo-lib/src/bitgo/transaction.ts
index badb3e919e..deadda0751 100644
--- a/modules/utxo-lib/src/bitgo/transaction.ts
+++ b/modules/utxo-lib/src/bitgo/transaction.ts
@@ -13,6 +13,7 @@ import { ZcashPsbt } from './zcash/ZcashPsbt';
 import { ZcashTransactionBuilder } from './zcash/ZcashTransactionBuilder';
 import { ZcashNetwork, ZcashTransaction } from './zcash/ZcashTransaction';
 import { LitecoinPsbt, LitecoinTransaction, LitecoinTransactionBuilder } from './litecoin';
+import { isPsbt, toPsbtBuffer } from './PsbtUtil';
 
 export function createTransactionFromBuffer(
   buf: Buffer,
@@ -87,7 +88,15 @@ export function createTransactionFromHex<TNumber extends number | bigint = numbe
   return createTransactionFromBuffer<TNumber>(Buffer.from(hex, 'hex'), network, p);
 }
 
+/**
+ * @param buf - must start with the PSBT magic bytes
+ * @param network
+ * @param bip32PathsAbsolute
+ */
 export function createPsbtFromBuffer(buf: Buffer, network: Network, bip32PathsAbsolute = false): UtxoPsbt {
+  if (!isPsbt(buf)) {
+    throw new Error(`invalid psbt (does not start with 'psbt' magic bytes)`);
+  }
   switch (getMainnet(network)) {
     case networks.bitcoin:
     case networks.bitcoincash:
@@ -108,6 +117,17 @@ export function createPsbtFromBuffer(buf: Buffer, network: Network, bip32PathsAb
   throw new Error(`invalid network`);
 }
 
+/**
+ * Like createPsbtFromBuffer, but attempts hex and base64 decoding as well.
+ *
+ * @param buf
+ * @param network
+ * @param bip32PathsAbsolute
+ */
+export function createPsbtDecode(buf: Buffer | string, network: Network, bip32PathsAbsolute = false): UtxoPsbt {
+  return createPsbtFromBuffer(toPsbtBuffer(buf), network, bip32PathsAbsolute);
+}
+
 export function createPsbtFromHex(hex: string, network: Network, bip32PathsAbsolute = false): UtxoPsbt {
   return createPsbtFromBuffer(Buffer.from(hex, 'hex'), network, bip32PathsAbsolute);
 }
diff --git a/modules/utxo-lib/test/bitgo/psbt/toPsbtBuffer.ts b/modules/utxo-lib/test/bitgo/psbt/toPsbtBuffer.ts
new file mode 100644
index 0000000000..03baaff220
--- /dev/null
+++ b/modules/utxo-lib/test/bitgo/psbt/toPsbtBuffer.ts
@@ -0,0 +1,27 @@
+import * as assert from 'assert';
+
+import { toPsbtBuffer } from '../../../src/bitgo';
+
+describe('bufferUtil', function () {
+  function variants(data: Buffer | string): (Buffer | string)[] {
+    return [
+      data,
+      data.toString('hex'),
+      Buffer.from(data.toString('hex')),
+      data.toString('base64'),
+      Buffer.from(data.toString('base64')),
+    ];
+  }
+
+  it('should convert a buffer to a string', function () {
+    const psbt = Buffer.from('psbt\xff', 'ascii');
+    for (const v of variants(psbt)) {
+      assert.ok(toPsbtBuffer(v).equals(psbt));
+    }
+
+    const nonPsbt = Buffer.from('hello world', 'ascii');
+    for (const v of variants(nonPsbt)) {
+      assert.throws(() => toPsbtBuffer(v));
+    }
+  });
+});