Merge pull request #8156 from BitGo/BTC-2768.custom-change-wallet.mockdata

feat(abstract-utxo): add customChangeXpubs to explainTransaction

Author: OttoAllmendinger

modules/abstract-utxo/src/abstractUtxoCoin.ts

@@ -243,6 +243,7 @@ export interface ExplainTransactionOptions<TNumber extends number | bigint = num
   txInfo?: TransactionInfo<TNumber>;
   feeInfo?: string;
   pubs?: Triple<string>;
+  customChangeXpubs?: Triple<string>;
   decodeWith?: SdkBackend;
 }
 

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

@@ -27,6 +27,7 @@ export function explainTx<TNumber extends number | bigint>(
   params: {
     wallet?: IWallet;
     pubs?: string[];
+    customChangeXpubs?: Triple<string>;
     txInfo?: { unspents?: Unspent<TNumber>[] };
     changeInfo?: fixedScript.ChangeAddressInfo[];
   },
@@ -49,7 +50,7 @@ export function explainTx<TNumber extends number | bigint>(
     throw new Error('legacy transactions are not supported for descriptor wallets');
   }
   if (tx instanceof utxolib.bitgo.UtxoPsbt) {
-    return fixedScript.explainPsbt(tx, params, coinName);
+    return fixedScript.explainPsbt(tx, { ...params, customChangePubs: params.customChangeXpubs }, coinName);
   } else if (tx instanceof fixedScriptWallet.BitGoPsbt) {
     const pubs = params.pubs;
     if (!pubs) {
@@ -68,6 +69,7 @@ export function explainTx<TNumber extends number | bigint>(
       replayProtection: {
         publicKeys: getReplayProtectionPubkeys(coinName),
       },
+      customChangeWalletXpubs: params.customChangeXpubs,
     });
   } else {
     return fixedScript.explainLegacyTx(tx, params, coinName);

modules/abstract-utxo/src/transaction/fixedScript/parseTransaction.ts

@@ -1,11 +1,19 @@
 import assert from 'assert';
 
 import _ from 'lodash';
-import { ITransactionRecipient, Triple, VerificationOptions, Wallet } from '@bitgo/sdk-core';
+import { ITransactionRecipient, KeyIndices, Triple, VerificationOptions, Wallet } from '@bitgo/sdk-core';
 
 import type { AbstractUtxoCoin, ParseTransactionOptions } from '../../abstractUtxoCoin';
 import type { FixedScriptWalletOutput, Output, ParsedTransaction } from '../types';
-import { fetchKeychains, getKeySignatures, toKeychainTriple, UtxoKeychain, UtxoNamedKeychains } from '../../keychains';
+import {
+  fetchKeychains,
+  getKeySignatures,
+  toKeychainTriple,
+  toXpubTriple,
+  UtxoKeychain,
+  UtxoNamedKeychains,
+} from '../../keychains';
+import { verifyKeySignature } from '../../verifyKey';
 import {
   assertValidTransactionRecipient,
   fromExtendedAddressFormatToScript,
@@ -112,6 +120,20 @@ function toExpectedOutputs(
   return expectedOutputs;
 }
 
+function verifyCustomChangeKeys(userKeychain: UtxoKeychain, customChange: CustomChangeOptions): void {
+  for (const keyIndex of [KeyIndices.USER, KeyIndices.BACKUP, KeyIndices.BITGO]) {
+    if (
+      !verifyKeySignature({
+        userKeychain,
+        keychainToVerify: customChange.keys[keyIndex],
+        keySignature: customChange.signatures[keyIndex],
+      })
+    ) {
+      throw new Error(`failed to verify custom change ${KeyIndices[keyIndex].toLowerCase()} key signature`);
+    }
+  }
+}
+
 export async function parseTransaction<TNumber extends bigint | number>(
   coin: AbstractUtxoCoin,
   params: ParseTransactionOptions<TNumber>
@@ -177,11 +199,18 @@ export async function parseTransaction<TNumber extends bigint | number>(
     }
   }
 
+  let customChangeXpubs: Triple<string> | undefined;
+  if (customChange) {
+    verifyCustomChangeKeys(keychainArray[KeyIndices.USER], customChange);
+    customChangeXpubs = toXpubTriple(customChange.keys);
+  }
+
   // obtain all outputs
   const explanation: TransactionExplanation = await coin.explainTransaction<TNumber>({
     txHex: txPrebuild.txHex,
     txInfo: txPrebuild.txInfo,
     pubs: keychainArray.map((k) => k.pub) as Triple<string>,
+    customChangeXpubs,
   });
 
   const allOutputs = [...explanation.outputs, ...explanation.changeOutputs];

modules/abstract-utxo/test/unit/customChangeWallet.ts

@@ -1,3 +1,5 @@
+import assert from 'assert';
+
 import should = require('should');
 import * as sinon from 'sinon';
 import { Wallet } from '@bitgo/sdk-core';
@@ -133,4 +135,66 @@ describe('Custom Change Wallets', () => {
     (coin.wallets as any).restore();
     (coin.keychains as any).restore();
   });
+
+  it('should reject invalid custom change key signatures before calling explainTransaction', async () => {
+    const wrongKey = coin.keychains().create();
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    const sign = async ({ key }) => (await coin.signMessage({ prv: wrongKey.prv }, key.pub!)).toString('hex');
+    const invalidSignatures = {
+      user: await sign(keys.change.user),
+      backup: await sign(keys.change.backup),
+      bitgo: await sign(keys.change.bitgo),
+    };
+
+    const signedSendingWallet = sinon.createStubInstance(Wallet, stubData.signedSendingWallet as any);
+    const changeWallet = sinon.createStubInstance(Wallet, stubData.changeWallet as any);
+
+    sinon.stub(coin, 'keychains').returns({
+      get: sinon.stub().callsFake(({ id }) => {
+        switch (id) {
+          case keys.send.user.id:
+            return Promise.resolve({ id, ...keys.send.user.key });
+          case keys.send.backup.id:
+            return Promise.resolve({ id, ...keys.send.backup.key });
+          case keys.send.bitgo.id:
+            return Promise.resolve({ id, ...keys.send.bitgo.key });
+          case keys.change.user.id:
+            return Promise.resolve({ id, ...keys.change.user.key });
+          case keys.change.backup.id:
+            return Promise.resolve({ id, ...keys.change.backup.key });
+          case keys.change.bitgo.id:
+            return Promise.resolve({ id, ...keys.change.bitgo.key });
+        }
+      }),
+    } as any);
+
+    sinon.stub(coin, 'wallets').returns({
+      get: sinon.stub().callsFake(() => Promise.resolve(changeWallet)),
+    } as any);
+
+    const explainStub = sinon.stub(coin, 'explainTransaction');
+
+    signedSendingWallet._wallet = signedSendingWallet._wallet || {
+      customChangeKeySignatures: invalidSignatures,
+    };
+
+    try {
+      await coin.parseTransaction({
+        txParams: { recipients: [] },
+        txPrebuild: { txHex: '' },
+        wallet: signedSendingWallet as any,
+        verification: {},
+      });
+      assert.fail('parseTransaction should have thrown for invalid custom change key signatures');
+    } catch (e) {
+      assert.ok(e instanceof Error);
+      assert.match(e.message, /failed to verify custom change .* key signature/);
+    }
+
+    assert.strictEqual(explainStub.called, false, 'explainTransaction should not have been called');
+
+    explainStub.restore();
+    (coin.wallets as any).restore();
+    (coin.keychains as any).restore();
+  });
 });

modules/abstract-utxo/test/unit/customChangeWalletFixtures.ts

@@ -0,0 +1,68 @@
+import assert from 'assert';
+
+import { BIP32, message } from '@bitgo/wasm-utxo';
+
+import { verifyKeySignature } from '../../src/verifyKey';
+
+const KEY_ROLES = ['user', 'backup', 'bitgo'] as const;
+type KeyRole = (typeof KEY_ROLES)[number];
+
+function signKeySignature(signerPrivateKey: BIP32, pubToSign: string): string {
+  assert(signerPrivateKey.privateKey, 'signer must have a private key');
+  return Buffer.from(message.signMessage(pubToSign, signerPrivateKey.privateKey)).toString('hex');
+}
+
+function createMockCustomChangeKeySignatures(
+  signerPrivateKey: BIP32,
+  changeKeychains: { pub: string }[]
+): Record<KeyRole, string> {
+  return {
+    user: signKeySignature(signerPrivateKey, changeKeychains[0].pub),
+    backup: signKeySignature(signerPrivateKey, changeKeychains[1].pub),
+    bitgo: signKeySignature(signerPrivateKey, changeKeychains[2].pub),
+  };
+}
+
+describe('Custom Change Wallet Key Signatures', function () {
+  const mainUserKey = BIP32.fromSeedSha256('main-user');
+  const changeKeys = [
+    BIP32.fromSeedSha256('change-user'),
+    BIP32.fromSeedSha256('change-backup'),
+    BIP32.fromSeedSha256('change-bitgo'),
+  ];
+  const changeKeychains = changeKeys.map((k) => ({ pub: k.neutered().toBase58() }));
+
+  it('should accept signatures from the correct key', function () {
+    const signatures = createMockCustomChangeKeySignatures(mainUserKey, changeKeychains);
+
+    for (const role of KEY_ROLES) {
+      const index = KEY_ROLES.indexOf(role);
+      assert.ok(
+        verifyKeySignature({
+          userKeychain: { pub: mainUserKey.neutered().toBase58() },
+          keychainToVerify: changeKeychains[index],
+          keySignature: signatures[role],
+        }),
+        `failed to verify mock custom change ${role} key signature`
+      );
+    }
+  });
+
+  it('should reject signatures from a different key', function () {
+    const wrongKey = BIP32.fromSeedSha256('wrong-key');
+    const badSignatures = createMockCustomChangeKeySignatures(wrongKey, changeKeychains);
+
+    for (const role of KEY_ROLES) {
+      const index = KEY_ROLES.indexOf(role);
+      assert.strictEqual(
+        verifyKeySignature({
+          userKeychain: { pub: mainUserKey.neutered().toBase58() },
+          keychainToVerify: changeKeychains[index],
+          keySignature: badSignatures[role],
+        }),
+        false,
+        `should have rejected custom change ${role} key signature from wrong key`
+      );
+    }
+  });
+});