refactor(utxo-core): improve descriptor handling with derivation checks

Issue: BTC-1826

Author: OttoAllmendinger

modules/utxo-core/src/descriptor/Output.ts

@@ -1,9 +1,11 @@
+import assert from 'assert';
+
 import { Descriptor } from '@bitgo/wasm-miniscript';
 
 import { getFixedOutputSum, MaxOutput, Output, PrevOutput } from '../Output';
 
 import { DescriptorMap } from './DescriptorMap';
-import { createScriptPubKeyFromDescriptor } from './address';
+import { getDescriptorAtIndexCheckScript } from './derive';
 
 export type WithDescriptor<T> = T & {
   descriptor: Descriptor;
@@ -31,7 +33,7 @@ export function getExternalFixedAmount(outputs: WithOptDescriptor<Output | MaxOu
 
 export type DescriptorWalletOutput = PrevOutput & {
   descriptorName: string;
-  descriptorIndex: number;
+  descriptorIndex: number | undefined;
 };
 
 export type DerivedDescriptorWalletOutput = WithDescriptor<PrevOutput>;
@@ -44,17 +46,17 @@ export function toDerivedDescriptorWalletOutput(
   if (!descriptor) {
     throw new Error(`Descriptor not found: ${output.descriptorName}`);
   }
-  const derivedDescriptor = descriptor.atDerivationIndex(output.descriptorIndex);
-  const script = createScriptPubKeyFromDescriptor(derivedDescriptor);
-  if (!script.equals(output.witnessUtxo.script)) {
-    throw new Error(
-      `Script mismatch: descriptor ${output.descriptorName} ${descriptor.toString()} script=${script.toString('hex')}`
-    );
-  }
+  assert(descriptor instanceof Descriptor);
+  const descriptorAtIndex = getDescriptorAtIndexCheckScript(
+    descriptor,
+    output.descriptorIndex,
+    output.witnessUtxo.script,
+    output.descriptorName
+  );
   return {
     hash: output.hash,
     index: output.index,
     witnessUtxo: output.witnessUtxo,
-    descriptor: descriptor.atDerivationIndex(output.descriptorIndex),
+    descriptor: descriptorAtIndex,
   };
 }

modules/utxo-core/src/descriptor/address.ts

@@ -1,11 +1,11 @@
 import { Descriptor } from '@bitgo/wasm-miniscript';
 import * as utxolib from '@bitgo/utxo-lib';
 
-export function createScriptPubKeyFromDescriptor(descriptor: Descriptor, index?: number): Buffer {
+export function createScriptPubKeyFromDescriptor(descriptor: Descriptor, index: number | undefined): Buffer {
   if (index === undefined) {
     return Buffer.from(descriptor.scriptPubkey());
   }
-  return createScriptPubKeyFromDescriptor(descriptor.atDerivationIndex(index));
+  return createScriptPubKeyFromDescriptor(descriptor.atDerivationIndex(index), undefined);
 }
 
 export function createAddressFromDescriptor(

modules/utxo-core/src/descriptor/derive.ts

@@ -0,0 +1,44 @@
+import assert from 'assert';
+
+import { Descriptor } from '@bitgo/wasm-miniscript';
+
+/**
+ * Get a descriptor at a specific derivation index.
+ * For wildcard descriptors (containing '*'), the index is required and used for derivation.
+ * For definite descriptors (not containing '*'), no index should be provided.
+ * @param descriptor - The descriptor to derive from
+ * @param index - The derivation index for wildcard descriptors
+ * @returns A new descriptor at the specified index for wildcard descriptors, or the original descriptor for definite ones
+ * @throws {Error} If index is undefined for a wildcard descriptor or if index is provided for a definite descriptor
+ */
+export function getDescriptorAtIndex(descriptor: Descriptor, index: number | undefined): Descriptor {
+  assert(descriptor instanceof Descriptor);
+  Descriptor.fromString(descriptor.toString(), 'derivable');
+  descriptor = Descriptor.fromStringDetectType(descriptor.toString());
+  if (descriptor.hasWildcard()) {
+    if (index === undefined) {
+      throw new Error('Derivable descriptor requires an index');
+    }
+    return descriptor.atDerivationIndex(index);
+  } else {
+    if (index !== undefined) {
+      throw new Error('Definite descriptor cannot be derived with index');
+    }
+    return descriptor;
+  }
+}
+
+export function getDescriptorAtIndexCheckScript(
+  descriptor: Descriptor,
+  index: number | undefined,
+  script: Buffer,
+  descriptorString = descriptor.toString()
+): Descriptor {
+  assert(descriptor instanceof Descriptor);
+  const descriptorAtIndex = getDescriptorAtIndex(descriptor, index);
+  if (!script.equals(descriptorAtIndex.scriptPubkey())) {
+    throw new Error(`Script mismatch: descriptor ${descriptorString} script=${script.toString('hex')}`);
+  }
+  assert(descriptorAtIndex instanceof Descriptor);
+  return descriptorAtIndex;
+}

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

@@ -33,7 +33,7 @@ export function mockDerivedDescriptorWalletOutput(
     hash,
     index: vout,
     witnessUtxo: {
-      script: createScriptPubKeyFromDescriptor(descriptor),
+      script: createScriptPubKeyFromDescriptor(descriptor, undefined),
       value,
     },
     descriptor,
@@ -67,7 +67,7 @@ export function mockPsbt(
     outputs.map((o) => {
       const derivedDescriptor = tryDeriveAtIndex(o.descriptor, o.index);
       return {
-        script: createScriptPubKeyFromDescriptor(derivedDescriptor),
+        script: createScriptPubKeyFromDescriptor(derivedDescriptor, undefined),
         value: o.value,
         descriptor: o.external ? undefined : derivedDescriptor,
       };

modules/utxo-core/test/descriptor/Output.ts

@@ -1,25 +1,48 @@
 import assert from 'assert';
 
-import { Descriptor } from '@bitgo/wasm-miniscript';
-
-import { isExternalOutput, isInternalOutput } from '../../src/descriptor';
+import { isExternalOutput, isInternalOutput, toDerivedDescriptorWalletOutput } from '../../src/descriptor/Output';
+import { getDescriptor } from '../../src/testutil/descriptor';
+import { createScriptPubKeyFromDescriptor } from '../../src/descriptor';
 
 describe('decscriptor.Output', function () {
-  const mockDescriptor = {} as Descriptor;
+  const descriptor = getDescriptor('Wsh2Of3');
 
   it('isInternalOutput correctly identifies internal outputs', function () {
-    const internalOutput = { value: 1n, descriptor: mockDescriptor };
+    const internalOutput = { value: 1n, descriptor };
     const externalOutput = { value: 1n };
 
     assert.strictEqual(isInternalOutput(internalOutput), true);
     assert.strictEqual(isInternalOutput(externalOutput), false);
   });
 
   it('isExternalOutput correctly identifies external outputs', function () {
-    const internalOutput = { value: 1n, descriptor: mockDescriptor };
+    const internalOutput = { value: 1n, descriptor };
     const externalOutput = { value: 1n };
 
     assert.strictEqual(isExternalOutput(internalOutput), false);
     assert.strictEqual(isExternalOutput(externalOutput), true);
   });
+
+  it('toDerivedDescriptorWalletOutput returns expected values', function () {
+    const derivable = descriptor;
+    const definite = derivable.atDerivationIndex(0);
+    for (const descriptor of [derivable, definite]) {
+      const descriptorIndex = descriptor === derivable ? 0 : undefined;
+      const descriptorMap = new Map([['desc', descriptor]]);
+      const descriptorWalletOutput = {
+        hash: Buffer.alloc(32).toString('hex'),
+        index: 0,
+        witnessUtxo: {
+          script: createScriptPubKeyFromDescriptor(descriptor, descriptorIndex),
+          value: 1n,
+        },
+        descriptorName: 'desc',
+        descriptorIndex,
+      };
+      assert.strictEqual(
+        toDerivedDescriptorWalletOutput(descriptorWalletOutput, descriptorMap).descriptor.toString(),
+        definite.toString()
+      );
+    }
+  });
 });

modules/utxo-core/test/descriptor/derive.ts

@@ -0,0 +1,25 @@
+import assert from 'assert';
+
+import { getDescriptor } from '../../src/testutil/descriptor';
+import { getDescriptorAtIndex, getDescriptorAtIndexCheckScript } from '../../src/descriptor/derive';
+
+describe('derive', function () {
+  const derivable = getDescriptor('Wsh2Of3');
+  const definite = derivable.atDerivationIndex(0);
+
+  it('getDescriptorAtIndex', function () {
+    assert(derivable.hasWildcard());
+    assert(!definite.hasWildcard());
+    assert.strictEqual(getDescriptorAtIndex(derivable, 0).toString(), definite.toString());
+    assert.strictEqual(getDescriptorAtIndex(definite, undefined).toString(), definite.toString());
+    assert.throws(() => getDescriptorAtIndex(derivable, undefined), /Derivable descriptor requires an index/);
+    assert.throws(() => getDescriptorAtIndex(definite, 0), /Definite descriptor cannot be derived with index/);
+  });
+
+  it('getDescriptorAtIndexCheckScript', function () {
+    const script0 = Buffer.from(derivable.atDerivationIndex(0).scriptPubkey());
+    const script1 = Buffer.from(derivable.atDerivationIndex(1).scriptPubkey());
+    assert.strictEqual(getDescriptorAtIndexCheckScript(derivable, 0, script0).toString(), definite.toString());
+    assert.throws(() => getDescriptorAtIndexCheckScript(derivable, 0, script1), /Script mismatch/);
+  });
+});