feat(utxo-lib): psbt util functions for output proprietary key vals

TICKET: BTC-2183

Author: christopher-fong

modules/utxo-lib/src/bitgo/PsbtUtil.ts

@@ -1,5 +1,5 @@
 import { decodeProprietaryKey, ProprietaryKey } from 'bip174/src/lib/proprietaryKeyVal';
-import { PsbtInput } from 'bip174/src/lib/interfaces';
+import { PsbtInput, PsbtOutput, KeyValue } from 'bip174/src/lib/interfaces';
 import { Psbt } from 'bitcoinjs-lib/src/psbt';
 
 /**
@@ -38,21 +38,14 @@ export interface ProprietaryKeySearch {
   identifierEncoding?: BufferEncoding;
 }
 
-/**
- * Search any data from psbt proprietary key value against keydata.
- * Default identifierEncoding is utf-8 for identifier.
- */
-export function getPsbtInputProprietaryKeyVals(
-  input: PsbtInput,
+function getProprietaryKeyValuesFromUnknownKeyValues(
+  unknownKeyVals: KeyValue[],
   keySearch?: ProprietaryKeySearch
 ): ProprietaryKeyValue[] {
-  if (!input.unknownKeyVals?.length) {
-    return [];
-  }
   if (keySearch && keySearch.subtype === undefined && Buffer.isBuffer(keySearch.keydata)) {
     throw new Error('invalid proprietary key search filter combination. subtype is required');
   }
-  const keyVals = input.unknownKeyVals.map(({ key, value }, i) => {
+  const keyVals = unknownKeyVals.map(({ key, value }, i) => {
     return { key: decodeProprietaryKey(key), value };
   });
   return keyVals.filter((keyVal) => {
@@ -65,7 +58,28 @@ export function getPsbtInputProprietaryKeyVals(
     );
   });
 }
-
+/**
+ * Search any data from psbt proprietary key value against keydata.
+ * Default identifierEncoding is utf-8 for identifier.
+ */
+export function getPsbtInputProprietaryKeyVals(
+  input: PsbtInput,
+  keySearch?: ProprietaryKeySearch
+): ProprietaryKeyValue[] {
+  if (!input.unknownKeyVals?.length) {
+    return [];
+  }
+  return getProprietaryKeyValuesFromUnknownKeyValues(input.unknownKeyVals, keySearch);
+}
+export function getPsbtOutputProprietaryKeyVals(
+  output: PsbtOutput,
+  keySearch?: ProprietaryKeySearch
+): ProprietaryKeyValue[] {
+  if (!output.unknownKeyVals?.length) {
+    return [];
+  }
+  return getProprietaryKeyValuesFromUnknownKeyValues(output.unknownKeyVals, keySearch);
+}
 /**
  * @return partialSig/tapScriptSig/MUSIG2_PARTIAL_SIG count iff input is not finalized
  */
@@ -82,14 +96,12 @@ export function getPsbtInputSignatureCount(input: PsbtInput): number {
     }).length
   );
 }
-
 /**
  * @return true iff PSBT input is finalized
  */
 export function isPsbtInputFinalized(input: PsbtInput): boolean {
   return Buffer.isBuffer(input.finalScriptSig) || Buffer.isBuffer(input.finalScriptWitness);
 }
-
 /**
  * @return true iff data starts with magic PSBT byte sequence
  * @param data byte array or hex string
@@ -105,7 +117,6 @@ 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.
@@ -123,11 +134,9 @@ export function toPsbtBuffer(data: Buffer | string): Buffer {
     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) {
@@ -141,13 +150,10 @@ export function toPsbtBuffer(data: Buffer | string): 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.
  */

modules/utxo-lib/src/bitgo/UtxoPsbt.ts

@@ -7,7 +7,7 @@ import {
   Transaction as ITransaction,
   TransactionFromBuffer,
 } from 'bip174/src/lib/interfaces';
-import { checkForInput } from 'bip174/src/lib/utils';
+import { checkForInput, checkForOutput } from 'bip174/src/lib/utils';
 import { BufferWriter, varuint } from 'bitcoinjs-lib/src/bufferutils';
 import { SessionKey } from '@brandonblack/musig';
 import { BIP32Factory, BIP32Interface } from 'bip32';
@@ -57,6 +57,7 @@ import { getTaprootOutputKey } from '../taproot';
 import {
   getPsbtInputProprietaryKeyVals,
   getPsbtInputSignatureCount,
+  getPsbtOutputProprietaryKeyVals,
   ProprietaryKeySearch,
   ProprietaryKeySubtype,
   ProprietaryKeyValue,
@@ -1109,7 +1110,7 @@ export class UtxoPsbt<Tx extends UtxoTransaction<bigint> = UtxoTransaction<bigin
   }
 
   /**
-   * To search any data from proprietary key value against keydata.
+   * To search any data from proprietary key value against keydata in the inputs.
    * Default identifierEncoding is utf-8 for identifier.
    */
   getProprietaryKeyVals(inputIndex: number, keySearch?: ProprietaryKeySearch): ProprietaryKeyValue[] {
@@ -1118,7 +1119,7 @@ export class UtxoPsbt<Tx extends UtxoTransaction<bigint> = UtxoTransaction<bigin
   }
 
   /**
-   * To delete any data from proprietary key value.
+   * To delete any data from proprietary key value in PSBT input.
    * Default identifierEncoding is utf-8 for identifier.
    */
   deleteProprietaryKeyVals(inputIndex: number, keysToDelete?: ProprietaryKeySearch): this {
@@ -1142,6 +1143,75 @@ export class UtxoPsbt<Tx extends UtxoTransaction<bigint> = UtxoTransaction<bigin
     return this;
   }
 
+  /**
+   * Adds a proprietary key value pair to PSBT output
+   * Default identifier is utf-8 for identifier
+   */
+  addProprietaryKeyValToOutput(outputIndex: number, keyValueData: ProprietaryKeyValue): this {
+    const output = checkForOutput(this.data.outputs, outputIndex);
+    assert(output.unknownKeyVals);
+    return this.addUnknownKeyValToOutput(outputIndex, {
+      key: encodeProprietaryKey(keyValueData.key),
+      value: keyValueData.value,
+    });
+  }
+
+  /**
+   * To search any data from proprietary key value against keydata in the PSBT outputs.
+   * Default identifierEncoding is utf-8 for identifier.
+   */
+  getOutputProprietaryKeyVals(outputIndex: number, keySearch?: ProprietaryKeySearch): ProprietaryKeyValue[] {
+    const output = checkForOutput(this.data.outputs, outputIndex);
+    return getPsbtOutputProprietaryKeyVals(output, keySearch);
+  }
+
+  /**
+   * Adds or updates (if exists) proprietary key value pair to PSBT output.
+   * Default identifierEncoding is utf-8 for identifier.
+   */
+  addOrUpdateProprietaryKeyValsToOutput(outputIndex: number, keyValueData: ProprietaryKeyValue): this {
+    const output = checkForOutput(this.data.outputs, outputIndex);
+    const key = encodeProprietaryKey(keyValueData.key);
+    const { value } = keyValueData;
+    if (output.unknownKeyVals?.length) {
+      const ukvIndex = output.unknownKeyVals.findIndex((ukv) => ukv.key.equals(key));
+      if (ukvIndex > -1) {
+        output.unknownKeyVals[ukvIndex] = { key, value };
+        return this;
+      }
+    }
+    this.addUnknownKeyValToOutput(outputIndex, {
+      key,
+      value,
+    });
+    return this;
+  }
+
+  /**
+   * To delete any data from proprietary key value in PSBT output.
+   * Default identifierEncoding is utf-8 for identifier.
+   */
+  deleteProprietaryKeyValsInOutput(outputIndex: number, keysToDelete?: ProprietaryKeySearch): this {
+    const output = checkForOutput(this.data.outputs, outputIndex);
+    if (!output.unknownKeyVals?.length) {
+      return this;
+    }
+    if (keysToDelete && keysToDelete.subtype === undefined && Buffer.isBuffer(keysToDelete.keydata)) {
+      throw new Error('invalid proprietary key search filter combination. subtype is required');
+    }
+    output.unknownKeyVals = output.unknownKeyVals.filter((keyValue, i) => {
+      const key = decodeProprietaryKey(keyValue.key);
+      return !(
+        keysToDelete === undefined ||
+        (keysToDelete.identifier === key.identifier &&
+          (keysToDelete.subtype === undefined ||
+            (keysToDelete.subtype === key.subtype &&
+              (!Buffer.isBuffer(keysToDelete.keydata) || keysToDelete.keydata.equals(key.keydata)))))
+      );
+    });
+    return this;
+  }
+
   private createMusig2NonceForInput(
     inputIndex: number,
     keyPair: BIP32Interface,