fix(abstract-utxo): check the descriptor policy spend rules when a descriptor wallet signs

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

@@ -1,9 +1,10 @@
 import { EnvironmentName, Triple } from '@bitgo/sdk-core';
 import { pox5 } from '@bitgo/utxo-descriptors';
-import { bip32, Descriptor, descriptorWallet } from '@bitgo/wasm-utxo';
+import { bip32, Descriptor, descriptorWallet, Psbt } from '@bitgo/wasm-utxo';
 
 import { DescriptorBuilder, parseDescriptor } from './builder';
 import { hasValidSignature, NamedDescriptor, NamedDescriptorNative, toNamedDescriptorNative } from './NamedDescriptor';
+import { assertPox5SpendPolicy } from './pox5SpendPolicy';
 
 import type { DescriptorMap } from './index';
 
@@ -13,6 +14,19 @@ export interface DescriptorValidationPolicy {
   name: string;
 
   validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean;
+
+  /**
+   * Throws if the PSBT spends an input of a descriptor this policy covers in a way that must not be signed.
+   * Must not throw for inputs the policy does not recognize.
+   */
+  validateSpend?(psbt: Psbt, descriptorMap: descriptorWallet.DescriptorMap): void;
+}
+
+/** Spend validation of a policy that is made up of other policies: each of them checks the inputs it recognizes. */
+function validateSpendOf(
+  validators: DescriptorValidationPolicy[]
+): (psbt: Psbt, descriptorMap: descriptorWallet.DescriptorMap) => void {
+  return (psbt, descriptorMap) => validators.forEach((v) => v.validateSpend?.(psbt, descriptorMap));
 }
 
 export const policyAllowAll: DescriptorValidationPolicy = {
@@ -52,6 +66,7 @@ export function getValidatorEvery(validators: DescriptorValidationPolicy[]): Des
     validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean {
       return validators.every((v) => v.validate(arr, walletKeys));
     },
+    validateSpend: validateSpendOf(validators),
   };
 }
 
@@ -61,6 +76,7 @@ export function getValidatorSome(validators: DescriptorValidationPolicy[]): Desc
     validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean {
       return validators.some((v) => v.validate(arr, walletKeys));
     },
+    validateSpend: validateSpendOf(validators),
   };
 }
 
@@ -110,6 +126,7 @@ export function getValidatorPox5Lockup(): DescriptorValidationPolicy {
     validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean {
       return arr.every((d) => isCanonicalPox5Lockup(d.value, walletKeys));
     },
+    validateSpend: assertPox5SpendPolicy,
   };
 }
 
@@ -123,6 +140,7 @@ export function getValidatorEachOneOf(validators: DescriptorValidationPolicy[]):
     validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean {
       return arr.every((d) => validators.some((v) => v.validate([d], walletKeys)));
     },
+    validateSpend: validateSpendOf(validators),
   };
 }
 
@@ -163,22 +181,24 @@ export function toDescriptorMapValidate(
 }
 
 export function getPolicyForEnv(env: EnvironmentName): DescriptorValidationPolicy {
+  const prodPolicy = getValidatorSome([
+    // allow 2-of-3-ish descriptor groups where the keys match the wallet keys
+    getValidatorDescriptorTemplate('Wsh2Of3'),
+    // allow descriptor groups where all keys match the wallet keys plus OP_DROP (coredao staking)
+    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 all descriptors signed by the user key
+    getValidatorSignedByUserKey(),
+  ]);
   switch (env) {
     case 'adminProd':
     case 'prod':
-      return getValidatorSome([
-        // allow 2-of-3-ish descriptor groups where the keys match the wallet keys
-        getValidatorDescriptorTemplate('Wsh2Of3'),
-        // allow descriptor groups where all keys match the wallet keys plus OP_DROP (coredao staking)
-        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 all descriptors signed by the user key
-        getValidatorSignedByUserKey(),
-      ]);
+      return prodPolicy;
     default:
-      return policyAllowAll;
+      // any descriptor is accepted outside of prod, but spends are still checked
+      return { ...policyAllowAll, validateSpend: prodPolicy.validateSpend };
   }
 }

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

@@ -13,8 +13,7 @@ import {
 import { explainPsbt, signPsbt } from '../../transaction/descriptor';
 import { TransactionExplanation } from '../TransactionExplanation';
 import { UtxoCoinName } from '../../names';
-
-import { assertPox5SpendPolicy } from './pox5';
+import { assertPox5SpendPolicy } from '../../descriptor/pox5SpendPolicy';
 
 export const DescriptorTransaction = t.intersection(
   [OfflineVaultSignable, t.type({ descriptors: t.array(NamedDescriptor) })],

modules/abstract-utxo/src/transaction/signTransaction.ts

@@ -68,7 +68,9 @@ export async function signTransaction<TNumber extends number | bigint>(
     }
     const psbt = decodeDescriptorPsbt(params.txPrebuild);
     const walletKeys = toBip32Triple(await fetchKeychains(coin, wallet));
-    const descriptorMap = getDescriptorMapFromWallet(wallet, walletKeys, getPolicyForEnv(bitgo.env));
+    const policy = getPolicyForEnv(bitgo.env);
+    const descriptorMap = getDescriptorMapFromWallet(wallet, walletKeys, policy);
+    policy.validateSpend?.(psbt, descriptorMap);
     descriptor.signPsbt(psbt, descriptorMap, signerKeychain, {
       onUnknownInput: 'throw',
     });

modules/abstract-utxo/test/unit/transaction/descriptor/signTransactionPox5.ts

@@ -0,0 +1,125 @@
+import 'mocha';
+import assert from 'assert';
+import { createHash } from 'crypto';
+
+import { pox5 } from '@bitgo/utxo-descriptors';
+import { Descriptor, descriptorWallet, Psbt } from '@bitgo/wasm-utxo';
+import { getKey, getKeyTriple } from '@bitgo/wasm-utxo/testutils';
+
+import { getPolicyForEnv } from '../../../../src/descriptor/validatePolicy';
+import type { UtxoWallet } from '../../../../src/wallet';
+import { getUtxoCoin } from '../../util/utxoCoins';
+import { nockBitGo } from '../../util/nockBitGo';
+
+const UNLOCK_HEIGHT = 840_000;
+const PRINCIPAL_PREIMAGE = Buffer.alloc(32, 0x42);
+
+type Pox5SpendParams = {
+  locktime: number;
+  sequence?: number;
+  /** adds the principal preimage to the lockup input, which makes the spend an early exit */
+  preimage?: Uint8Array;
+};
+
+// Hot descriptor wallets sign through coin.signTransaction. The PoX-5 spend policy that applies to the
+// offline vault must apply here as well: the SDK must not sign an exit that spends the lockup incorrectly.
+describe('signTransaction: PoX-5 lockup spend of a descriptor wallet', function () {
+  const coin = getUtxoCoin('btc');
+  const keys = getKeyTriple('pox5-hot-wallet');
+  const keyIds = ['kU', 'kB', 'kG'];
+
+  const lockup: Descriptor = pox5.createPox5LockupDescriptor({
+    unlockHeight: UNLOCK_HEIGHT,
+    stakerCommitment: createHash('sha256').update(PRINCIPAL_PREIMAGE).digest(),
+    earlyExitKey: Buffer.from(getKey('pox5-hot-wallet-early-exit').publicKey),
+    stakerKeys: [keys[0].neutered(), keys[1].neutered(), keys[2].neutered()],
+  });
+
+  function createPsbt({ locktime, sequence = 0xfffffffe, preimage }: Pox5SpendParams): Psbt {
+    const lockupAtIndex = lockup.atDerivationIndex(0);
+    const psbt = descriptorWallet.createPsbt(
+      { version: 2, locktime },
+      [
+        {
+          hash: '01'.repeat(32),
+          index: 0,
+          witnessUtxo: { script: lockupAtIndex.scriptPubkey(), value: 100_000n },
+          descriptor: lockupAtIndex,
+          sequence,
+        },
+      ],
+      [{ script: lockupAtIndex.scriptPubkey(), value: 90_000n }]
+    );
+    if (preimage) {
+      psbt.addSha256Preimage(0, preimage);
+    }
+    return psbt;
+  }
+
+  async function sign(params: Pox5SpendParams): Promise<Psbt> {
+    const wallet = {
+      coinSpecific: () => ({ descriptors: [{ name: 'pox5-lockup', value: lockup.toString() }] }),
+      keyIds: () => keyIds,
+    } as unknown as UtxoWallet;
+    keyIds.forEach((id, i) => {
+      nockBitGo().get(`/api/v2/${coin.getChain()}/key/${id}`).reply(200, { pub: keys[i].neutered().toBase58() });
+    });
+    const result = await coin.signTransaction({
+      txPrebuild: { txHex: Buffer.from(createPsbt(params).serialize()).toString('hex'), decodeWith: 'wasm-utxo' },
+      prv: keys[0].toBase58(),
+      wallet,
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    } as any);
+    assert.ok('txHex' in result);
+    return Psbt.deserialize(Buffer.from(result.txHex, 'hex'));
+  }
+
+  function assertSigned(psbt: Psbt): void {
+    const sigs = psbt.getPartialSignatures(0);
+    assert.strictEqual(sigs.length, 1);
+    assert.ok(psbt.validateSignatureAtInput(0, sigs[0].pubkey));
+  }
+
+  describe('late exit', function () {
+    it('signs with nLockTime at the unlock height', async function () {
+      assertSigned(await sign({ locktime: UNLOCK_HEIGHT }));
+    });
+
+    it('rejects nLockTime below the unlock height', async function () {
+      await assert.rejects(sign({ locktime: UNLOCK_HEIGHT - 1 }), /nLockTime must be at least/);
+    });
+
+    it('rejects a final sequence', async function () {
+      await assert.rejects(sign({ locktime: UNLOCK_HEIGHT, sequence: 0xffffffff }), /non-final sequence/);
+    });
+  });
+
+  describe('early exit', function () {
+    it('signs with a zero nLockTime and keeps the principal preimage', async function () {
+      const signed = await sign({ locktime: 0, preimage: PRINCIPAL_PREIMAGE });
+      assertSigned(signed);
+      assert.ok(signed.getInputKeyValues(0).some((r) => r.type === 'known' && r.key === 'PSBT_IN_SHA256'));
+    });
+
+    it('rejects a non-zero nLockTime', async function () {
+      await assert.rejects(sign({ locktime: UNLOCK_HEIGHT, preimage: PRINCIPAL_PREIMAGE }), /nLockTime must be 0/);
+    });
+
+    it('rejects a principal preimage that does not match the descriptor commitment', async function () {
+      await assert.rejects(
+        sign({ locktime: 0, preimage: Buffer.alloc(32, 0x43) }),
+        /expected exactly one PSBT_IN_SHA256 record/
+      );
+    });
+  });
+
+  describe('policy', function () {
+    for (const env of ['prod', 'adminProd', 'test', 'mock'] as const) {
+      it(`checks the spend in the ${env} policy`, function () {
+        const spend = createPsbt({ locktime: UNLOCK_HEIGHT - 1 });
+        const descriptorMap = new Map([['pox5-lockup', lockup]]);
+        assert.throws(() => getPolicyForEnv(env).validateSpend?.(spend, descriptorMap), /nLockTime must be at least/);
+      });
+    }
+  });
+});