From 4878592fb2708a0ca009b6f27b8e61326af745de Mon Sep 17 00:00:00 2001 From: Abhishek Kumar Date: Mon, 17 Aug 2026 15:32:54 +0530 Subject: [PATCH] fix: add batching for keys accept decrypt Ticket: POL-929 --- modules/sdk-core/src/bitgo/wallet/wallets.ts | 125 ++++++++++-------- .../bitgo/wallet/walletsEncryptionVersion.ts | 53 ++++++++ 2 files changed, 120 insertions(+), 58 deletions(-) diff --git a/modules/sdk-core/src/bitgo/wallet/wallets.ts b/modules/sdk-core/src/bitgo/wallet/wallets.ts index 55959f3c5a..770f00c602 100644 --- a/modules/sdk-core/src/bitgo/wallet/wallets.ts +++ b/modules/sdk-core/src/bitgo/wallet/wallets.ts @@ -1308,67 +1308,76 @@ export class Wallets implements IWallets { }); const newWalletPassphrase = params.newWalletPassphrase || params.userLoginPassword; const webauthnInfo = params.webauthnInfo; - const keysForWalletShares = ( - await Promise.all( - walletShares.map(async (walletShare) => { - // Handle userMultiKeyRotationRequired case - these shares don't have keychains - if (walletShare.userMultiKeyRotationRequired) { - if (!params.userLoginPassword) { - throw new Error('userLoginPassword param must be provided to generate user keychain'); - } - const walletKeychain = this.baseCoin.keychains().create(); - const encryptedPrv = await this.bitgo.encrypt({ - password: newWalletPassphrase, - input: walletKeychain.prv, - encryptionVersion: params.encryptionVersion, - }); - return [ - { - walletShareId: walletShare.id, - encryptedPrv: encryptedPrv, - pub: walletKeychain.pub, - }, - ]; - } - // Standard case: shares with keychains - if (!walletShare.keychain) { - return []; - } - const secret = getSharedSecret( - bip32.fromBase58(sharingKeychain.prv).derivePath(sanitizeLegacyPath(walletShare.keychain.path)), - Buffer.from(walletShare.keychain.fromPubKey, 'hex') - ).toString('hex'); - - const decryptedSharedWalletPrv = await this.bitgo.decrypt({ - password: secret, - input: walletShare.keychain.encryptedPrv, - }); - const newEncryptedPrv = await this.bitgo.encrypt({ - password: newWalletPassphrase, + // Each decrypt/encrypt call runs Argon2id inside a WebAssembly instance that reserves ~2 GiB of + // virtual address space. Running all shares concurrently via Promise.all exhausts the browser's + // WASM memory at scale (e.g. 96 wallets). Process in small batches so only a bounded number of + // WASM instances are alive at once. + const BATCH_SIZE = 16; + + const processShare = async (walletShare: WalletShare): Promise => { + // Handle userMultiKeyRotationRequired case - these shares don't have keychains + if (walletShare.userMultiKeyRotationRequired) { + if (!params.userLoginPassword) { + throw new Error('userLoginPassword param must be provided to generate user keychain'); + } + const walletKeychain = this.baseCoin.keychains().create(); + const encryptedPrv = await this.bitgo.encrypt({ + password: newWalletPassphrase, + input: walletKeychain.prv, + encryptionVersion: params.encryptionVersion, + }); + return [ + { + walletShareId: walletShare.id, + encryptedPrv: encryptedPrv, + pub: walletKeychain.pub, + }, + ]; + } + + // Standard case: shares with keychains + if (!walletShare.keychain) { + return []; + } + const secret = getSharedSecret( + bip32.fromBase58(sharingKeychain.prv).derivePath(sanitizeLegacyPath(walletShare.keychain.path)), + Buffer.from(walletShare.keychain.fromPubKey, 'hex') + ).toString('hex'); + + const decryptedSharedWalletPrv = await this.bitgo.decrypt({ + password: secret, + input: walletShare.keychain.encryptedPrv, + }); + const newEncryptedPrv = await this.bitgo.encrypt({ + password: newWalletPassphrase, + input: decryptedSharedWalletPrv, + encryptionVersion: params.encryptionVersion, + }); + const entry: AcceptShareOptionsRequest = { + walletShareId: walletShare.id, + encryptedPrv: newEncryptedPrv, + }; + if (webauthnInfo) { + entry.webauthnInfo = { + otpDeviceId: webauthnInfo.otpDeviceId, + prfSalt: webauthnInfo.prfSalt, + encryptedPrv: await this.bitgo.encrypt({ + password: webauthnInfo.passphrase, input: decryptedSharedWalletPrv, encryptionVersion: params.encryptionVersion, - }); - const entry: AcceptShareOptionsRequest = { - walletShareId: walletShare.id, - encryptedPrv: newEncryptedPrv, - }; - if (webauthnInfo) { - entry.webauthnInfo = { - otpDeviceId: webauthnInfo.otpDeviceId, - prfSalt: webauthnInfo.prfSalt, - encryptedPrv: await this.bitgo.encrypt({ - password: webauthnInfo.passphrase, - input: decryptedSharedWalletPrv, - encryptionVersion: params.encryptionVersion, - adata: walletShare.enterprise, - }), - }; - } - return [entry]; - }) - ) - ).flat(); + adata: walletShare.enterprise, + }), + }; + } + return [entry]; + }; + + const keysForWalletShares: AcceptShareOptionsRequest[] = []; + for (const batch of _.chunk(walletShares, BATCH_SIZE)) { + const batchResults = await Promise.all(batch.map((walletShare) => processShare(walletShare))); + keysForWalletShares.push(...batchResults.flat()); + } return this.bulkAcceptShareRequest(keysForWalletShares); } diff --git a/modules/sdk-core/test/unit/bitgo/wallet/walletsEncryptionVersion.ts b/modules/sdk-core/test/unit/bitgo/wallet/walletsEncryptionVersion.ts index 13f8170b06..2c2ae030e1 100644 --- a/modules/sdk-core/test/unit/bitgo/wallet/walletsEncryptionVersion.ts +++ b/modules/sdk-core/test/unit/bitgo/wallet/walletsEncryptionVersion.ts @@ -135,6 +135,59 @@ describe('Wallets - encryptionVersion threading', function () { const call = mockBitGo.encrypt.firstCall; assert.strictEqual(call.args[0].encryptionVersion, undefined); }); + + it('processes shares in batches of 16 to avoid WASM memory exhaustion', async function () { + // 20 shares → 2 batches (16 + 4), verifying the batch boundary is crossed + const manyShares = Array.from({ length: 20 }, (_, i) => ({ + id: `share-${i}`, + userMultiKeyRotationRequired: true, + keychain: null, + permissions: ['spend'], + })); + mockBitGo.get.returns({ + result: sinon.stub().resolves({ incoming: manyShares, outgoing: [] }), + }); + + await wallets.bulkAcceptShare({ + walletShareIds: manyShares.map((s) => s.id), + userLoginPassword: 'login-password', + }); + + // All 20 shares should have been encrypted (one encrypt call per share) + assert.strictEqual(mockBitGo.encrypt.callCount, 20); + }); + + it('never runs more than 16 shares concurrently', async function () { + let inFlight = 0; + let maxInFlight = 0; + + mockBitGo.encrypt.callsFake(() => { + inFlight++; + maxInFlight = Math.max(maxInFlight, inFlight); + return Promise.resolve('encrypted').then((r) => { + inFlight--; + return r; + }); + }); + + const manyShares = Array.from({ length: 20 }, (_, i) => ({ + id: `share-${i}`, + userMultiKeyRotationRequired: true, + keychain: null, + permissions: ['spend'], + })); + mockBitGo.get.returns({ + result: sinon.stub().resolves({ incoming: manyShares, outgoing: [] }), + }); + + await wallets.bulkAcceptShare({ + walletShareIds: manyShares.map((s) => s.id), + userLoginPassword: 'login-password', + }); + + assert.ok(maxInFlight <= 16, `expected max concurrency <= 16, got ${maxInFlight}`); + assert.strictEqual(mockBitGo.encrypt.callCount, 20); + }); }); describe('Wallet.shareWallet / createBulkWalletShare', function () {