feat: do taproot checks for addOutput()

Author: motorina0

src/psbt.d.ts

@@ -1,6 +1,6 @@
 /// <reference types="node" />
 import { Psbt as PsbtBase } from 'bip174';
-import { KeyValue, PsbtGlobalUpdate, PsbtInput, PsbtInputUpdate, PsbtOutput, PsbtOutputUpdate } from 'bip174/src/lib/interfaces';
+import { KeyValue, PsbtGlobalUpdate, PsbtInput, PsbtInputUpdate, PsbtOutput, PsbtOutputUpdate, TapInternalKey, TapTree } from 'bip174/src/lib/interfaces';
 import { Network } from './networks';
 import { Transaction } from './transaction';
 export interface TransactionInput {
@@ -125,7 +125,7 @@ interface PsbtOptsOptional {
 }
 interface PsbtInputExtended extends PsbtInput, TransactionInput {
 }
-declare type PsbtOutputExtended = PsbtOutputExtendedAddress | PsbtOutputExtendedScript;
+declare type PsbtOutputExtended = PsbtOutputExtendedAddress | PsbtOutputExtendedScript | PsbtOutputExtendedTaproot;
 interface PsbtOutputExtendedAddress extends PsbtOutput {
     address: string;
     value: number;
@@ -134,6 +134,12 @@ interface PsbtOutputExtendedScript extends PsbtOutput {
     script: Buffer;
     value: number;
 }
+interface PsbtOutputExtendedTaproot extends PsbtOutput {
+    tapInternalKey: TapInternalKey;
+    tapTree?: TapTree;
+    script?: Buffer;
+    value: number;
+}
 interface HDSignerBase {
     /**
      * DER format compressed publicKey buffer

src/psbt.js

@@ -226,11 +226,13 @@ class Psbt {
       arguments.length > 1 ||
       !outputData ||
       outputData.value === undefined ||
-      (outputData.address === undefined && outputData.script === undefined)
+      (outputData.address === undefined &&
+        outputData.script === undefined &&
+        outputData.tapInternalKey === undefined)
     ) {
       throw new Error(
         `Invalid arguments for Psbt.addOutput. ` +
-          `Requires single object with at least [script or address] and [value]`,
+          `Requires single object with at least [script, address or tapInternalKey] and [value]`,
       );
     }
     checkInputsForPartialSig(this.data.inputs, 'addOutput');
@@ -240,6 +242,20 @@ class Psbt {
       const script = (0, address_1.toOutputScript)(address, network);
       outputData = Object.assign(outputData, { script });
     }
+    if ((0, bip371_1.isTaprootOutput)(outputData)) {
+      (0, bip371_1.checkTaprootOutputFields)(
+        outputData,
+        outputData,
+        'addOutput',
+      );
+      const scriptAndAddress = (0, bip371_1.getNewTaprootScriptAndAddress)(
+        outputData,
+        outputData,
+        this.opts.network,
+      );
+      if (scriptAndAddress)
+        outputData = Object.assign(outputData, scriptAndAddress);
+    }
     const c = this.__CACHE;
     this.data.addOutput(outputData);
     c.__FEE = undefined;
@@ -1028,6 +1044,7 @@ function checkFees(psbt, cache, opts) {
   }
 }
 function checkInputsForPartialSig(inputs, action) {
+  // todo: add for taproot
   inputs.forEach(input => {
     let throws = false;
     let pSigs = [];

src/psbt/bip371.d.ts

@@ -1,6 +1,7 @@
 /// <reference types="node" />
 import { Taptree } from '../types';
-import { PsbtInput, TapLeaf } from 'bip174/src/lib/interfaces';
+import { PsbtInput, PsbtOutput, TapLeaf } from 'bip174/src/lib/interfaces';
+import { Network } from '../networks';
 export declare const toXOnly: (pubKey: Buffer) => Buffer;
 /**
  * Default tapscript finalizer. It searches for the `tapLeafHashToFinalize` if provided.
@@ -16,7 +17,13 @@ export declare function tapScriptFinalizer(inputIndex: number, input: PsbtInput,
 };
 export declare function serializeTaprootSignature(sig: Buffer, sighashType?: number): Buffer;
 export declare function isTaprootInput(input: PsbtInput): boolean;
+export declare function isTaprootOutput(output: PsbtOutput, script?: Buffer): boolean;
 export declare function checkTaprootInputFields(inputData: PsbtInput, newInputData: PsbtInput, action: string): void;
+export declare function checkTaprootOutputFields(outputData: PsbtOutput, newOutputData: PsbtOutput, action: string): void;
+export declare function getNewTaprootScriptAndAddress(outputData: PsbtOutput, newOutputData: PsbtOutput, network?: Network): {
+    script: Buffer;
+    address: string;
+} | undefined;
 export declare function tweakInternalPubKey(inputIndex: number, input: PsbtInput): Buffer;
 /**
  * Convert a binary tree to a BIP371 type list. Each element of the list is (according to BIP371):

src/psbt/bip371.js

@@ -1,9 +1,10 @@
 'use strict';
 Object.defineProperty(exports, '__esModule', { value: true });
-exports.tapTreeFromList = exports.tapTreeToList = exports.tweakInternalPubKey = exports.checkTaprootInputFields = exports.isTaprootInput = exports.serializeTaprootSignature = exports.tapScriptFinalizer = exports.toXOnly = void 0;
+exports.tapTreeFromList = exports.tapTreeToList = exports.tweakInternalPubKey = exports.getNewTaprootScriptAndAddress = exports.checkTaprootOutputFields = exports.checkTaprootInputFields = exports.isTaprootOutput = exports.isTaprootInput = exports.serializeTaprootSignature = exports.tapScriptFinalizer = exports.toXOnly = void 0;
 const types_1 = require('../types');
 const psbtutils_1 = require('./psbtutils');
 const taprootutils_1 = require('../payments/taprootutils');
+const payments_1 = require('../payments');
 const toXOnly = pubKey => (pubKey.length === 32 ? pubKey : pubKey.slice(1, 33));
 exports.toXOnly = toXOnly;
 /**
@@ -52,11 +53,54 @@ function isTaprootInput(input) {
   );
 }
 exports.isTaprootInput = isTaprootInput;
+function isTaprootOutput(output, script) {
+  return (
+    output &&
+    !!(
+      output.tapInternalKey ||
+      output.tapTree ||
+      (output.tapBip32Derivation && output.tapBip32Derivation.length) ||
+      (script && (0, psbtutils_1.isP2TR)(script))
+    )
+  );
+}
+exports.isTaprootOutput = isTaprootOutput;
 function checkTaprootInputFields(inputData, newInputData, action) {
-  checkMixedTaprootAndNonTaprootFields(inputData, newInputData, action);
+  checkMixedTaprootAndNonTaprootInputFields(inputData, newInputData, action);
   checkIfTapLeafInTree(inputData, newInputData, action);
 }
 exports.checkTaprootInputFields = checkTaprootInputFields;
+function checkTaprootOutputFields(outputData, newOutputData, action) {
+  checkMixedTaprootAndNonTaprootOutputFields(outputData, newOutputData, action);
+}
+exports.checkTaprootOutputFields = checkTaprootOutputFields;
+function getNewTaprootScriptAndAddress(outputData, newOutputData, network) {
+  if (!newOutputData.tapTree && !newOutputData.tapInternalKey) return;
+  const tapInternalKey =
+    newOutputData.tapInternalKey || outputData.tapInternalKey;
+  const tapTree = newOutputData.tapTree || outputData.tapTree;
+  if (tapInternalKey) {
+    const { script, address } = getTaprootScriptAndAddress(
+      tapInternalKey,
+      tapTree,
+      network,
+    );
+    const { script: newScript } = newOutputData;
+    if (newScript && !newScript.equals(script))
+      throw new Error('Error adding output. Script or address missmatch.');
+    return { script, address };
+  }
+}
+exports.getNewTaprootScriptAndAddress = getNewTaprootScriptAndAddress;
+function getTaprootScriptAndAddress(tapInternalKey, tapTree, network) {
+  const scriptTree = tapTree && tapTreeFromList(tapTree.leaves);
+  const { output, address } = (0, payments_1.p2tr)({
+    internalPubkey: tapInternalKey,
+    scriptTree,
+    network,
+  });
+  return { script: output, address: address };
+}
 function tweakInternalPubKey(inputIndex, input) {
   const tapInternalKey = input.tapInternalKey;
   const outputKey =
@@ -144,14 +188,36 @@ function instertLeafInTree(leaf, tree, depth = 0) {
   const rightSide = instertLeafInTree(leaf, tree && tree[1], depth + 1);
   if (rightSide) return [tree && tree[0], rightSide];
 }
-function checkMixedTaprootAndNonTaprootFields(inputData, newInputData, action) {
+function checkMixedTaprootAndNonTaprootInputFields(
+  inputData,
+  newInputData,
+  action,
+) {
   const isBadTaprootUpdate =
-    isTaprootInput(inputData) && hasNonTaprootInputFields(newInputData);
+    isTaprootInput(inputData) && hasNonTaprootFields(newInputData);
   const isBadNonTaprootUpdate =
-    hasNonTaprootInputFields(inputData) && isTaprootInput(newInputData);
+    hasNonTaprootFields(inputData) && isTaprootInput(newInputData);
   const hasMixedFields =
     inputData === newInputData &&
-    (isTaprootInput(newInputData) && hasNonTaprootInputFields(newInputData));
+    (isTaprootInput(newInputData) && hasNonTaprootFields(newInputData)); // todo: bad? use !===
+  if (isBadTaprootUpdate || isBadNonTaprootUpdate || hasMixedFields)
+    throw new Error(
+      `Invalid arguments for Psbt.${action}. ` +
+        `Cannot use both taproot and non-taproot fields.`,
+    );
+}
+function checkMixedTaprootAndNonTaprootOutputFields(
+  inputData,
+  newInputData,
+  action,
+) {
+  const isBadTaprootUpdate =
+    isTaprootOutput(inputData) && hasNonTaprootFields(newInputData);
+  const isBadNonTaprootUpdate =
+    hasNonTaprootFields(inputData) && isTaprootOutput(newInputData);
+  const hasMixedFields =
+    inputData === newInputData &&
+    (isTaprootOutput(newInputData) && hasNonTaprootFields(newInputData));
   if (isBadTaprootUpdate || isBadNonTaprootUpdate || hasMixedFields)
     throw new Error(
       `Invalid arguments for Psbt.${action}. ` +
@@ -244,13 +310,13 @@ function canFinalizeLeaf(leaf, tapScriptSig, hash) {
     tapScriptSig.find(tss => tss.leafHash.equals(leafHash)) !== undefined
   );
 }
-function hasNonTaprootInputFields(input) {
+function hasNonTaprootFields(io) {
   return (
-    input &&
+    io &&
     !!(
-      input.redeemScript ||
-      input.witnessScript ||
-      (input.bip32Derivation && input.bip32Derivation.length)
+      io.redeemScript ||
+      io.witnessScript ||
+      (io.bip32Derivation && io.bip32Derivation.length)
     )
   );
 }