Merge pull request #6425 from BitGo/BTC-2143.add-selectTapLeafScript

feat(utxo-core): add selectTapLeafScript option to createPsbt

Author: OttoAllmendinger

modules/utxo-core/src/descriptor/psbt/createPsbt.ts

@@ -1,5 +1,6 @@
+import { TapLeafScript } from 'bip174/src/lib/interfaces';
 import * as utxolib from '@bitgo/utxo-lib';
-import { Descriptor } from '@bitgo/wasm-miniscript';
+import { Descriptor, Miniscript } from '@bitgo/wasm-miniscript';
 
 import { DerivedDescriptorWalletOutput, WithOptDescriptor } from '../Output';
 import { Output } from '../../Output';
@@ -13,18 +14,43 @@ import { assertSatisfiable } from './assertSatisfiable';
  * */
 export const MAX_BIP125_RBF_SEQUENCE = 0xffffffff - 2;
 
-function updateInputsWithDescriptors(psbt: utxolib.bitgo.UtxoPsbt, descriptors: Descriptor[]) {
-  if (psbt.txInputs.length !== descriptors.length) {
-    throw new Error(`Input count mismatch (psbt=${psbt.txInputs.length}, descriptors=${descriptors.length})`);
+export function findTapLeafScript(input: TapLeafScript[], script: Buffer | Miniscript): TapLeafScript {
+  if (!Buffer.isBuffer(script)) {
+    script = Buffer.from(script.encode());
+  }
+  const matches = input.filter((leaf) => {
+    return leaf.script.equals(script);
+  });
+  if (matches.length === 0) {
+    throw new Error(`No tapLeafScript found for script: ${script.toString('hex')}`);
+  }
+  if (matches.length > 1) {
+    throw new Error(`Multiple tapLeafScripts found for script: ${script.toString('hex')}`);
+  }
+  return matches[0];
+}
+
+function updateInputsWithDescriptors(
+  psbt: utxolib.bitgo.UtxoPsbt,
+  inputParams: Array<{ descriptor: Descriptor; selectTapLeafScript?: Miniscript }>
+) {
+  if (psbt.txInputs.length !== inputParams.length) {
+    throw new Error(`Input count mismatch (psbt=${psbt.txInputs.length}, inputParams=${inputParams.length})`);
   }
   const wrappedPsbt = toWrappedPsbt(psbt);
-  for (const [inputIndex, descriptor] of descriptors.entries()) {
-    assertSatisfiable(psbt, inputIndex, descriptor);
-    wrappedPsbt.updateInputWithDescriptor(inputIndex, descriptor);
+  for (const [inputIndex, v] of inputParams.entries()) {
+    assertSatisfiable(psbt, inputIndex, v.descriptor);
+    wrappedPsbt.updateInputWithDescriptor(inputIndex, v.descriptor);
   }
   const unwrappedPsbt = toUtxoPsbt(wrappedPsbt, psbt.network);
   for (const inputIndex in psbt.txInputs) {
-    psbt.data.inputs[inputIndex] = unwrappedPsbt.data.inputs[inputIndex];
+    const preparedInput = unwrappedPsbt.data.inputs[inputIndex];
+    const v = inputParams[inputIndex];
+    if (v.selectTapLeafScript && preparedInput.tapLeafScript) {
+      const selected = findTapLeafScript(preparedInput.tapLeafScript, v.selectTapLeafScript);
+      preparedInput.tapLeafScript = [selected];
+    }
+    psbt.data.inputs[inputIndex] = preparedInput;
   }
 }
 
@@ -49,6 +75,7 @@ export type PsbtParams = {
 };
 
 export type DerivedDescriptorTransactionInput = DerivedDescriptorWalletOutput & {
+  selectTapLeafScript?: Miniscript;
   sequence?: number;
 };
 
@@ -64,7 +91,10 @@ export function createPsbt(
   psbt.addOutputs(outputs);
   updateInputsWithDescriptors(
     psbt,
-    inputs.map((i) => i.descriptor)
+    inputs.map(({ descriptor, selectTapLeafScript }) => ({
+      descriptor,
+      selectTapLeafScript,
+    }))
   );
   updateOutputsWithDescriptors(psbt, outputs);
   return psbt;

modules/utxo-core/src/testutil/descriptor/descriptors.ts

@@ -1,6 +1,6 @@
 import assert from 'assert';
 
-import { Descriptor, ast } from '@bitgo/wasm-miniscript';
+import { Miniscript, Descriptor, ast } from '@bitgo/wasm-miniscript';
 import { bip32, BIP32Interface } from '@bitgo/utxo-lib';
 
 import { DescriptorMap, PsbtParams } from '../../descriptor';
@@ -134,6 +134,50 @@ function getDescriptorNode(
   throw new Error(`Unknown descriptor template: ${template}`);
 }
 
+type TapTree = [TapTree, TapTree] | ast.MiniscriptNode;
+
+function getTapLeafScriptNodes(t: ast.DescriptorNode | TapTree): ast.MiniscriptNode[] {
+  if (Array.isArray(t)) {
+    if (t.length !== 2) {
+      throw new Error(`expected tuple, got: ${JSON.stringify(t)}`);
+    }
+    return t.map((v) => (Array.isArray(v) ? getTapLeafScriptNodes(v) : v)).flat();
+  }
+
+  if (typeof t === 'object') {
+    const node = t;
+    if (!('tr' in node)) {
+      throw new Error(`TapLeafScripts are only supported for Taproot descriptors, got: ${t}`);
+    }
+    if (!Array.isArray(node.tr) || node.tr.length !== 2) {
+      throw new Error(`expected tuple, got: ${JSON.stringify(node.tr)}`);
+    }
+    const tapscript = node.tr[1];
+    if (!Array.isArray(tapscript)) {
+      throw new Error(`expected tapscript to be an array, got: ${JSON.stringify(tapscript)}`);
+    }
+    return getTapLeafScriptNodes(tapscript);
+  }
+
+  throw new Error(`Invalid input: ${JSON.stringify(t)}`);
+}
+
+export function containsKey(script: Miniscript | ast.MiniscriptNode, key: BIP32Interface | string): boolean {
+  if (script instanceof Miniscript) {
+    script = ast.fromMiniscript(script);
+  }
+  if ('pk' in script) {
+    return script.pk === toXOnly(key);
+  }
+  throw new Error(`Unsupported script type: ${JSON.stringify(script)}`);
+}
+
+export function getTapLeafScripts(d: Descriptor): string[] {
+  return getTapLeafScriptNodes(ast.fromDescriptor(d)).map((n) =>
+    Miniscript.fromString(ast.formatNode(n), 'tap').toString()
+  );
+}
+
 export function getDescriptor(
   template: DescriptorTemplate,
   keys: KeyTriple | string[] = getDefaultXPubs(),

modules/utxo-core/src/testutil/descriptor/mock.utils.ts

@@ -1,4 +1,4 @@
-import { Descriptor } from '@bitgo/wasm-miniscript';
+import { Descriptor, Miniscript } from '@bitgo/wasm-miniscript';
 import * as utxolib from '@bitgo/utxo-lib';
 
 import {
@@ -17,6 +17,7 @@ type BaseMockDescriptorOutputParams = {
   index?: number;
   value?: bigint;
   sequence?: number;
+  selectTapLeafScript?: Miniscript;
 };
 
 function mockOutputId(id?: MockOutputIdParams): {
@@ -42,13 +43,15 @@ export function mockDerivedDescriptorWalletOutput(
       value,
     },
     descriptor,
+    selectTapLeafScript: outputParams.selectTapLeafScript,
     sequence: outputParams.sequence,
   };
 }
 
 type MockInput = BaseMockDescriptorOutputParams & {
   index: number;
   descriptor: Descriptor;
+  selectTapLeafScript?: Miniscript;
 };
 
 type MockOutput = {