Merge pull request #6688 from BitGo/CSI-889-fix-bulk-accept-share-request-entity-too-large

feat(sdk-core): add retry logic for bulk accept share requests

Author: sdoshibitgo

modules/bitgo/test/v2/unit/wallets.ts

@@ -22,6 +22,7 @@ import {
   makeRandomKey,
   getSharedSecret,
   BulkWalletShareOptions,
+  AcceptShareOptionsRequest,
   KeychainWithEncryptedPrv,
   WalletWithKeychains,
   multisigTypes,
@@ -1872,6 +1873,283 @@ describe('V2 Wallets:', function () {
           ],
         });
       });
+
+      it('should handle 413 payload too large error with smart retry', async () => {
+        const walletPassphrase = 'bitgo1234';
+        const fromUserPrv = Math.random();
+        const keychainTest: OptionalKeychainEncryptedKey = {
+          encryptedPrv: bitgo.encrypt({ input: fromUserPrv.toString(), password: walletPassphrase }),
+        };
+        const userPrv = decryptKeychainPrivateKey(bitgo, keychainTest, walletPassphrase);
+        if (!userPrv) {
+          throw new Error('Unable to decrypt user keychain');
+        }
+
+        const toKeychain = utxoLib.bip32.fromSeed(Buffer.from('deadbeef02deadbeef02deadbeef02deadbeef02', 'hex'));
+        const path = 'm/999999/1/1';
+        const pubkey = toKeychain.derivePath(path).publicKey.toString('hex');
+
+        const eckey = makeRandomKey();
+        const secret = getSharedSecret(eckey, Buffer.from(pubkey, 'hex')).toString('hex');
+        // Pad the private key with additional data to make it larger before encrypting
+        const newEncryptedPrv = bitgo.encrypt({ password: secret, input: userPrv });
+        const keychain = {
+          path: path,
+          fromPubKey: eckey.publicKey.toString('hex'),
+          encryptedPrv: newEncryptedPrv,
+          toPubKey: pubkey,
+          pub: pubkey,
+        };
+        const shareIds = Array.from({ length: 20 }, (_, i) => `share${i + 1}`);
+
+        // Mock listSharesV2 to return 25 shares
+        const shares = shareIds.map((id, index) => ({
+          id,
+          coin: 'tsol',
+          walletLabel: `testing${index}`,
+          fromUser: 'dummyFromUser',
+          toUser: 'dummyToUser',
+          wallet: `wallet${index}`,
+          permissions: ['spend'],
+          state: 'active' as const,
+          keychain: keychain,
+        }));
+
+        sinon.stub(Wallets.prototype, 'listSharesV2').resolves({
+          incoming: shares,
+          outgoing: [],
+        });
+
+        const myEcdhKeychain = await bitgo.keychains().create();
+        sinon.stub(bitgo, 'getECDHKeychain').resolves({
+          encryptedXprv: bitgo.encrypt({ input: myEcdhKeychain.xprv, password: walletPassphrase }),
+        });
+
+        const prvKey = bitgo.decrypt({
+          password: walletPassphrase,
+          input: bitgo.encrypt({ input: myEcdhKeychain.xprv, password: walletPassphrase }),
+        });
+        sinon.stub(bitgo, 'decrypt').returns(prvKey);
+        sinon.stub(bitgo, 'encrypt').returns(userPrv + 'X'.repeat(100000));
+        sinon.stub(moduleBitgo, 'getSharedSecret').resolves('fakeSharedSecret');
+
+        // Mock bulkAcceptShareRequestWithRetry to track batch sizes
+        const batchSizes: number[] = [];
+
+        nock(bgUrl)
+          .persist() // This ensures the interceptor remains active for multiple requests
+          .put('/api/v2/walletshares/accept')
+          .reply(function (_, requestBody, cb) {
+            const params = requestBody['keysForWalletShares'] as AcceptShareOptionsRequest[];
+            batchSizes.push(params.length);
+            if (Buffer.byteLength(JSON.stringify(requestBody), 'utf8') > 950000) {
+              // Simulate 413 error
+              return cb(null, [413, { error: 'Request Entity Too Large' }]);
+            }
+            // Return success for smaller batches
+            return cb(null, [
+              200,
+              {
+                acceptedWalletShares: params.map((param) => ({
+                  walletShareId: param.walletShareId,
+                })),
+              },
+            ]);
+          });
+
+        const result = await wallets.bulkAcceptShare({
+          walletShareIds: shareIds,
+          userLoginPassword: walletPassphrase,
+        });
+
+        // Should have tried with 20 (initial batch size for 25 items), then retried with smaller batches
+        batchSizes.length.should.be.greaterThan(1);
+        batchSizes.should.deepEqual([9, 9, 2]); // Initial batch size// Retry batches should be smaller
+
+        result.should.have.property('acceptedWalletShares');
+        result.acceptedWalletShares.should.be.an.Array();
+        result.acceptedWalletShares.length.should.equal(20);
+        result.acceptedWalletShares.forEach((share) => {
+          share.should.have.property('walletShareId');
+          share.walletShareId.should.match(/^share\d+$/);
+        });
+      });
+
+      it('should retry with progressively smaller batch sizes on 413 errors', async () => {
+        const walletPassphrase = 'bitgo1234';
+        const fromUserPrv = Math.random();
+        const keychainTest: OptionalKeychainEncryptedKey = {
+          encryptedPrv: bitgo.encrypt({ input: fromUserPrv.toString(), password: walletPassphrase }),
+        };
+        const userPrv = decryptKeychainPrivateKey(bitgo, keychainTest, walletPassphrase);
+        if (!userPrv) {
+          throw new Error('Unable to decrypt user keychain');
+        }
+
+        const toKeychain = utxoLib.bip32.fromSeed(Buffer.from('deadbeef02deadbeef02deadbeef02deadbeef02', 'hex'));
+        const path = 'm/999999/1/1';
+        const pubkey = toKeychain.derivePath(path).publicKey.toString('hex');
+
+        const eckey = makeRandomKey();
+        const secret = getSharedSecret(eckey, Buffer.from(pubkey, 'hex')).toString('hex');
+        const newEncryptedPrv = bitgo.encrypt({ password: secret, input: userPrv });
+        const keychain = {
+          path: path,
+          fromPubKey: eckey.publicKey.toString('hex'),
+          encryptedPrv: newEncryptedPrv,
+          toPubKey: pubkey,
+          pub: pubkey,
+        };
+        const shareIds = Array.from({ length: 20 }, (_, i) => `share${i + 1}`);
+
+        // Mock listSharesV2
+        const shares = shareIds.map((id, index) => ({
+          id,
+          coin: 'tsol',
+          walletLabel: `testing${index}`,
+          fromUser: 'dummyFromUser',
+          toUser: 'dummyToUser',
+          wallet: `wallet${index}`,
+          permissions: ['spend'],
+          state: 'active' as const,
+          keychain: keychain,
+        }));
+
+        sinon.stub(Wallets.prototype, 'listSharesV2').resolves({
+          incoming: shares,
+          outgoing: [],
+        });
+
+        const myEcdhKeychain = await bitgo.keychains().create();
+        sinon.stub(bitgo, 'getECDHKeychain').resolves({
+          encryptedXprv: bitgo.encrypt({ input: myEcdhKeychain.xprv, password: walletPassphrase }),
+        });
+
+        const prvKey = bitgo.decrypt({
+          password: walletPassphrase,
+          input: bitgo.encrypt({ input: myEcdhKeychain.xprv, password: walletPassphrase }),
+        });
+        sinon.stub(bitgo, 'decrypt').returns(prvKey);
+        sinon.stub(bitgo, 'encrypt').returns(userPrv + 'X'.repeat(100000));
+        sinon.stub(moduleBitgo, 'getSharedSecret').resolves('fakeSharedSecret');
+
+        // Track the sequence of batch sizes attempted
+        const batchSizeAttempts: number[] = [];
+
+        nock(bgUrl)
+          .persist() // This ensures the interceptor remains active for multiple requests
+          .put('/api/v2/walletshares/accept')
+          .reply(function (_, requestBody: any, cb) {
+            const params = requestBody['keysForWalletShares'] as AcceptShareOptionsRequest[];
+            batchSizeAttempts.push(params.length);
+
+            // Simulate 413 for batches > 5, success for batches <= 5
+            if (Buffer.byteLength(JSON.stringify(requestBody), 'utf8') > 600000) {
+              // Simulate 413 error
+              return cb(null, [413, { error: 'Request Entity Too Large' }]);
+            }
+
+            // Return success for smaller batches
+            return cb(null, [
+              200,
+              {
+                acceptedWalletShares: params.map((param) => ({
+                  walletShareId: param.walletShareId || 'test',
+                })),
+              },
+            ]);
+          });
+
+        const result = await wallets.bulkAcceptShare({
+          walletShareIds: shareIds,
+          userLoginPassword: walletPassphrase,
+        });
+
+        // Should see progressive batch size reduction: 20 -> 10 -> 5 (success)
+        batchSizeAttempts.should.containDeep([9, 4, 4, 4, 4, 4]);
+
+        result.should.have.property('acceptedWalletShares');
+        result.acceptedWalletShares.should.be.an.Array();
+        result.acceptedWalletShares.length.should.equal(20);
+        result.acceptedWalletShares.forEach((share) => {
+          share.should.have.property('walletShareId');
+          share.walletShareId.should.match(/^share\d+$/);
+        });
+      });
+
+      it('should throw error when batch size cannot be reduced further', async () => {
+        const walletPassphrase = 'bitgo1234';
+        const fromUserPrv = Math.random();
+        const keychainTest: OptionalKeychainEncryptedKey = {
+          encryptedPrv: bitgo.encrypt({ input: fromUserPrv.toString(), password: walletPassphrase }),
+        };
+        const userPrv = decryptKeychainPrivateKey(bitgo, keychainTest, walletPassphrase);
+        if (!userPrv) {
+          throw new Error('Unable to decrypt user keychain');
+        }
+
+        const toKeychain = utxoLib.bip32.fromSeed(Buffer.from('deadbeef02deadbeef02deadbeef02deadbeef02', 'hex'));
+        const path = 'm/999999/1/1';
+        const pubkey = toKeychain.derivePath(path).publicKey.toString('hex');
+
+        const eckey = makeRandomKey();
+        const secret = getSharedSecret(eckey, Buffer.from(pubkey, 'hex')).toString('hex');
+        const newEncryptedPrv = bitgo.encrypt({ password: secret, input: userPrv });
+        const keychain = {
+          path: path,
+          fromPubKey: eckey.publicKey.toString('hex'),
+          encryptedPrv: newEncryptedPrv,
+          toPubKey: pubkey,
+          pub: pubkey,
+        };
+        const shareIds = ['share1'];
+
+        // Mock listSharesV2
+        sinon.stub(Wallets.prototype, 'listSharesV2').resolves({
+          incoming: [
+            {
+              id: 'share1',
+              coin: 'tsol',
+              walletLabel: 'testing',
+              fromUser: 'dummyFromUser',
+              toUser: 'dummyToUser',
+              wallet: 'wallet1',
+              permissions: ['spend'],
+              state: 'active',
+              keychain: keychain,
+            },
+          ],
+          outgoing: [],
+        });
+
+        const myEcdhKeychain = await bitgo.keychains().create();
+        sinon.stub(bitgo, 'getECDHKeychain').resolves({
+          encryptedXprv: bitgo.encrypt({ input: myEcdhKeychain.xprv, password: walletPassphrase }),
+        });
+
+        const prvKey = bitgo.decrypt({
+          password: walletPassphrase,
+          input: bitgo.encrypt({ input: myEcdhKeychain.xprv, password: walletPassphrase }),
+        });
+        sinon.stub(bitgo, 'decrypt').returns(prvKey);
+        sinon.stub(moduleBitgo, 'getSharedSecret').resolves('fakeSharedSecret');
+
+        // Always throw 413 error, even for batch size 1
+        nock(bgUrl)
+          .persist()
+          .put('/api/v2/walletshares/accept')
+          .reply(function (_, _requestBody, cb) {
+            // Always respond with 413 error
+            return cb(null, [413, { error: 'Request Entity Too Large' }]);
+          });
+
+        await wallets
+          .bulkAcceptShare({
+            walletShareIds: shareIds,
+            userLoginPassword: walletPassphrase,
+          })
+          .should.be.rejectedWith('Request Entity Too Large');
+      });
     });
 
     describe('bulkUpdateWalletShare', function () {

modules/sdk-core/src/bitgo/wallet/wallets.ts

@@ -639,12 +639,69 @@ export class Wallets implements IWallets {
    * @returns {Promise<BulkAcceptShareResponse>}
    */
   async bulkAcceptShareRequest(params: AcceptShareOptionsRequest[]): Promise<BulkAcceptShareResponse> {
-    return await this.bitgo
-      .put(this.bitgo.url('/walletshares/accept', 2))
-      .send({
-        keysForWalletShares: params,
-      })
-      .result();
+    return await this.bulkAcceptShareRequestWithRetry(params);
+  }
+
+  private async bulkAcceptShareRequestWithRetry(params: AcceptShareOptionsRequest[]): Promise<BulkAcceptShareResponse> {
+    // Server has a limit of approximately 1MB for payload size
+    let MAX_PAYLOAD_SIZE = 950000; // ~950KB to leave some buffer
+
+    // Function to calculate the size of a payload
+    const calculatePayloadSize = (items: AcceptShareOptionsRequest[]): number => {
+      return Buffer.byteLength(JSON.stringify({ keysForWalletShares: items }), 'utf8');
+    };
+
+    const results: any[] = [];
+    const remainingParams = [...params];
+
+    while (remainingParams.length > 0) {
+      // Build optimal batch by adding items until we reach size limit
+      const batch: AcceptShareOptionsRequest[] = [];
+      // Start with empty batch
+
+      // Add items one by one while monitoring payload size
+      while (remainingParams.length > 0) {
+        // Test adding the next item
+        const testBatch = [...batch, remainingParams[0]];
+        const testSize = calculatePayloadSize(testBatch);
+
+        // If adding this item would exceed the size limit, stop adding
+        if (testSize > MAX_PAYLOAD_SIZE && batch.length > 0) {
+          break;
+        }
+
+        // Otherwise, add the item to the batch
+        batch.push(remainingParams.shift()!);
+      }
+
+      // Handle case where even a single item is too large
+      if (batch.length === 0 && remainingParams.length > 0) {
+        // Send just the first item even if it's oversized
+        batch.push(remainingParams.shift()!);
+      }
+
+      const payloadObj = { keysForWalletShares: batch };
+
+      try {
+        const result = await this.bitgo.put(this.bitgo.url('/walletshares/accept', 2)).send(payloadObj).result();
+
+        if (result.acceptedWalletShares && Array.isArray(result.acceptedWalletShares)) {
+          results.push(...result.acceptedWalletShares);
+        }
+      } catch (error: any) {
+        if (error.status === 413 && batch.length > 1) {
+          // If we still get 413 with multiple items, put them back and try with half the batch size
+          remainingParams.unshift(...batch);
+          MAX_PAYLOAD_SIZE = Math.floor(MAX_PAYLOAD_SIZE / 2); // Reduce size limit for next attempt
+          continue;
+        }
+        throw error;
+      }
+    }
+
+    return {
+      acceptedWalletShares: results,
+    };
   }
 
   async bulkUpdateWalletShareRequest(