Merge pull request #237 from BitGo/follow-up/polling-worker-WCN-746

chore: update async worker to use callbacks

Author: pranishnepal

src/__tests__/api/master/asyncJobWorker.test.ts

@@ -21,7 +21,7 @@ const POLL_INTERVAL_MS = 1000;
 function makeUserKeychain() {
   return {
     id: 'user-key-id',
-    pub: 'xpub_user',
+    pub: 'xpub661MyMwAqRbcEvJQx6spkkHLRgtjxmVdyDSvbDt2m9NFpbkHdcu5WJsHHHqFxNATbNHnhMWJiwckoMqF75EpcNhU9xeVM4oDS7urM3os4BH',
     encryptedPrv: 'encrypted-user-prv',
     type: 'independent' as const,
     source: 'user' as const,
@@ -32,7 +32,7 @@ function makeUserKeychain() {
 function makeBackupKeychain() {
   return {
     id: 'backup-key-id',
-    pub: 'xpub_backup',
+    pub: 'xpub661MyMwAqRbcFnihegj1Mo2ePZoMQyLbBYpW7gDXZ7qzqxF3FBAkNAP8Gki8Mxx2BVLjN3RRa75pt5apD2g3ewXPrCfdssAJ7VupXqucLsb',
     encryptedPrv: 'encrypted-backup-prv',
     type: 'independent' as const,
     source: 'backup' as const,
@@ -215,19 +215,27 @@ describe('asyncJobWorker', () => {
         .query({ status: 'awaiting_bitgo' })
         .reply(200, { jobs: [job1, job2] });
 
-      nockBitgoKeychainRegistration({ pub: 'xpub_user', source: 'user', keyId: 'user-key-id' });
       nockBitgoKeychainRegistration({
-        pub: 'xpub_backup',
+        pub: makeUserKeychain().pub,
+        source: 'user',
+        keyId: 'user-key-id',
+      });
+      nockBitgoKeychainRegistration({
+        pub: makeBackupKeychain().pub,
         source: 'backup',
         keyId: 'backup-key-id',
       });
       nockBitgoKeyCreate('bitgo-key-id');
       nockWalletAdd('wallet-1');
       nockUpdateJobComplete('job-1', 'wallet-1');
 
-      nockBitgoKeychainRegistration({ pub: 'xpub_user', source: 'user', keyId: 'user-key-id' });
       nockBitgoKeychainRegistration({
-        pub: 'xpub_backup',
+        pub: makeUserKeychain().pub,
+        source: 'user',
+        keyId: 'user-key-id',
+      });
+      nockBitgoKeychainRegistration({
+        pub: makeBackupKeychain().pub,
         source: 'backup',
         keyId: 'backup-key-id',
       });
@@ -254,9 +262,13 @@ describe('asyncJobWorker', () => {
 
       nockUpdateJobFailed('job-bad');
 
-      nockBitgoKeychainRegistration({ pub: 'xpub_user', source: 'user', keyId: 'user-key-id' });
       nockBitgoKeychainRegistration({
-        pub: 'xpub_backup',
+        pub: makeUserKeychain().pub,
+        source: 'user',
+        keyId: 'user-key-id',
+      });
+      nockBitgoKeychainRegistration({
+        pub: makeBackupKeychain().pub,
         source: 'backup',
         keyId: 'backup-key-id',
       });
@@ -289,12 +301,12 @@ describe('asyncJobWorker', () => {
       const walletId = 'new-wallet-id';
 
       const userKeyNock = nockBitgoKeychainRegistration({
-        pub: 'xpub_user',
+        pub: makeUserKeychain().pub,
         source: 'user',
         keyId: 'user-key-id',
       });
       const backupKeyNock = nockBitgoKeychainRegistration({
-        pub: 'xpub_backup',
+        pub: makeBackupKeychain().pub,
         source: 'backup',
         keyId: 'backup-key-id',
       });
@@ -367,9 +379,13 @@ describe('asyncJobWorker', () => {
       });
       const walletId = 'ent-wallet-id';
 
-      nockBitgoKeychainRegistration({ pub: 'xpub_user', source: 'user', keyId: 'user-key-id' });
       nockBitgoKeychainRegistration({
-        pub: 'xpub_backup',
+        pub: makeUserKeychain().pub,
+        source: 'user',
+        keyId: 'user-key-id',
+      });
+      nockBitgoKeychainRegistration({
+        pub: makeBackupKeychain().pub,
         source: 'backup',
         keyId: 'backup-key-id',
       });

src/__tests__/integration/asyncJobWorker.integ.test.ts

@@ -40,7 +40,7 @@ function makeAwaitingBitgoJob(overrides: Partial<BridgeJobResponse> = {}): Bridg
       status: 200,
       body: {
         id: 'user-key-id',
-        pub: 'xpub_user',
+        pub: 'xpub661MyMwAqRbcEvJQx6spkkHLRgtjxmVdyDSvbDt2m9NFpbkHdcu5WJsHHHqFxNATbNHnhMWJiwckoMqF75EpcNhU9xeVM4oDS7urM3os4BH',
         type: 'independent',
         source: 'user',
         coin: COIN,
@@ -50,7 +50,7 @@ function makeAwaitingBitgoJob(overrides: Partial<BridgeJobResponse> = {}): Bridg
       status: 200,
       body: {
         id: 'backup-key-id',
-        pub: 'xpub_backup',
+        pub: 'xpub661MyMwAqRbcFnihegj1Mo2ePZoMQyLbBYpW7gDXZ7qzqxF3FBAkNAP8Gki8Mxx2BVLjN3RRa75pt5apD2g3ewXPrCfdssAJ7VupXqucLsb',
         type: 'independent',
         source: 'backup',
         coin: COIN,
@@ -80,7 +80,7 @@ describe('asyncJobWorker: end-to-end polling', () => {
     services.bridge.calls.length = 0;
   });
 
-  it('picks up an awaiting_bitgo keygen job, registers keychains, creates wallet, and PATCHes complete', async () => {
+  it('picks up an awaiting_bitgo keygen job, creates wallet, and PATCHes complete', async () => {
     const jobId = 'integ-job-123';
     assert(services.bridge, 'bridge service should be defined');
     services.bridge.setPendingJobs([makeAwaitingBitgoJob()]);

src/masterBitgoExpress/handlers/utils/walletCreationUtils.ts

@@ -1,85 +1,3 @@
-import {
-  Keychain,
-  KeychainsTriplet,
-  promiseProps,
-  RequestTracer,
-  SupplementGenerateWalletOptions,
-  Wallet,
-  WalletWithKeychains,
-} from '@bitgo-beta/sdk-core';
-import { BitGoAPI } from '@bitgo-beta/sdk-api';
-import _ from 'lodash';
-import { IndependentKeychainResponse } from '../../clients/advancedWalletManagerClient';
-import coinFactory from '../../../shared/coinFactory';
-
 export function getBaseWalletParams(multisigType: 'onchain' | 'tss') {
   return { m: 2, n: 3, keys: [] as string[], type: 'advanced', multisigType } as const;
 }
-
-export interface RegisterKeychainsAndCreateWalletParams {
-  coin: string;
-  bitgo: BitGoAPI;
-  userKeychain: IndependentKeychainResponse;
-  backupKeychain: IndependentKeychainResponse;
-  walletParams: SupplementGenerateWalletOptions;
-  isDistributedCustody?: boolean;
-}
-
-export async function registerKeychainsAndCreateWallet({
-  bitgo,
-  coin,
-  walletParams,
-  userKeychain,
-  backupKeychain,
-  isDistributedCustody,
-}: RegisterKeychainsAndCreateWalletParams): Promise<WalletWithKeychains> {
-  const baseCoin = await coinFactory.getCoin(coin, bitgo);
-  const reqId = new RequestTracer();
-
-  const registerKeychain = async (keyChain: IndependentKeychainResponse): Promise<Keychain> => {
-    const registered = await baseCoin.keychains().add({
-      pub: keyChain.pub,
-      keyType: keyChain.type,
-      source: keyChain.source,
-      reqId,
-    });
-    return _.extend({}, registered, keyChain);
-  };
-
-  const {
-    userKeychain: registeredUser,
-    backupKeychain: registeredBackup,
-    bitgoKeychain,
-  }: KeychainsTriplet = await promiseProps({
-    userKeychain: registerKeychain(userKeychain),
-    backupKeychain: registerKeychain(backupKeychain),
-    bitgoKeychain: baseCoin.keychains().createBitGo({
-      enterprise: walletParams.enterprise,
-      keyType: 'independent',
-      reqId,
-      isDistributedCustody,
-    }),
-  });
-
-  const keychains: KeychainsTriplet = {
-    userKeychain: registeredUser,
-    backupKeychain: registeredBackup,
-    bitgoKeychain,
-  };
-
-  const finalWalletParams = await baseCoin.supplementGenerateWallet(
-    { ...walletParams, keys: [registeredUser.id, registeredBackup.id, bitgoKeychain.id] },
-    keychains,
-  );
-
-  bitgo.setRequestTracer(reqId);
-  const newWallet = await bitgo.post(baseCoin.url('/wallet/add')).send(finalWalletParams).result();
-
-  return {
-    wallet: new Wallet(bitgo, baseCoin, newWallet),
-    userKeychain: registeredUser,
-    backupKeychain: registeredBackup,
-    bitgoKeychain,
-    responseType: 'WalletWithKeychains',
-  };
-}

src/masterBitgoExpress/handlers/walletGenerationCallbacks.ts

@@ -1,6 +1,9 @@
 import { CreateKeychainCallback } from '@bitgo-beta/sdk-core';
 import { KeySource } from '../../shared/types';
-import { AdvancedWalletManagerClient } from '../clients/advancedWalletManagerClient';
+import {
+  AdvancedWalletManagerClient,
+  IndependentKeychainResponse,
+} from '../clients/advancedWalletManagerClient';
 
 export function createOnchainKeyGenCallback(
   awmUserClient: AdvancedWalletManagerClient,
@@ -20,3 +23,20 @@ export function createOnchainKeyGenCallback(
     return keychain as { pub: string; type: 'independent'; source: typeof source };
   };
 }
+
+export function createOnchainKeyGenCallbackForPreGeneratedKeychains(
+  preGeneratedKeychains: Record<KeySource.USER | KeySource.BACKUP, IndependentKeychainResponse>,
+): CreateKeychainCallback {
+  return async ({ source, coin: _ }) => {
+    if (!(source in preGeneratedKeychains)) {
+      throw new Error(`${source} keychain not available for onchain key generation`);
+    }
+
+    const keychain = preGeneratedKeychains[source];
+    return {
+      source,
+      pub: keychain.pub,
+      type: 'independent',
+    };
+  };
+}

src/masterBitgoExpress/workers/asyncJobWorker.ts

@@ -1,17 +1,14 @@
-import { CreateBitGoOptions, SupplementGenerateWalletOptions } from '@bitgo-beta/sdk-core';
 import { BitGoAPI } from '@bitgo-beta/sdk-api';
 import { OsoBridgeClient } from '../clients/bridgeClient';
 import { AwmResponseSchema, BridgeJobResponse } from '../clients/bridgeClient.types';
 import {
   IndependentKeychainResponseSchema,
   type IndependentKeychainResponse,
 } from '../clients/advancedWalletManagerClient';
-import {
-  getBaseWalletParams,
-  registerKeychainsAndCreateWallet,
-} from '../handlers/utils/walletCreationUtils';
+import coinFactory from '../../shared/coinFactory';
 import { MasterExpressConfig } from '../../shared/types';
 import logger from '../../shared/logger';
+import { createOnchainKeyGenCallbackForPreGeneratedKeychains } from '../handlers/walletGenerationCallbacks';
 
 const ASYNC_OPERATIONS_TO_HANDLERS: Partial<
   Record<
@@ -116,31 +113,26 @@ export async function handleKeyGenerationOperation(
   const backupKeychain = parseKeychainFromAwmResponse(job.awmBackupResponse, 'awmBackupResponse');
   const { jobId, coin, version } = job;
 
-  const requestBody = (job.request?.body ?? {}) as unknown as SupplementGenerateWalletOptions &
-    Pick<CreateBitGoOptions, 'isDistributedCustody'>;
-
-  const walletParams: SupplementGenerateWalletOptions = {
-    ...requestBody,
-    ...getBaseWalletParams('onchain'),
-  };
-
-  const result = await registerKeychainsAndCreateWallet({
-    bitgo,
-    walletParams,
-    userKeychain,
-    backupKeychain,
-    coin,
-    isDistributedCustody: requestBody.isDistributedCustody,
+  const baseCoin = await coinFactory.getCoin(coin, bitgo);
+  const result = await baseCoin.wallets().generateWallet({
+    ...(job.request?.body ?? {}),
+    type: 'advanced',
+    multisigType: 'onchain',
+    createKeychainCallback: createOnchainKeyGenCallbackForPreGeneratedKeychains({
+      user: userKeychain,
+      backup: backupKeychain,
+    }),
   });
 
   logger.info(`${logPrefix} job ${jobId} created wallet - updating job status to complete`);
+  const walletId = result.wallet.toJSON().id;
 
   await bridge.updateJob({
     jobId,
     version,
     status: 'complete',
-    result: { walletId: result.wallet.id() },
+    result: { walletId },
   });
 
-  logger.info(`${logPrefix} job ${jobId} complete, walletId ${result.wallet.id()}`);
+  logger.info(`${logPrefix} job ${jobId} complete, walletId ${walletId}`);
 }