From ef520da47bc84090b4401c5b672f4711835de60c Mon Sep 17 00:00:00 2001
From: Alejandro Busse <alejandrobusse@bitgo.com>
Date: Fri, 10 Jan 2025 14:25:08 -0300
Subject: [PATCH] fix(root): add retry mechanism for MPC sendSignatureShareV2

added a retry mechanism with a backoff to sendSignatureShareV2 for MPC when the response is a 429
rate limit error

WP-3392

TICKET: WP-3392
---
 .../tssUtils/ecdsaMPCv2/signTxRequest.ts      | 66 ++++++++++++++++++-
 modules/sdk-core/src/bitgo/tss/common.ts      | 23 ++++++-
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/modules/bitgo/test/v2/unit/internal/tssUtils/ecdsaMPCv2/signTxRequest.ts b/modules/bitgo/test/v2/unit/internal/tssUtils/ecdsaMPCv2/signTxRequest.ts
index 4c6ec0c68a..9fe7614ad5 100644
--- a/modules/bitgo/test/v2/unit/internal/tssUtils/ecdsaMPCv2/signTxRequest.ts
+++ b/modules/bitgo/test/v2/unit/internal/tssUtils/ecdsaMPCv2/signTxRequest.ts
@@ -123,6 +123,9 @@ describe('signTxRequest:', function () {
         },
       ],
       curve: 'secp256k1',
+      config: {
+        rejectCurves: new Set(),
+      },
     });
     const constants = {
       mpc: {
@@ -238,6 +241,47 @@ describe('signTxRequest:', function () {
     nockPromises[1].isDone().should.be.true();
     nockPromises[2].isDone().should.be.true();
   });
+
+  it('successfully signs a txRequest for a dkls hot wallet after receiving multiple 429 errors', async function () {
+    const nockPromises = [
+      await nockTxRequestResponseSignatureShareRoundOne(bitgoParty, txRequest, bitgoGpgKey),
+      await nockTxRequestResponseSignatureShareRoundTwo(bitgoParty, txRequest, bitgoGpgKey, 0, 3),
+      await nockTxRequestResponseSignatureShareRoundThree(txRequest),
+      await nockSendTxRequest(txRequest),
+    ];
+    await Promise.all(nockPromises);
+
+    const userShare = fs.readFileSync(shareFiles[vector.party1]);
+    const userPrvBase64 = Buffer.from(userShare).toString('base64');
+    await tssUtils.signTxRequest({
+      txRequest,
+      prv: userPrvBase64,
+      reqId,
+    });
+    nockPromises[0].isDone().should.be.true();
+    nockPromises[1].isDone().should.be.true();
+    nockPromises[2].isDone().should.be.true();
+  });
+
+  it('fails to signs a txRequest for a dkls hot wallet after receiving over 3 429 errors', async function () {
+    const nockPromises = [
+      await nockTxRequestResponseSignatureShareRoundOne(bitgoParty, txRequest, bitgoGpgKey),
+      await nockTxRequestResponseSignatureShareRoundTwo(bitgoParty, txRequest, bitgoGpgKey, 0, 4),
+    ];
+    await Promise.all(nockPromises);
+
+    const userShare = fs.readFileSync(shareFiles[vector.party1]);
+    const userPrvBase64 = Buffer.from(userShare).toString('base64');
+    await tssUtils
+      .signTxRequest({
+        txRequest,
+        prv: userPrvBase64,
+        reqId,
+      })
+      .should.be.rejectedWith('Too many requests, slow down!');
+    nockPromises[0].isDone().should.be.true();
+    nockPromises[1].isDone().should.be.false();
+  });
 });
 
 export function getBitGoPartyGpgKeyPrv(key: openpgp.SerializedKeyPair<string>): DklsTypes.PartyGpgKey {
@@ -336,11 +380,27 @@ async function nockTxRequestResponseSignatureShareRoundTwo(
   bitgoSession: DklsDsg.Dsg,
   txRequest: TxRequest,
   bitgoGpgKey: openpgp.SerializedKeyPair<string>,
-  partyId: 0 | 1 = 0
+  partyId: 0 | 1 = 0,
+  rateLimitErrorCount = 0
 ): Promise<nock.Scope> {
   const transactions = getRoute('ecdsa');
-  return nock('https://bitgo.fakeurl')
-    .persist(true)
+  const scope = nock('https://bitgo.fakeurl');
+
+  if (rateLimitErrorCount > 0) {
+    scope
+      .post(
+        `/api/v2/wallet/${txRequest.walletId}/txrequests/${txRequest.txRequestId + transactions}/sign`,
+        (body) => (JSON.parse(body.signatureShares[0].share) as MPCv2SignatureShareRound2Input).type === 'round2Input'
+      )
+      .times(rateLimitErrorCount)
+      .reply(429, {
+        error: 'Too many requests, slow down!',
+        name: 'TooManyRequests',
+        requestId: 'cm5qx01lh0013b2ek2sxl4w00',
+        context: {},
+      });
+  }
+  return scope
     .post(
       `/api/v2/wallet/${txRequest.walletId}/txrequests/${txRequest.txRequestId + transactions}/sign`,
       (body) => (JSON.parse(body.signatureShares[0].share) as MPCv2SignatureShareRound2Input).type === 'round2Input'
diff --git a/modules/sdk-core/src/bitgo/tss/common.ts b/modules/sdk-core/src/bitgo/tss/common.ts
index c35f137b21..22a34048e6 100644
--- a/modules/sdk-core/src/bitgo/tss/common.ts
+++ b/modules/sdk-core/src/bitgo/tss/common.ts
@@ -15,6 +15,8 @@ import {
 } from '../utils';
 import { IRequestTracer } from '../../api';
 
+const debug = require('debug')('bitgo:tss:common');
+
 /**
  * Gets the latest Tx Request by id
  *
@@ -140,7 +142,26 @@ export async function sendSignatureShareV2(
   };
   const reqTracer = reqId || new RequestTracer();
   bitgo.setRequestTracer(reqTracer);
-  return bitgo.post(bitgo.url(urlPath, 2)).send(requestBody).result();
+
+  let attempts = 0;
+  const maxAttempts = 3;
+
+  while (attempts < maxAttempts) {
+    try {
+      return await bitgo.post(bitgo.url(urlPath, 2)).send(requestBody).result();
+    } catch (err) {
+      if (err?.status === 429) {
+        const sleepTime = 1000 * (attempts + 1);
+        debug(`MPC Signing rate limit error - retrying in ${sleepTime / 1000} seconds`);
+        // sleep for a bit before retrying
+        await new Promise((resolve) => setTimeout(resolve, sleepTime));
+        attempts++;
+      } else {
+        throw err;
+      }
+    }
+  }
+  return await bitgo.post(bitgo.url(urlPath, 2)).send(requestBody).result();
 }
 
 /**