From 2b0166c14719060148754478efefedc1c2c05001 Mon Sep 17 00:00:00 2001 From: Pranish Nepal Date: Fri, 12 Jun 2026 14:55:19 -0400 Subject: [PATCH] chore: update async worker to use callbacks This commit refactors the async job worker to utilize callbacks for keychain creation and wallet addition. Ticket: WCN-746 --- .../api/master/asyncJobWorker.test.ts | 40 ++++++--- .../integration/asyncJobWorker.integ.test.ts | 6 +- .../handlers/utils/walletCreationUtils.ts | 82 ------------------- .../handlers/walletGenerationCallbacks.ts | 22 ++++- .../workers/asyncJobWorker.ts | 36 ++++---- 5 files changed, 66 insertions(+), 120 deletions(-) diff --git a/src/__tests__/api/master/asyncJobWorker.test.ts b/src/__tests__/api/master/asyncJobWorker.test.ts index a4e7b203..035484a7 100644 --- a/src/__tests__/api/master/asyncJobWorker.test.ts +++ b/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,9 +215,13 @@ 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', }); @@ -225,9 +229,13 @@ describe('asyncJobWorker', () => { 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', }); diff --git a/src/__tests__/integration/asyncJobWorker.integ.test.ts b/src/__tests__/integration/asyncJobWorker.integ.test.ts index e64fc2bb..138ac622 100644 --- a/src/__tests__/integration/asyncJobWorker.integ.test.ts +++ b/src/__tests__/integration/asyncJobWorker.integ.test.ts @@ -40,7 +40,7 @@ function makeAwaitingBitgoJob(overrides: Partial = {}): 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 = {}): 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()]); diff --git a/src/masterBitgoExpress/handlers/utils/walletCreationUtils.ts b/src/masterBitgoExpress/handlers/utils/walletCreationUtils.ts index 6cfa6e19..1fddb9e6 100644 --- a/src/masterBitgoExpress/handlers/utils/walletCreationUtils.ts +++ b/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 { - const baseCoin = await coinFactory.getCoin(coin, bitgo); - const reqId = new RequestTracer(); - - const registerKeychain = async (keyChain: IndependentKeychainResponse): Promise => { - 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', - }; -} diff --git a/src/masterBitgoExpress/handlers/walletGenerationCallbacks.ts b/src/masterBitgoExpress/handlers/walletGenerationCallbacks.ts index 94a2faaf..9f8bf37a 100644 --- a/src/masterBitgoExpress/handlers/walletGenerationCallbacks.ts +++ b/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, +): 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', + }; + }; +} diff --git a/src/masterBitgoExpress/workers/asyncJobWorker.ts b/src/masterBitgoExpress/workers/asyncJobWorker.ts index 4c59cae0..aede98b4 100644 --- a/src/masterBitgoExpress/workers/asyncJobWorker.ts +++ b/src/masterBitgoExpress/workers/asyncJobWorker.ts @@ -1,4 +1,3 @@ -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'; @@ -6,12 +5,10 @@ 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; - - 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}`); }