refactor(abstract-utxo): offline vault applies the descriptor policy spend rules

modules/abstract-utxo/src/descriptor/pox5SpendPolicy.ts

@@ -10,13 +10,23 @@ function getPox5InputMatches(psbt: Psbt, descriptorMap: descriptorWallet.Descrip
 /** WP builds early exits with a zero nLockTime; the early-exit branch does not evaluate the CLTV. */
 const POX5_EARLY_EXIT_LOCKTIME = 0;
 
+/**
+ * The early-exit checks beyond what the shared `pox5.assertPox5EarlyExitSpend` does (it only checks that the
+ * input matches the descriptor): nLockTime 0 and exactly one principal preimage that the descriptor commits to.
+ * `classifyPox5Spend` has already checked the descriptor match by the time this runs.
+ */
 function assertPox5EarlyExitPolicy(psbt: Psbt, match: pox5.Pox5InputMatch): void {
-  pox5.assertPox5EarlyExitSpend(psbt, match);
   const lockTime = psbt.lockTime();
   if (lockTime !== POX5_EARLY_EXIT_LOCKTIME) {
     throw new Error(`PoX-5 early exit nLockTime must be ${POX5_EARLY_EXIT_LOCKTIME}, got ${lockTime}`);
   }
-  // throws unless exactly one PSBT_IN_SHA256 record carries the descriptor's commitment and its preimage
+  const sha256Records = psbt
+    .getInputKeyValues(match.inputIndex)
+    .filter((record) => record.type === 'known' && record.key === 'PSBT_IN_SHA256');
+  if (sha256Records.length !== 1) {
+    throw new Error(`PoX-5 early exit requires exactly one PSBT_IN_SHA256 record, found ${sha256Records.length}`);
+  }
+  // throws unless the record carries the descriptor's commitment and its preimage
   pox5.getPox5PrincipalPreimage(psbt, match.inputIndex, match.info);
 }
 
@@ -37,6 +47,8 @@ export function assertPox5SpendPolicy(psbt: Psbt, descriptorMap: descriptorWalle
     throw new Error(`expected exactly one PoX-5 lockup input, found ${matches.length}`);
   }
   const [match] = matches;
+  // Any PSBT_IN_SHA256 record makes this an early exit, even one the descriptor does not commit to: a stray
+  // record is then refused by the early-exit checks instead of being carried along in a late exit.
   const branch = pox5.classifyPox5Spend(psbt, match);
   switch (branch) {
     case 'locktime':
@@ -45,7 +57,9 @@ export function assertPox5SpendPolicy(psbt: Psbt, descriptorMap: descriptorWalle
     case 'early-exit':
       assertPox5EarlyExitPolicy(psbt, match);
       return;
-    default:
-      throw new Error(`unsupported PoX-5 spend branch: ${branch}`);
+    default: {
+      const unsupported: never = branch;
+      throw new Error(`unsupported PoX-5 spend branch: ${String(unsupported)}`);
+    }
   }
 }

modules/abstract-utxo/src/offlineVault/descriptor/transaction.ts

@@ -4,6 +4,7 @@ import { bip32, Psbt } from '@bitgo/wasm-utxo';
 import { DescriptorMap, NamedDescriptor } from '../../descriptor';
 import { OfflineVaultSignable, toKeyTriple } from '../OfflineVaultSignable';
 import {
+  DescriptorValidationPolicy,
   getValidatorOneOfTemplates,
   getValidatorPox5Wallet,
   getValidatorSignedByUserKey,
@@ -13,7 +14,6 @@ import {
 import { explainPsbt, signPsbt } from '../../transaction/descriptor';
 import { TransactionExplanation } from '../TransactionExplanation';
 import { UtxoCoinName } from '../../names';
-import { assertPox5SpendPolicy } from '../../descriptor/pox5SpendPolicy';
 
 export const DescriptorTransaction = t.intersection(
   [OfflineVaultSignable, t.type({ descriptors: t.array(NamedDescriptor) })],
@@ -22,24 +22,30 @@ export const DescriptorTransaction = t.intersection(
 
 export type DescriptorTransaction = t.TypeOf<typeof DescriptorTransaction>;
 
-export function getDescriptorsFromDescriptorTransaction(tx: DescriptorTransaction): DescriptorMap {
-  const { descriptors, xpubsWithDerivationPath } = tx;
-  const pubkeys = toKeyTriple(xpubsWithDerivationPath);
-  const policy = getValidatorSome([
+function getDescriptorPolicy(): DescriptorValidationPolicy {
+  return getValidatorSome([
     // allow all 2-of-3-ish descriptors where the keys match the wallet keys
     getValidatorOneOfTemplates(['Wsh2Of3', 'Wsh2Of3CltvDrop', 'ShWsh2Of3CltvDrop']),
     // allow 2-of-3 wallet descriptors together with PoX-5 lockup descriptors where the keys match the wallet keys
     getValidatorPox5Wallet(),
     // allow all descriptors signed by the user key
     getValidatorSignedByUserKey(),
   ]);
-  return toDescriptorMapValidate(descriptors, pubkeys, policy);
+}
+
+export function getDescriptorsFromDescriptorTransaction(
+  tx: DescriptorTransaction,
+  policy = getDescriptorPolicy()
+): DescriptorMap {
+  const { descriptors, xpubsWithDerivationPath } = tx;
+  return toDescriptorMapValidate(descriptors, toKeyTriple(xpubsWithDerivationPath), policy);
 }
 
 export function getHalfSignedPsbt(tx: DescriptorTransaction, prv: bip32.BIP32Interface, coinName: UtxoCoinName): Psbt {
   const psbt = Psbt.deserialize(Buffer.from(tx.coinSpecific.txHex, 'hex'));
-  const descriptorMap = getDescriptorsFromDescriptorTransaction(tx);
-  assertPox5SpendPolicy(psbt, descriptorMap);
+  const policy = getDescriptorPolicy();
+  const descriptorMap = getDescriptorsFromDescriptorTransaction(tx, policy);
+  policy.validateSpend?.(psbt, descriptorMap);
   signPsbt(psbt, descriptorMap, prv, { onUnknownInput: 'throw' });
   return psbt;
 }

modules/abstract-utxo/test/unit/offlineVault/pox5.ts

@@ -62,6 +62,8 @@ type Pox5TxParams = {
   lockupInputs?: number;
   /** add the principal preimage of the lockup input, which makes the spend an early exit */
   earlyExitPreimage?: Uint8Array;
+  /** an additional principal preimage record on the lockup input */
+  extraPreimage?: Uint8Array;
 };
 
 function createPox5Tx(params: Pox5TxParams = {}): { tx: DescriptorTransaction; rootPrv: BIP32 } {
@@ -90,6 +92,9 @@ function createPox5Tx(params: Pox5TxParams = {}): { tx: DescriptorTransaction; r
   if (params.earlyExitPreimage) {
     psbt.addSha256Preimage(0, params.earlyExitPreimage);
   }
+  if (params.extraPreimage) {
+    psbt.addSha256Preimage(0, params.extraPreimage);
+  }
 
   const tx: DescriptorTransaction = {
     xpubsWithDerivationPath: {
@@ -195,6 +200,20 @@ describe('OfflineVault PoX-5 lockup descriptor', function () {
           );
         });
 
+        it('rejects a second principal preimage record that the descriptor does not commit to', function () {
+          assert.throws(
+            () => halfSign(coinName, earlyExit({ extraPreimage: Buffer.alloc(32, 0x43) })),
+            /requires exactly one PSBT_IN_SHA256 record/
+          );
+        });
+
+        it('rejects a late exit that carries a stray principal preimage record', function () {
+          assert.throws(
+            () => halfSign(coinName, createPox5Tx({ extraPreimage: Buffer.alloc(32, 0x43) })),
+            /nLockTime must be 0/
+          );
+        });
+
         it('rejects lockup descriptors with foreign staker keys', function () {
           assert.throws(() => halfSign(coinName, earlyExit({ lockupKeys: getForeignKeys() })), /does not match policy/);
         });

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

@@ -55,7 +55,13 @@ export function assertPox5LocktimeSpend(psbt: Psbt, inputs: readonly Pox5InputMa
   }
 }
 
-/** Validate that a canonical PoX-5 input can use the principal-preimage branch. */
+/**
+ * Validate that a canonical PoX-5 input can use the principal-preimage branch.
+ *
+ * This only checks that the descriptor match belongs to the PSBT input. It does not check the nLockTime or the
+ * principal preimage, so callers that sign or finalize an early exit must check those themselves
+ * (see `assertPox5SpendPolicy` in `@bitgo/abstract-utxo`, which requires nLockTime 0 and the committed preimage).
+ */
 export function assertPox5EarlyExitSpend(psbt: Psbt, input: Pox5InputMatch): void {
   getMatchedTransactionInput(psbt, input);
 }