fix(abstract-utxo): add PoX-5 lockup validator to descriptor policy

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

@@ -1,5 +1,6 @@
 import { EnvironmentName, Triple } from '@bitgo/sdk-core';
-import { bip32, descriptorWallet } from '@bitgo/wasm-utxo';
+import { pox5 } from '@bitgo/utxo-descriptors';
+import { BIP32, bip32, descriptorWallet } from '@bitgo/wasm-utxo';
 
 import { parseDescriptor } from './builder';
 import { hasValidSignature, NamedDescriptor, NamedDescriptorNative, toNamedDescriptorNative } from './NamedDescriptor';
@@ -19,18 +20,89 @@ export const policyAllowAll: DescriptorValidationPolicy = {
   validate: () => true,
 };
 
+function tryParseDescriptor(descriptor: NamedDescriptorNative['value']) {
+  try {
+    return parseDescriptor(descriptor);
+  } catch {
+    // probing a template is fallible: a descriptor that cannot be parsed is
+    // simply not a match so that `some(...)` falls through to the next validator
+    return undefined;
+  }
+}
+
+function matchesDescriptorTemplate(
+  descriptor: NamedDescriptorNative['value'],
+  name: string,
+  walletKeys: KeyTriple
+): boolean {
+  const parsed = tryParseDescriptor(descriptor);
+  return (
+    parsed !== undefined &&
+    parsed.name === name &&
+    parsed.keys.length === walletKeys.length &&
+    parsed.keys.every((k, i) => k.toBase58() === walletKeys[i].neutered().toBase58())
+  );
+}
+
 export function getValidatorDescriptorTemplate(name: string): DescriptorValidationPolicy {
   return {
     name: 'descriptorTemplate(' + name + ')',
     validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean {
-      return arr.every((d) => {
-        const parsed = parseDescriptor(d.value);
-        return (
-          parsed.name === name &&
-          parsed.keys.length === walletKeys.length &&
-          parsed.keys.every((k, i) => k.toBase58() === walletKeys[i].neutered().toBase58())
-        );
-      });
+      return arr.every((d) => matchesDescriptorTemplate(d.value, name, walletKeys));
+    },
+  };
+}
+
+function tryParsePox5Lockup(descriptor: NamedDescriptorNative['value']): pox5.Pox5DescriptorInfo | undefined {
+  try {
+    return pox5.parsePox5LockupDescriptor(descriptor);
+  } catch {
+    // malformed lockups are not canonical PoX-5 descriptors
+    return undefined;
+  }
+}
+
+function hasPox5WalletStakerKeys(info: pox5.Pox5DescriptorInfo, walletKeys: KeyTriple): boolean {
+  return (
+    info.stakerKeyStrings.length === walletKeys.length &&
+    info.stakerKeyStrings.every((keyString, i) => {
+      // staker keys are `<xpub>/<path>` entries; the wallet key identity is the xpub root
+      const [xpub] = keyString.split('/');
+      try {
+        return BIP32.fromBase58(xpub).toBase58() === walletKeys[i].neutered().toBase58();
+      } catch {
+        // raw pubkey staker keys cannot be matched against the wallet key triple
+        return false;
+      }
+    })
+  );
+}
+
+/**
+ * Accepts PoX-5 staking descriptor groups: at least one canonical PoX-5 lockup
+ * descriptor whose staker keys match the wallet key triple, with every remaining
+ * descriptor matching one of `templateNames` (e.g. the wallet's external and
+ * internal descriptors).
+ */
+export function getValidatorPox5Lockup(templateNames: string[]): DescriptorValidationPolicy {
+  const templatesPolicy = getValidatorOneOfTemplates(templateNames);
+  return {
+    name: 'pox5Lockup(' + templateNames.join(',') + ')',
+    validate(arr: NamedDescriptorNative[], walletKeys: KeyTriple): boolean {
+      const rest: NamedDescriptorNative[] = [];
+      let lockupCount = 0;
+      for (const d of arr) {
+        const info = tryParsePox5Lockup(d.value);
+        if (info === undefined) {
+          rest.push(d);
+          continue;
+        }
+        lockupCount++;
+        if (!hasPox5WalletStakerKeys(info, walletKeys)) {
+          return false;
+        }
+      }
+      return lockupCount > 0 && templatesPolicy.validate(rest, walletKeys);
     },
   };
 }
@@ -106,6 +178,9 @@ 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 PoX-5 staking groups: canonical lockup descriptors with wallet
+        // staker keys alongside the wallet's regular descriptors
+        getValidatorPox5Lockup(['Wsh2Of3', 'Wsh2Of3CltvDrop', 'SbtcDeposit']),
         // allow all descriptors signed by the user key
         getValidatorSignedByUserKey(),
       ]);

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

@@ -1,7 +1,8 @@
 import assert from 'assert';
 
 import { Triple } from '@bitgo/sdk-core';
-import { bip32 } from '@bitgo/wasm-utxo';
+import { pox5 } from '@bitgo/utxo-descriptors';
+import { bip32, Descriptor } from '@bitgo/wasm-utxo';
 import * as testutils from '@bitgo/wasm-utxo/testutils';
 
 import {
@@ -10,6 +11,7 @@ import {
   DescriptorValidationPolicy,
   getPolicyForEnv,
   getValidatorDescriptorTemplate,
+  getValidatorPox5Lockup,
 } from '../../../../src/descriptor/validatePolicy';
 import {
   NamedDescriptor,
@@ -89,3 +91,94 @@ describe('assertDescriptorPolicy', function () {
     testAssertDescriptorPolicy([stripSignature(getNamedDescriptor('Wsh2Of2'))], getPolicyForEnv('test'), keys, null);
   });
 });
+
+describe('PoX-5 lockup descriptor policy', function () {
+  const keys = testutils.getKeyTriple('default');
+  const stakerCommitment = Buffer.alloc(32, 0x42);
+  const earlyExitKey = Buffer.from(keys[0].derive(9).publicKey);
+
+  function getLockupDescriptor(stakerKeys: Triple<bip32.BIP32Interface> = keys): Descriptor {
+    return pox5.createPox5LockupDescriptor({
+      unlockHeight: 840_000,
+      stakerCommitment,
+      earlyExitKey,
+      stakerKeys,
+    });
+  }
+
+  function toNamed(value: string, name = 'descriptor'): NamedDescriptor {
+    return { name, value };
+  }
+
+  const external = toNamed(getDescriptor('Wsh2Of3').toString(), 'external');
+  const internal = toNamed(getDescriptor('Wsh2Of3').toString(), 'internal');
+  const lockup = toNamed(getLockupDescriptor().toString(), 'pox5-lockup-u840000');
+  const policy = getValidatorPox5Lockup(['Wsh2Of3', 'Wsh2Of3CltvDrop', 'SbtcDeposit']);
+
+  function assertAccepts(ds: NamedDescriptor[]): void {
+    assert.doesNotThrow(() =>
+      assertDescriptorPolicy(
+        ds.map((d) => toNamedDescriptorNative(d, 'derivable')),
+        policy,
+        keys
+      )
+    );
+    assert.doesNotThrow(() =>
+      assertDescriptorPolicy(
+        ds.map((d) => toNamedDescriptorNative(d, 'derivable')),
+        getPolicyForEnv('prod'),
+        keys
+      )
+    );
+  }
+
+  function assertRejects(ds: NamedDescriptor[]): void {
+    assert.throws(() =>
+      assertDescriptorPolicy(
+        ds.map((d) => toNamedDescriptorNative(d, 'derivable')),
+        policy,
+        keys
+      )
+    );
+    assert.throws(() =>
+      assertDescriptorPolicy(
+        ds.map((d) => toNamedDescriptorNative(d, 'derivable')),
+        getPolicyForEnv('prod'),
+        keys
+      )
+    );
+  }
+
+  it('accepts wallet descriptors mixed with a canonical lockup', function () {
+    assertAccepts([lockup]);
+    assertAccepts([external, internal, lockup]);
+  });
+
+  it('requires at least one lockup', function () {
+    // the template policy already accepts plain wallet descriptor groups
+    assert.throws(() =>
+      assertDescriptorPolicy(
+        [external, internal].map((d) => toNamedDescriptorNative(d, 'derivable')),
+        policy,
+        keys
+      )
+    );
+  });
+
+  it('rejects lockups with foreign staker keys', function () {
+    assertRejects([external, internal, toNamed(getLockupDescriptor(testutils.getKeyTriple('foreign')).toString())]);
+  });
+
+  it('rejects malformed and non-canonical lockups', function () {
+    const malformed = getLockupDescriptor().toString().split('#')[0].replace('multi(2,', 'multi(3,');
+    assertRejects([external, internal, toNamed(malformed)]);
+    assertRejects([external, internal, toNamed(`wsh(pk(${keys[0].neutered().toBase58()}/*))`)]);
+    assertRejects([toNamed(malformed)]);
+  });
+
+  it('returns false instead of throwing for descriptors it cannot parse', function () {
+    const nativeLockup = toNamedDescriptorNative(lockup, 'derivable');
+    assert.strictEqual(getValidatorDescriptorTemplate('Wsh2Of3').validate([nativeLockup], keys), false);
+    assert.strictEqual(getValidatorDescriptorTemplate('Wsh2Of3CltvDrop').validate([nativeLockup], keys), false);
+  });
+});