feat(utxo-staking): add PoX-5 recovery policy

modules/utxo-descriptors/src/pox5/descriptor.ts

@@ -57,7 +57,7 @@ function validateParams(params: Pox5LockupDescriptorParams): void {
  * Build the canonical PoX-5 P2WSH descriptor. The post-CLTV and early-exit
  * paths share BitGo's standard 2-of-3 compressed-key multisig tail.
  */
-export function createPox5LockupDescriptor(params: Pox5LockupDescriptorParams): string {
+export function createPox5LockupDescriptor(params: Pox5LockupDescriptorParams): Descriptor {
   validateParams(params);
   const stakerKeys = params.stakerKeys.map((key, index) => asDescriptorKey(key, `stakerKeys[${index}]`));
   const miniscript: ast.MiniscriptNode = {
@@ -76,18 +76,17 @@ export function createPox5LockupDescriptor(params: Pox5LockupDescriptorParams):
       { multi: [2, ...stakerKeys] },
     ],
   };
-  return ast.formatNode({ wsh: miniscript });
+  const descriptorString = ast.formatNode({ wsh: miniscript });
+  return Descriptor.fromString(descriptorString, isBip32Triple(params.stakerKeys) ? 'derivable' : 'definite');
 }
 
 /** Compile the PoX-5 P2WSH scriptPubKey at a BIP32 derivation index. */
 export function createPox5LockupScriptPubKey(params: Pox5LockupDescriptorParams, derivationIndex = 0): Buffer {
   const descriptor = createPox5LockupDescriptor(params);
   if (isBip32Triple(params.stakerKeys)) {
-    return Buffer.from(
-      Descriptor.fromString(descriptor, 'derivable').atDerivationIndex(derivationIndex).scriptPubkey()
-    );
+    return Buffer.from(descriptor.atDerivationIndex(derivationIndex).scriptPubkey());
   }
-  return Buffer.from(Descriptor.fromString(descriptor, 'definite').scriptPubkey());
+  return Buffer.from(descriptor.scriptPubkey());
 }
 
 /** Derive the compressed staker keys needed to prepare a witness at an index. */

modules/utxo-descriptors/src/pox5/index.ts

@@ -1,2 +1,4 @@
 export * from './descriptor';
 export * from './parseDescriptor';
+export * from './input';
+export * from './validation';

modules/utxo-descriptors/src/pox5/input.ts

@@ -0,0 +1,63 @@
+import { Descriptor, Psbt, descriptorWallet } from '@bitgo/wasm-utxo';
+
+import { Pox5DescriptorInfo, parsePox5LockupDescriptor } from './parseDescriptor';
+
+type ResolvedPox5DescriptorInfo = Pox5DescriptorInfo & {
+  stakerKeys: [Buffer, Buffer, Buffer];
+};
+
+export type Pox5InputMatch = {
+  /** The concrete descriptor whose script matches the input. */
+  descriptor: Descriptor;
+  /** The derivation index used for a wildcard descriptor, if any. */
+  index: number | undefined;
+  info: ResolvedPox5DescriptorInfo;
+};
+
+function getConcreteDescriptor(descriptor: Descriptor, index: number | undefined): Descriptor {
+  return index === undefined ? descriptor : descriptor.atDerivationIndex(index);
+}
+
+/**
+ * Find and parse a canonical PoX-5 descriptor for a native PSBT input projection.
+ *
+ * Foreign root derivations are ignored by the shared native descriptor matcher. A
+ * match is returned only when all three staker keys can be resolved at the matched
+ * descriptor index.
+ */
+export function findPox5DescriptorForInput(
+  input: descriptorWallet.PsbtInput,
+  descriptors: descriptorWallet.DescriptorMap
+): Pox5InputMatch | undefined {
+  try {
+    const matched = descriptorWallet.findDescriptorForInput(input, descriptors);
+    if (!matched) {
+      return undefined;
+    }
+    const descriptor = getConcreteDescriptor(matched.descriptor, matched.index);
+    const info = parsePox5LockupDescriptor(descriptor);
+    if (!info?.stakerKeys) {
+      return undefined;
+    }
+    return {
+      descriptor,
+      index: matched.index,
+      info: {
+        ...info,
+        stakerKeys: info.stakerKeys,
+      },
+    };
+  } catch {
+    return undefined;
+  }
+}
+
+/** Find and parse a canonical PoX-5 descriptor for one native PSBT input. */
+export function matchPox5Input(
+  psbt: Psbt,
+  inputIndex: number,
+  descriptors: descriptorWallet.DescriptorMap
+): Pox5InputMatch | undefined {
+  const input = psbt.getInputs()[inputIndex];
+  return input ? findPox5DescriptorForInput(input, descriptors) : undefined;
+}

modules/utxo-descriptors/src/pox5/parseDescriptor.ts

@@ -1,7 +1,8 @@
-import { BIP32, Descriptor, ast } from '@bitgo/wasm-utxo';
-import { Pattern, PatternMatcher } from '@bitgo/utxo-core/descriptor';
+import { BIP32, Descriptor, ast, descriptorWallet } from '@bitgo/wasm-utxo';
 
-export type ParsedPox5LockupDescriptor = {
+type Pattern = descriptorWallet.Pattern;
+
+export type Pox5DescriptorInfo = {
   unlockHeight: number;
   stakerCommitment: Buffer;
   earlyExitKey: Buffer;
@@ -10,6 +11,9 @@ export type ParsedPox5LockupDescriptor = {
   miniscriptNode: ast.MiniscriptNode;
 };
 
+/** @deprecated Use Pox5DescriptorInfo instead. */
+export type ParsedPox5LockupDescriptor = Pox5DescriptorInfo;
+
 const COMPRESSED_KEY = /^(02|03)[0-9a-fA-F]{64}$/;
 const XPUB_WITH_INDEX = /^([1-9A-HJ-NP-Za-km-z]+)\/(\d+)$/;
 
@@ -53,17 +57,15 @@ function resolveStakerKey(value: string): Buffer | undefined {
 
 /**
  * Parse only the canonical PoX-5 descriptor template. Other descriptors return
- * null; malformed fields within the template throw so callers cannot finalize
+ * undefined; malformed fields within the template throw so callers cannot finalize
  * a script under an ambiguous policy.
  */
-export function parsePox5LockupDescriptor(
-  descriptor: Descriptor | ast.DescriptorNode
-): ParsedPox5LockupDescriptor | null {
-  const matcher = new PatternMatcher();
+export function parsePox5LockupDescriptor(descriptor: Descriptor | ast.DescriptorNode): Pox5DescriptorInfo | undefined {
+  const matcher = new descriptorWallet.PatternMatcher();
   const descriptorNode = descriptor instanceof Descriptor ? ast.fromDescriptor(descriptor) : descriptor;
   const matched = matcher.match(descriptorNode, { wsh: { $var: 'miniscript' } });
   if (!matched) {
-    return null;
+    return undefined;
   }
 
   const miniscriptNode = matched.miniscript as ast.MiniscriptNode;
@@ -80,7 +82,7 @@ export function parsePox5LockupDescriptor(
   };
   const fields = matcher.match(miniscriptNode, pattern);
   if (!fields) {
-    return null;
+    return undefined;
   }
 
   const unlockHeight = asNumber(fields.unlockHeight, 'after argument');

modules/utxo-descriptors/src/pox5/validation.ts

@@ -0,0 +1,84 @@
+import { createHash } from 'crypto';
+
+import { Psbt, type Descriptor, type PsbtInputKeyValue } from '@bitgo/wasm-utxo';
+
+import { Pox5DescriptorInfo, parsePox5LockupDescriptor } from './parseDescriptor';
+
+const SHA256_INPUT_KEY = 'PSBT_IN_SHA256';
+
+type Sha256InputKeyValue = Extract<PsbtInputKeyValue, { type: 'known' }>;
+
+function isSha256InputKeyValue(keyValue: PsbtInputKeyValue): keyValue is Sha256InputKeyValue {
+  return keyValue.type === 'known' && keyValue.key === SHA256_INPUT_KEY;
+}
+
+/**
+ * Read and validate the unique PoX-5 principal preimage from a native PSBT input.
+ * The PSBT_IN_SHA256 key data is the digest and its value is the preimage.
+ */
+export function getPox5PrincipalPreimage(psbt: Psbt, inputIndex: number): Buffer {
+  const records = psbt.getInputKeyValues(inputIndex).filter(isSha256InputKeyValue);
+  if (records.length !== 1) {
+    throw new Error(`expected exactly one ${SHA256_INPUT_KEY} record, found ${records.length}`);
+  }
+
+  const [record] = records;
+  if (record.keyData.length !== 32) {
+    throw new Error(`${SHA256_INPUT_KEY} digest must be 32 bytes`);
+  }
+  if (record.value.length !== 32) {
+    throw new Error(`${SHA256_INPUT_KEY} preimage must be 32 bytes`);
+  }
+
+  const preimage = Buffer.from(record.value);
+  const digest = createHash('sha256').update(preimage).digest();
+  if (!digest.equals(Buffer.from(record.keyData))) {
+    throw new Error(`${SHA256_INPUT_KEY} digest does not match its preimage`);
+  }
+  return preimage;
+}
+
+function isPox5DescriptorInfo(value: unknown): value is Pox5DescriptorInfo {
+  return (
+    value !== null &&
+    typeof value === 'object' &&
+    'stakerCommitment' in value &&
+    Buffer.isBuffer(value.stakerCommitment)
+  );
+}
+
+function getPox5DescriptorInfo(
+  descriptor: Pox5DescriptorInfo | Descriptor | import('@bitgo/wasm-utxo').ast.DescriptorNode
+): Pox5DescriptorInfo {
+  if (isPox5DescriptorInfo(descriptor)) {
+    return descriptor;
+  }
+  const info = parsePox5LockupDescriptor(descriptor);
+  if (!info) {
+    throw new Error('descriptor is not a canonical PoX-5 lockup descriptor');
+  }
+  return info;
+}
+
+/** Verify that a principal preimage is committed by a canonical PoX-5 descriptor. */
+export function assertPox5PrincipalPreimage(
+  descriptor: Pox5DescriptorInfo | Descriptor | import('@bitgo/wasm-utxo').ast.DescriptorNode,
+  principalPreimage: Uint8Array
+): void {
+  if (principalPreimage.length !== 32) {
+    throw new Error('principalPreimage must be 32 bytes');
+  }
+  const info = getPox5DescriptorInfo(descriptor);
+  const digest = createHash('sha256').update(principalPreimage).digest();
+  if (!digest.equals(info.stakerCommitment)) {
+    throw new Error('principalPreimage does not match the descriptor stakerCommitment');
+  }
+}
+
+/** Alias for callers that prefer validation terminology. */
+export function validatePox5PrincipalPreimage(
+  descriptor: Pox5DescriptorInfo | Descriptor | import('@bitgo/wasm-utxo').ast.DescriptorNode,
+  principalPreimage: Uint8Array
+): void {
+  assertPox5PrincipalPreimage(descriptor, principalPreimage);
+}

modules/utxo-descriptors/test/unit/pox5/descriptor.ts

@@ -93,8 +93,7 @@ describe('PoX-5 lockup descriptors', function () {
       0
     );
     const definite = { ...derivable, stakerKeys };
-    const descriptorString = createPox5LockupDescriptor(definite);
-    const descriptor = Descriptor.fromString(descriptorString, 'definite');
+    const descriptor = createPox5LockupDescriptor(definite);
     const localWitnessScript = asmToScript(descriptor.toAsmString());
     const unlockBytes = encodeTwoOfThreeUnlock(stakerKeys);
     const earlyUnlockBytes = buildUnlockScript(definite.earlyExitKey);
@@ -137,7 +136,7 @@ describe('PoX-5 lockup descriptors', function () {
 
   it('supports derivation and preserves wildcard keys until an index is selected', function () {
     const value = params();
-    const descriptor = Descriptor.fromString(createPox5LockupDescriptor(value), 'derivable');
+    const descriptor = createPox5LockupDescriptor(value);
     const wildcard = parsePox5LockupDescriptor(descriptor);
     const derived = parsePox5LockupDescriptor(descriptor.atDerivationIndex(4));
 
@@ -162,6 +161,6 @@ describe('PoX-5 lockup descriptors', function () {
       )
     );
     const validKey = params().earlyExitKey.toString('hex');
-    assert.strictEqual(parsePox5LockupDescriptor(Descriptor.fromString(`wsh(pk(${validKey}))`, 'definite')), null);
+    assert.strictEqual(parsePox5LockupDescriptor(Descriptor.fromString(`wsh(pk(${validKey}))`, 'definite')), undefined);
   });
 });

modules/utxo-descriptors/test/unit/pox5/input.ts

@@ -0,0 +1,105 @@
+import * as assert from 'assert';
+import { createHash } from 'crypto';
+
+import { bip32, descriptorWallet, Psbt } from '@bitgo/wasm-utxo';
+import { getKeyTriple } from '@bitgo/wasm-utxo/testutils';
+
+import {
+  createPox5LockupDescriptor,
+  derivePox5StakerKeys,
+  findPox5DescriptorForInput,
+  matchPox5Input,
+  Pox5LockupDescriptorParams,
+} from '../../../src/pox5';
+
+const UNLOCK_HEIGHT = 840_000;
+type BIP32Interface = bip32.BIP32Interface;
+type Pox5Bip32Params = Omit<Pox5LockupDescriptorParams, 'stakerKeys'> & {
+  stakerKeys: [BIP32Interface, BIP32Interface, BIP32Interface];
+};
+
+function sha256(value: Uint8Array): Buffer {
+  return createHash('sha256').update(value).digest();
+}
+
+function getParams(): Pox5Bip32Params {
+  const stakerKeys = getKeyTriple('utxo-descriptors-pox5');
+  return {
+    unlockHeight: UNLOCK_HEIGHT,
+    stakerCommitment: sha256(Buffer.alloc(32, 0x42)),
+    earlyExitKey: Buffer.from(stakerKeys[0].derive(9).publicKey),
+    stakerKeys,
+  };
+}
+
+function createInputPsbt(descriptor: ReturnType<typeof createPox5LockupDescriptor>): Psbt {
+  return descriptorWallet.createPsbt(
+    { version: 2, locktime: 0 },
+    [
+      {
+        hash: '01'.repeat(32),
+        index: 0,
+        witnessUtxo: { script: descriptor.scriptPubkey(), value: 100_000n },
+        descriptor,
+      },
+    ],
+    []
+  );
+}
+
+describe('PoX-5 input resolution', function () {
+  it('matches a definite descriptor without derivation metadata', function () {
+    const descriptor = createPox5LockupDescriptor({
+      ...getParams(),
+      stakerKeys: derivePox5StakerKeys(getParams().stakerKeys, 0),
+    });
+    const psbt = Psbt.create(2, 0);
+    psbt.addInput('01'.repeat(32), 0, 100_000n, descriptor.scriptPubkey());
+    const match = matchPox5Input(psbt, 0, new Map([['lockup', descriptor]]));
+
+    assert.ok(match);
+    assert.strictEqual(match.index, undefined);
+    assert.strictEqual(match.descriptor.toString(), descriptor.toString());
+    assert.deepStrictEqual(match.info.stakerKeys, derivePox5StakerKeys(getParams().stakerKeys, 0));
+  });
+
+  it('matches a derived descriptor while ignoring foreign root derivations', function () {
+    const params = getParams();
+    const descriptor = createPox5LockupDescriptor(params);
+    const concreteDescriptor = descriptor.atDerivationIndex(4);
+    const psbt = createInputPsbt(concreteDescriptor);
+    const match = matchPox5Input(psbt, 0, new Map([['lockup', descriptor]]));
+
+    assert.ok(match);
+    assert.strictEqual(match.index, 4);
+    assert.strictEqual(match.descriptor.toString(), concreteDescriptor.toString());
+    assert.deepStrictEqual(match.info.stakerKeys, derivePox5StakerKeys(params.stakerKeys, 4));
+    assert.ok(psbt.getInputs()[0]?.bip32Derivation.some((derivation) => derivation.path === ''));
+  });
+
+  it('returns no match when the input has no usable derivation metadata', function () {
+    const params = getParams();
+    const descriptor = createPox5LockupDescriptor(params);
+    const concreteDescriptor = descriptor.atDerivationIndex(4);
+    const input = {
+      witnessUtxo: { script: concreteDescriptor.scriptPubkey(), value: 100_000n },
+      bip32Derivation: [],
+      tapBip32Derivation: [],
+    };
+
+    assert.strictEqual(findPox5DescriptorForInput(input, new Map([['lockup', descriptor]])), undefined);
+  });
+
+  it('returns no match for a noncanonical descriptor', function () {
+    const params = getParams();
+    const pox5Descriptor = createPox5LockupDescriptor(params);
+    const psbt = Psbt.create(2, 0);
+    psbt.addInput('01'.repeat(32), 0, 100_000n, pox5Descriptor.atDerivationIndex(0).scriptPubkey());
+    const nonPox5Descriptor = createPox5LockupDescriptor({
+      ...params,
+      unlockHeight: UNLOCK_HEIGHT + 1,
+    }).atDerivationIndex(0);
+
+    assert.strictEqual(matchPox5Input(psbt, 0, new Map([['other', nonPox5Descriptor]])), undefined);
+  });
+});