refactor(utxo-descriptors): make PoX-5 a descriptor extension

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

@@ -0,0 +1,22 @@
+import { DescriptorExtension, pox5 } from '@bitgo/utxo-descriptors';
+import type { descriptorWallet, Psbt } from '@bitgo/wasm-utxo';
+
+/**
+ * Special-purpose descriptors that descriptor wallets accept next to their regular wallet descriptors.
+ * This is the one place that lists them; the descriptor flow itself only deals with `DescriptorExtension`.
+ */
+export const defaultDescriptorExtensions: readonly DescriptorExtension[] = [pox5.descriptorExtension];
+
+/**
+ * Apply the spend policy of every extension to a PSBT that is about to be signed.
+ * Throws if an extension refuses the spend.
+ */
+export function assertExtensionSpendPolicy(
+  psbt: Psbt,
+  descriptorMap: descriptorWallet.DescriptorMap,
+  extensions: readonly DescriptorExtension[] = defaultDescriptorExtensions
+): void {
+  for (const extension of extensions) {
+    extension.assertSpendPolicy(psbt, descriptorMap);
+  }
+}

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

@@ -1,8 +1,9 @@
 import { EnvironmentName, Triple } from '@bitgo/sdk-core';
-import { pox5 } from '@bitgo/utxo-descriptors';
+import type { DescriptorExtension } from '@bitgo/utxo-descriptors';
 import { bip32, Descriptor, descriptorWallet } from '@bitgo/wasm-utxo';
 
 import { DescriptorBuilder, parseDescriptor } from './builder';
+import { defaultDescriptorExtensions } from './extensions';
 import { hasValidSignature, NamedDescriptor, NamedDescriptorNative, toNamedDescriptorNative } from './NamedDescriptor';
 
 import type { DescriptorMap } from './index';
@@ -78,60 +79,39 @@ export function getValidatorSignedByUserKey(): DescriptorValidationPolicy {
   };
 }
 
-function isCanonicalPox5Lockup(descriptor: Descriptor, walletKeys: KeyTriple): boolean {
-  try {
-    const info = pox5.parsePox5LockupDescriptor(descriptor);
-    if (!info) {
-      return false;
-    }
-    // Rebuilding the canonical descriptor from the wallet keys checks both that the
-    // staker keys are the wallet keys (with the `/*` derivation) and that the
-    // descriptor has no non-canonical encoding.
-    const canonical = pox5.createPox5LockupDescriptor({
-      unlockHeight: info.unlockHeight,
-      stakerCommitment: info.stakerCommitment,
-      earlyExitKey: info.earlyExitKey,
-      stakerKeys: [walletKeys[0], walletKeys[1], walletKeys[2]],
-    });
-    return canonical.toString() === descriptor.toString();
-  } catch {
-    // malformed lockup parameters
-    return false;
-  }
-}
-
 /**
- * Accepts descriptors that are canonical PoX-5 lockup descriptors whose staker keys
- * are the wallet keys.
+ * Accepts descriptor groups where each descriptor individually satisfies at least one of
+ * the validators. Unlike `getValidatorSome`, different descriptors may match different validators.
  */
-export function getValidatorPox5Lockup(): DescriptorValidationPolicy {
+export function getValidatorEachOneOf(validators: DescriptorValidationPolicy[]): DescriptorValidationPolicy {
   return {
-    name: 'pox5Lockup',
+    name: 'eachOneOf(' + validators.map((v) => v.name).join(',') + ')',
     validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean {
-      return arr.every((d) => isCanonicalPox5Lockup(d.value, walletKeys));
+      return arr.every((d) => validators.some((v) => v.validate([d], walletKeys)));
     },
   };
 }
 
 /**
- * Accepts descriptor groups where each descriptor individually satisfies at least one of
- * the validators. Unlike `getValidatorSome`, different descriptors may match different validators.
+ * Accepts descriptors that the extension considers valid descriptors of a wallet with the given keys.
  */
-export function getValidatorEachOneOf(validators: DescriptorValidationPolicy[]): DescriptorValidationPolicy {
+export function getValidatorExtension(extension: DescriptorExtension): DescriptorValidationPolicy {
   return {
-    name: 'eachOneOf(' + validators.map((v) => v.name).join(',') + ')',
+    name: extension.name,
     validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean {
-      return arr.every((d) => validators.some((v) => v.validate([d], walletKeys)));
+      return arr.every((d) => extension.isWalletDescriptor(d.value, walletKeys));
     },
   };
 }
 
 /**
- * Accepts the descriptors of a PoX-5 staking wallet: the regular 2-of-3 wallet descriptors
- * together with PoX-5 lockup descriptors, all using the wallet keys.
+ * Accepts the descriptors of a wallet with extensions: the regular 2-of-3 wallet descriptors together with
+ * the descriptors of any of the extensions, all using the wallet keys.
  */
-export function getValidatorPox5Wallet(): DescriptorValidationPolicy {
-  return getValidatorEachOneOf([getValidatorDescriptorTemplate('Wsh2Of3'), getValidatorPox5Lockup()]);
+export function getValidatorWalletWithExtensions(
+  extensions: readonly DescriptorExtension[] = defaultDescriptorExtensions
+): DescriptorValidationPolicy {
+  return getValidatorEachOneOf([getValidatorDescriptorTemplate('Wsh2Of3'), ...extensions.map(getValidatorExtension)]);
 }
 
 export class DescriptorPolicyValidationError extends Error {
@@ -173,8 +153,8 @@ export function getPolicyForEnv(env: EnvironmentName): DescriptorValidationPolic
         getValidatorDescriptorTemplate('Wsh2Of3CltvDrop'),
         // allow sBTC peg-in deposit descriptors where the reclaim keys match the wallet keys
         getValidatorDescriptorTemplate('SbtcDeposit'),
-        // allow 2-of-3 wallet descriptors together with PoX-5 lockup descriptors (stacks staking)
-        getValidatorPox5Wallet(),
+        // allow 2-of-3 wallet descriptors together with the descriptors of extensions (e.g. staking lockups)
+        getValidatorWalletWithExtensions(),
         // allow all descriptors signed by the user key
         getValidatorSignedByUserKey(),
       ]);

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

@@ -2,20 +2,19 @@ import * as t from 'io-ts';
 import { bip32, Psbt } from '@bitgo/wasm-utxo';
 
 import { DescriptorMap, NamedDescriptor } from '../../descriptor';
+import { assertExtensionSpendPolicy } from '../../descriptor/extensions';
 import { OfflineVaultSignable, toKeyTriple } from '../OfflineVaultSignable';
 import {
   getValidatorOneOfTemplates,
-  getValidatorPox5Wallet,
   getValidatorSignedByUserKey,
   getValidatorSome,
+  getValidatorWalletWithExtensions,
   toDescriptorMapValidate,
 } from '../../descriptor/validatePolicy';
 import { explainPsbt, signPsbt } from '../../transaction/descriptor';
 import { TransactionExplanation } from '../TransactionExplanation';
 import { UtxoCoinName } from '../../names';
 
-import { assertPox5SpendPolicy } from './pox5';
-
 export const DescriptorTransaction = t.intersection(
   [OfflineVaultSignable, t.type({ descriptors: t.array(NamedDescriptor) })],
   'DescriptorTransaction'
@@ -29,8 +28,9 @@ export function getDescriptorsFromDescriptorTransaction(tx: DescriptorTransactio
   const policy = 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 2-of-3 wallet descriptors together with the descriptors of extensions (e.g. staking lockups)
+    // where the keys match the wallet keys
+    getValidatorWalletWithExtensions(),
     // allow all descriptors signed by the user key
     getValidatorSignedByUserKey(),
   ]);
@@ -40,7 +40,7 @@ export function getDescriptorsFromDescriptorTransaction(tx: DescriptorTransactio
 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);
+  assertExtensionSpendPolicy(psbt, descriptorMap);
   signPsbt(psbt, descriptorMap, prv, { onUnknownInput: 'throw' });
   return psbt;
 }

modules/utxo-descriptors/src/extension.ts

@@ -0,0 +1,22 @@
+import type { bip32, Descriptor, descriptorWallet, Psbt } from '@bitgo/wasm-utxo';
+
+/** The user, backup and BitGo keys of a wallet. */
+export type WalletKeyTriple = [bip32.BIP32Interface, bip32.BIP32Interface, bip32.BIP32Interface];
+
+/**
+ * A special-purpose descriptor (for example a staking lockup) that a descriptor wallet can hold next to its
+ * regular wallet descriptors. The generic descriptor wallet flow validates descriptors and spends through
+ * registered extensions and does not know what any of them are.
+ */
+export interface DescriptorExtension {
+  name: string;
+
+  /** Whether the descriptor is a valid descriptor of this extension for a wallet with the given keys. */
+  isWalletDescriptor(descriptor: Descriptor, walletKeys: WalletKeyTriple): boolean;
+
+  /**
+   * Throws if the PSBT spends inputs of this extension in a way that must not be signed.
+   * Does nothing when none of the inputs belongs to the extension.
+   */
+  assertSpendPolicy(psbt: Psbt, descriptorMap: descriptorWallet.DescriptorMap): void;
+}

modules/utxo-descriptors/src/index.ts

@@ -1,2 +1,3 @@
 export * as sbtc from './sbtc';
 export * as pox5 from './pox5';
+export * from './extension';

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

@@ -0,0 +1,39 @@
+import type { Descriptor } from '@bitgo/wasm-utxo';
+
+import type { DescriptorExtension, WalletKeyTriple } from '../extension';
+
+import { createPox5LockupDescriptor } from './descriptor';
+import { parsePox5LockupDescriptor } from './parseDescriptor';
+import { assertPox5SpendPolicy } from './spendPolicy';
+
+/**
+ * Whether the descriptor is a canonical PoX-5 lockup descriptor whose staker keys are the wallet keys.
+ */
+export function isWalletPox5Lockup(descriptor: Descriptor, walletKeys: WalletKeyTriple): boolean {
+  try {
+    const info = parsePox5LockupDescriptor(descriptor);
+    if (!info) {
+      return false;
+    }
+    // Rebuilding the canonical descriptor from the wallet keys checks both that the
+    // staker keys are the wallet keys (with the `/*` derivation) and that the
+    // descriptor has no non-canonical encoding.
+    const canonical = createPox5LockupDescriptor({
+      unlockHeight: info.unlockHeight,
+      stakerCommitment: info.stakerCommitment,
+      earlyExitKey: info.earlyExitKey,
+      stakerKeys: [walletKeys[0], walletKeys[1], walletKeys[2]],
+    });
+    return canonical.toString() === descriptor.toString();
+  } catch {
+    // malformed lockup parameters
+    return false;
+  }
+}
+
+/** PoX-5 staking lockups as a descriptor wallet extension. */
+export const descriptorExtension: DescriptorExtension = {
+  name: 'pox5',
+  isWalletDescriptor: isWalletPox5Lockup,
+  assertSpendPolicy: assertPox5SpendPolicy,
+};

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

@@ -3,3 +3,5 @@ export * from './parseDescriptor';
 export * from './input';
 export * from './validation';
 export * from './recovery';
+export * from './extension';
+export * from './spendPolicy';

modules/abstract-utxo/src/offlineVault/descriptor/pox5.ts → modules/utxo-descriptors/src/pox5/spendPolicy.ts

@@ -1,9 +1,11 @@
-import { pox5 } from '@bitgo/utxo-descriptors';
 import { descriptorWallet, Psbt } from '@bitgo/wasm-utxo';
 
-function getPox5InputMatches(psbt: Psbt, descriptorMap: descriptorWallet.DescriptorMap): pox5.Pox5InputMatch[] {
-  return Array.from({ length: psbt.inputCount() }, (_, i) => pox5.matchPox5Input(psbt, i, descriptorMap)).filter(
-    (match): match is pox5.Pox5InputMatch => match !== undefined
+import { Pox5InputMatch, matchPox5Input } from './input';
+import { assertPox5LocktimeSpend, classifyPox5Spend } from './recovery';
+
+function getPox5InputMatches(psbt: Psbt, descriptorMap: descriptorWallet.DescriptorMap): Pox5InputMatch[] {
+  return Array.from({ length: psbt.inputCount() }, (_, i) => matchPox5Input(psbt, i, descriptorMap)).filter(
+    (match): match is Pox5InputMatch => match !== undefined
   );
 }
 
@@ -22,9 +24,9 @@ export function assertPox5SpendPolicy(psbt: Psbt, descriptorMap: descriptorWalle
     throw new Error(`expected exactly one PoX-5 lockup input, found ${matches.length}`);
   }
   const [match] = matches;
-  const branch = pox5.classifyPox5Spend(psbt, match);
+  const branch = classifyPox5Spend(psbt, match);
   if (branch !== 'locktime') {
     throw new Error(`unsupported PoX-5 spend branch: ${branch}`);
   }
-  pox5.assertPox5LocktimeSpend(psbt, matches);
+  assertPox5LocktimeSpend(psbt, matches);
 }

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

@@ -0,0 +1,93 @@
+import assert from 'assert/strict';
+import { createHash } from 'crypto';
+
+import { Descriptor, descriptorWallet, Psbt } from '@bitgo/wasm-utxo';
+import { getKey, getKeyTriple } from '@bitgo/wasm-utxo/testutils';
+
+import {
+  createPox5LockupDescriptor,
+  descriptorExtension,
+  isWalletPox5Lockup,
+  assertPox5SpendPolicy,
+} from '../../../src/pox5';
+import type { WalletKeyTriple } from '../../../src/extension';
+
+const UNLOCK_HEIGHT = 840_000;
+
+function lockupFor(keys: WalletKeyTriple): Descriptor {
+  return createPox5LockupDescriptor({
+    unlockHeight: UNLOCK_HEIGHT,
+    stakerCommitment: createHash('sha256').update(Buffer.alloc(32, 0x42)).digest(),
+    earlyExitKey: Buffer.from(getKey('pox5-extension-early-exit').publicKey),
+    stakerKeys: [keys[0].neutered(), keys[1].neutered(), keys[2].neutered()],
+  });
+}
+
+function toWalletKeys(label: string): WalletKeyTriple {
+  const [a, b, c] = getKeyTriple(label);
+  return [a, b, c];
+}
+
+function createSpend(
+  lockup: Descriptor,
+  locktime: number
+): { psbt: Psbt; descriptorMap: descriptorWallet.DescriptorMap } {
+  const concrete = lockup.atDerivationIndex(0);
+  const psbt = Psbt.create(2, locktime);
+  psbt.addInput('01'.repeat(32), 0, 100_000n, concrete.scriptPubkey(), 0xfffffffe);
+  psbt.addOutput(concrete.scriptPubkey(), 90_000n);
+  psbt.updateInputWithDescriptor(0, concrete);
+  return { psbt, descriptorMap: new Map([['pox5', lockup]]) };
+}
+
+describe('PoX-5 descriptor extension', function () {
+  const walletKeys = toWalletKeys('pox5-extension');
+
+  it('is named pox5', function () {
+    assert.equal(descriptorExtension.name, 'pox5');
+  });
+
+  describe('isWalletDescriptor', function () {
+    it('accepts the canonical lockup descriptor over the wallet keys', function () {
+      assert.ok(isWalletPox5Lockup(lockupFor(walletKeys), walletKeys));
+      assert.ok(descriptorExtension.isWalletDescriptor(lockupFor(walletKeys), walletKeys));
+    });
+
+    it('rejects a lockup descriptor over other keys', function () {
+      assert.equal(isWalletPox5Lockup(lockupFor(toWalletKeys('pox5-extension-foreign')), walletKeys), false);
+    });
+
+    it('rejects a lockup descriptor with permuted wallet keys', function () {
+      assert.equal(isWalletPox5Lockup(lockupFor([walletKeys[1], walletKeys[0], walletKeys[2]]), walletKeys), false);
+    });
+
+    it('rejects descriptors that are not lockup descriptors', function () {
+      const xpubs = walletKeys.map((k) => k.neutered().toBase58() + '/*');
+      const other = Descriptor.fromString(`wsh(multi(2,${xpubs.join(',')}))`, 'derivable');
+      assert.equal(isWalletPox5Lockup(other, walletKeys), false);
+    });
+  });
+
+  describe('assertSpendPolicy', function () {
+    it('does nothing for a PSBT without PoX-5 lockup inputs', function () {
+      const xpubs = walletKeys.map((k) => k.neutered().toBase58() + '/*');
+      const other = Descriptor.fromString(`wsh(multi(2,${xpubs.join(',')}))`, 'derivable');
+      const { psbt, descriptorMap } = createSpend(other, 0);
+      assert.doesNotThrow(() => descriptorExtension.assertSpendPolicy(psbt, descriptorMap));
+    });
+
+    it('accepts a late exit at the unlock height', function () {
+      const { psbt, descriptorMap } = createSpend(lockupFor(walletKeys), UNLOCK_HEIGHT);
+      assert.doesNotThrow(() => descriptorExtension.assertSpendPolicy(psbt, descriptorMap));
+    });
+
+    it('refuses a late exit below the unlock height', function () {
+      const { psbt, descriptorMap } = createSpend(lockupFor(walletKeys), UNLOCK_HEIGHT - 1);
+      assert.throws(() => descriptorExtension.assertSpendPolicy(psbt, descriptorMap), /nLockTime must be at least/);
+    });
+
+    it('delegates to the PoX-5 spend policy', function () {
+      assert.equal(descriptorExtension.assertSpendPolicy, assertPox5SpendPolicy);
+    });
+  });
+});