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

Move native PoX-5 branch classification, locktime and sequence policy, preimage
preparation, and witness finalization into @bitgo/utxo-staking.

Refs: WAL-2022

Author: OttoAllmendinger

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

@@ -1 +1,2 @@
 export * from './witness';
+export * from './recovery';

modules/utxo-staking/src/pox5/recovery.ts

@@ -0,0 +1,80 @@
+import { Psbt, Transaction } from '@bitgo/wasm-utxo';
+import { pox5 } from '@bitgo/utxo-descriptors';
+
+export const POX5_MAX_UNLOCK_HEIGHT = 500_000_000;
+
+export type Pox5SpendBranch = 'locktime' | 'early-exit';
+
+export type Pox5SpendInput = pox5.Pox5InputMatch | pox5.Pox5DescriptorInfo;
+
+function getPox5DescriptorInfo(input: Pox5SpendInput): pox5.Pox5DescriptorInfo {
+  return 'info' in input ? input.info : input;
+}
+
+function assertPox5UnlockHeight(unlockHeight: number): void {
+  if (!Number.isSafeInteger(unlockHeight) || unlockHeight <= 0 || unlockHeight >= POX5_MAX_UNLOCK_HEIGHT) {
+    throw new Error(`PoX-5 unlock height must be a positive block height below ${POX5_MAX_UNLOCK_HEIGHT}`);
+  }
+}
+
+function assertPox5BlockHeightLocktime(lockTime: number): void {
+  if (!Number.isSafeInteger(lockTime) || lockTime < 0 || lockTime >= POX5_MAX_UNLOCK_HEIGHT) {
+    throw new Error(`PoX-5 nLockTime must be a block height below ${POX5_MAX_UNLOCK_HEIGHT}`);
+  }
+}
+
+function hasFinalInput(psbt: Psbt): boolean {
+  return Transaction.fromBytes(psbt.getUnsignedTx())
+    .getInputs()
+    .some((input) => input.sequence === 0xffffffff);
+}
+
+/** Classify a canonical PoX-5 input by the transaction branch it can spend. */
+export function classifyPox5Spend(psbt: Psbt, input: pox5.Pox5InputMatch): Pox5SpendBranch {
+  const { unlockHeight } = input.info;
+  assertPox5UnlockHeight(unlockHeight);
+  const lockTime = psbt.lockTime();
+  assertPox5BlockHeightLocktime(lockTime);
+  return lockTime >= unlockHeight ? 'locktime' : 'early-exit';
+}
+
+/** Validate the post-CLTV policy for all canonical PoX-5 inputs in a recovery PSBT. */
+export function assertPox5LocktimeSpend(psbt: Psbt, inputs: readonly Pox5SpendInput[]): void {
+  if (inputs.length === 0) {
+    throw new Error('PoX-5 lockup descriptor match is required');
+  }
+
+  const unlockHeights = inputs.map((input) => {
+    const { unlockHeight } = getPox5DescriptorInfo(input);
+    assertPox5UnlockHeight(unlockHeight);
+    return unlockHeight;
+  });
+  const lockTime = psbt.lockTime();
+  assertPox5BlockHeightLocktime(lockTime);
+  const requiredLockTime = Math.max(...unlockHeights);
+  if (lockTime < requiredLockTime) {
+    throw new Error(`PoX-5 nLockTime must be at least ${requiredLockTime}`);
+  }
+  if (hasFinalInput(psbt)) {
+    throw new Error('PoX-5 locktime spend inputs must use non-final sequences');
+  }
+}
+
+/** Validate that a canonical PoX-5 input uses the principal-preimage branch. */
+export function assertPox5EarlyExitSpend(psbt: Psbt, input: pox5.Pox5InputMatch): void {
+  if (classifyPox5Spend(psbt, input) !== 'early-exit') {
+    throw new Error('PoX-5 input is not an early-exit spend');
+  }
+}
+
+/** Add validated principal-preimage metadata for an early-exit spend. */
+export function preparePox5EarlyExit(
+  psbt: Psbt,
+  inputIndex: number,
+  input: pox5.Pox5InputMatch,
+  principalPreimage: Uint8Array
+): void {
+  assertPox5EarlyExitSpend(psbt, input);
+  pox5.assertPox5PrincipalPreimage(input.info, principalPreimage);
+  psbt.addSha256Preimage(inputIndex, principalPreimage);
+}

modules/utxo-staking/src/pox5/witness.ts

@@ -1,58 +1,29 @@
-import { createHash } from 'crypto';
-
-import { ast, Descriptor, Psbt } from '@bitgo/wasm-utxo';
+import { Psbt } from '@bitgo/wasm-utxo';
 import { pox5 } from '@bitgo/utxo-descriptors';
 
-type Pox5Descriptor = Descriptor | ast.DescriptorNode;
+import { assertPox5EarlyExitSpend, assertPox5LocktimeSpend, preparePox5EarlyExit } from './recovery';
 
 export type Pox5FinalizerParams = {
-  /** A definite or derivation-indexed canonical PoX-5 descriptor. */
-  descriptor: Pox5Descriptor;
-  /** The derived user, backup, and BitGo keys in descriptor order. */
-  stakerKeys: [Buffer, Buffer, Buffer];
+  /** The canonical descriptor match for the input being finalized. */
+  match: pox5.Pox5InputMatch;
 };
 
-function getParsedDescriptor(params: Pox5FinalizerParams) {
-  const parsed = pox5.parsePox5LockupDescriptor(params.descriptor);
-  if (!parsed || !parsed.stakerKeys) {
-    throw new Error('descriptor must be a definite or derivation-indexed canonical PoX-5 descriptor');
-  }
-  if (!parsed.stakerKeys.every((key, index) => key.equals(params.stakerKeys[index]))) {
-    throw new Error('stakerKeys must match the canonical descriptor order');
-  }
-  return parsed;
-}
-
-function getDescriptor(descriptor: Pox5Descriptor): Descriptor {
-  return descriptor instanceof Descriptor ? descriptor : Descriptor.fromString(ast.formatNode(descriptor), 'definite');
-}
-
-function prepareInput(psbt: Psbt, inputIndex: number, params: Pox5FinalizerParams) {
-  const parsed = getParsedDescriptor(params);
-  psbt.updateInputWithDescriptor(inputIndex, getDescriptor(params.descriptor));
-  return parsed;
-}
-
 /** Finalize the post-CLTV 2-of-3 PoX-5 spend branch. */
 export function finalizePox5LocktimePath(psbt: Psbt, inputIndex: number, params: Pox5FinalizerParams): void {
-  const parsed = prepareInput(psbt, inputIndex, params);
-  if (psbt.lockTime() < parsed.unlockHeight) {
-    throw new Error(`transaction locktime must be at least ${parsed.unlockHeight}`);
-  }
+  assertPox5LocktimeSpend(psbt, [params.match]);
+  psbt.updateInputWithDescriptor(inputIndex, params.match.descriptor);
   psbt.finalizeInput(inputIndex);
 }
 
 /** Finalize the principal-preimage early-exit 2-of-3 PoX-5 spend branch. */
 export function finalizePox5EarlyExitPath(
   psbt: Psbt,
   inputIndex: number,
-  params: Pox5FinalizerParams & { principalPreimage: Buffer }
+  params: Pox5FinalizerParams & { principalPreimage: Uint8Array }
 ): void {
-  const parsed = prepareInput(psbt, inputIndex, params);
-  const preimageHash = createHash('sha256').update(params.principalPreimage).digest();
-  if (!preimageHash.equals(parsed.stakerCommitment)) {
-    throw new Error('principalPreimage does not match the descriptor stakerCommitment');
-  }
-  psbt.addSha256Preimage(inputIndex, params.principalPreimage);
+  assertPox5EarlyExitSpend(psbt, params.match);
+  pox5.assertPox5PrincipalPreimage(params.match.info, params.principalPreimage);
+  preparePox5EarlyExit(psbt, inputIndex, params.match, params.principalPreimage);
+  psbt.updateInputWithDescriptor(inputIndex, params.match.descriptor);
   psbt.finalizeInput(inputIndex);
 }

modules/utxo-staking/test/unit/pox5/recovery.ts

@@ -0,0 +1,101 @@
+import assert from 'assert/strict';
+import { createHash } from 'crypto';
+
+import { pox5 } from '@bitgo/utxo-descriptors';
+import { Psbt, type Descriptor } from '@bitgo/wasm-utxo';
+import { getKey, getKeyTriple } from '@bitgo/wasm-utxo/testutils';
+
+import {
+  assertPox5EarlyExitSpend,
+  assertPox5LocktimeSpend,
+  classifyPox5Spend,
+  POX5_MAX_UNLOCK_HEIGHT,
+  preparePox5EarlyExit,
+} from '../../../src/pox5';
+
+type Pox5InputMatch = pox5.Pox5InputMatch;
+
+const UNLOCK_HEIGHT = 840_000;
+
+function sha256(value: Uint8Array): Buffer {
+  return createHash('sha256').update(value).digest();
+}
+
+function createPox5RecoveryPsbt(
+  lockTime: number,
+  sequence = 0xfffffffe,
+  unlockHeight = UNLOCK_HEIGHT
+): {
+  psbt: Psbt;
+  match: Pox5InputMatch;
+  principalPreimage: Buffer;
+} {
+  const [user, backup, bitgo] = getKeyTriple('utxo-staking-pox5-recovery');
+  const earlyExit = getKey('utxo-staking-pox5-recovery-early-exit');
+  const principalPreimage = Buffer.alloc(32, 0x42);
+  const descriptor = pox5.createPox5LockupDescriptor({
+    unlockHeight,
+    stakerCommitment: sha256(principalPreimage),
+    earlyExitKey: Buffer.from(earlyExit.publicKey),
+    stakerKeys: [Buffer.from(user.publicKey), Buffer.from(backup.publicKey), Buffer.from(bitgo.publicKey)],
+  });
+  const concreteDescriptor = descriptor as Descriptor;
+  const psbt = Psbt.create(2, lockTime);
+  psbt.addInput('01'.repeat(32), 0, 100_000n, concreteDescriptor.scriptPubkey(), sequence);
+  psbt.addOutput(concreteDescriptor.scriptPubkey(), 90_000n);
+  psbt.updateInputWithDescriptor(0, concreteDescriptor);
+
+  const match = pox5.matchPox5Input(psbt, 0, new Map([['pox5', concreteDescriptor]]));
+  assert.ok(match);
+  return { psbt, match: match as Pox5InputMatch, principalPreimage };
+}
+
+describe('PoX-5 spend policy', function () {
+  it('classifies locktime and early-exit branches from native transaction data', function () {
+    const locktimeSpend = createPox5RecoveryPsbt(UNLOCK_HEIGHT);
+    const earlyExitSpend = createPox5RecoveryPsbt(0);
+
+    assert.equal(classifyPox5Spend(locktimeSpend.psbt, locktimeSpend.match), 'locktime');
+    assert.equal(classifyPox5Spend(earlyExitSpend.psbt, earlyExitSpend.match), 'early-exit');
+    assert.doesNotThrow(() => assertPox5LocktimeSpend(locktimeSpend.psbt, [locktimeSpend.match]));
+    assert.doesNotThrow(() => assertPox5EarlyExitSpend(earlyExitSpend.psbt, earlyExitSpend.match));
+    assert.throws(() => assertPox5EarlyExitSpend(locktimeSpend.psbt, locktimeSpend.match), /not an early-exit spend/);
+  });
+
+  it('enforces the block-height and unlock-height boundaries', function () {
+    const atHeight = createPox5RecoveryPsbt(UNLOCK_HEIGHT);
+    const aboveHeight = createPox5RecoveryPsbt(UNLOCK_HEIGHT + 1);
+    const belowHeight = createPox5RecoveryPsbt(UNLOCK_HEIGHT - 1);
+    const timestampLocktime = createPox5RecoveryPsbt(POX5_MAX_UNLOCK_HEIGHT);
+
+    assert.doesNotThrow(() => assertPox5LocktimeSpend(atHeight.psbt, [atHeight.match]));
+    assert.doesNotThrow(() => assertPox5LocktimeSpend(aboveHeight.psbt, [aboveHeight.match]));
+    assert.throws(() => assertPox5LocktimeSpend(belowHeight.psbt, [belowHeight.match]), /at least/);
+    assert.throws(
+      () => assertPox5LocktimeSpend(timestampLocktime.psbt, [timestampLocktime.match]),
+      /block height below/
+    );
+  });
+
+  it('requires non-final sequences for locktime spends', function () {
+    const final = createPox5RecoveryPsbt(UNLOCK_HEIGHT, 0xffffffff);
+    const nonFinal = createPox5RecoveryPsbt(UNLOCK_HEIGHT, 0xfffffffe);
+
+    assert.throws(() => assertPox5LocktimeSpend(final.psbt, [final.match]), /non-final sequences/);
+    assert.doesNotThrow(() => assertPox5LocktimeSpend(nonFinal.psbt, [nonFinal.match]));
+  });
+
+  it('adds a validated principal preimage through the native PSBT API', function () {
+    const earlyExitSpend = createPox5RecoveryPsbt(0);
+
+    preparePox5EarlyExit(earlyExitSpend.psbt, 0, earlyExitSpend.match, earlyExitSpend.principalPreimage);
+
+    const records = earlyExitSpend.psbt
+      .getInputKeyValues(0)
+      .filter((record) => record.type === 'known' && record.key === 'PSBT_IN_SHA256');
+    assert.equal(records.length, 1);
+    const [record] = records;
+    assert.deepStrictEqual(Buffer.from(record.keyData), sha256(earlyExitSpend.principalPreimage));
+    assert.deepStrictEqual(Buffer.from(record.value), earlyExitSpend.principalPreimage);
+  });
+});

modules/utxo-staking/test/unit/pox5/witness.ts

@@ -5,7 +5,9 @@ import { pox5 } from '@bitgo/utxo-descriptors';
 import { Psbt, type Descriptor } from '@bitgo/wasm-utxo';
 import { getKey, getKeyTriple } from '@bitgo/wasm-utxo/testutils';
 
-import { finalizePox5EarlyExitPath, finalizePox5LocktimePath, Pox5FinalizerParams } from '../../../src/pox5';
+import { finalizePox5EarlyExitPath, finalizePox5LocktimePath, type Pox5FinalizerParams } from '../../../src/pox5';
+
+type Pox5InputMatch = pox5.Pox5InputMatch;
 
 const UNLOCK_HEIGHT = 840_000;
 
@@ -29,21 +31,22 @@ function createPox5Psbt(
     Buffer,
     Buffer
   ];
-  const params: Pox5FinalizerParams = {
-    descriptor: pox5.createPox5LockupDescriptor({
-      unlockHeight: UNLOCK_HEIGHT,
-      stakerCommitment: sha256(principalPreimage),
-      earlyExitKey: Buffer.from(earlyExit.publicKey),
-      stakerKeys,
-    }),
+  const descriptor = pox5.createPox5LockupDescriptor({
+    unlockHeight: UNLOCK_HEIGHT,
+    stakerCommitment: sha256(principalPreimage),
+    earlyExitKey: Buffer.from(earlyExit.publicKey),
     stakerKeys,
-  };
-  const descriptor = params.descriptor as Descriptor;
-  const scriptPubKey = descriptor.scriptPubkey();
+  });
+  const paramsDescriptor = descriptor as Descriptor;
+  const scriptPubKey = paramsDescriptor.scriptPubkey();
   const psbt = Psbt.create(2, lockTime);
   psbt.addInput('01'.repeat(32), 0, 100_000n, scriptPubKey, 0xfffffffe);
   psbt.addOutput(scriptPubKey, 90_000n);
-  psbt.updateInputWithDescriptor(0, descriptor);
+  psbt.updateInputWithDescriptor(0, paramsDescriptor);
+
+  const match = pox5.matchPox5Input(psbt, 0, new Map([['pox5', paramsDescriptor]]));
+  assert.ok(match);
+  const params: Pox5FinalizerParams = { match: match as Pox5InputMatch };
 
   for (const key of includeEarlyExitSignature ? [user, backup, earlyExit] : [user, backup]) {
     assert.ok(key.privateKey, 'test key must include private key material');