feat(awm): wiring in backup key clients into recovery flows

Ticket: WCN-363

Author: margueriteblair

src/__tests__/api/advancedWalletManager/postMpcV2Key.test.ts

@@ -42,6 +42,10 @@ describe('postMpcV2Key', () => {
       clientCertAllowSelfSigned: true,
     };
 
+    // Restore any existing stub from other test suites before re-stubbing
+    if (typeof (configModule.initConfig as any).restore === 'function') {
+      (configModule.initConfig as any).restore();
+    }
     configStub = sinon.stub(configModule, 'initConfig').returns(cfg);
 
     // app setup
@@ -93,7 +97,7 @@ describe('postMpcV2Key', () => {
   });
 
   after(() => {
-    configStub.restore();
+    configStub?.restore();
   });
 
   it('should be able to create a new MPC V2 wallet', async () => {

src/__tests__/api/advancedWalletManager/recoveryMpcV2.test.ts

@@ -20,7 +20,8 @@ describe('recoveryMpcV2', async () => {
   const cosmosLikeCoin = 'tsei';
   const accessToken = 'test-token';
 
-  // sinon stubs
+  // sinon sandbox
+  const sandbox = sinon.createSandbox();
   let configStub: sinon.SinonStub;
 
   // key provider nocks setup
@@ -67,7 +68,11 @@ describe('recoveryMpcV2', async () => {
       recoveryMode: true,
     };
 
-    configStub = sinon.stub(configModule, 'initConfig').returns(cfg);
+    // Restore any existing stub from other test suites before re-stubbing
+    if (typeof (configModule.initConfig as any).restore === 'function') {
+      (configModule.initConfig as any).restore();
+    }
+    configStub = sandbox.stub(configModule, 'initConfig').returns(cfg);
 
     // app setup
     app = advancedWalletManagerApp(cfg);
@@ -79,7 +84,7 @@ describe('recoveryMpcV2', async () => {
   });
 
   after(() => {
-    configStub.restore();
+    sandbox.restore();
   });
 
   // happy path test
@@ -133,15 +138,30 @@ describe('recoveryMpcV2', async () => {
   });
 
   it('should route backup key retrieval to backup KMS when configured', async () => {
+    const kmsUrl = 'http://kms.invalid';
     const backupKmsUrl = 'http://backup-kms.invalid';
 
+    const mockKmsUserResponse = {
+      prv: JSON.stringify(userKeyShare),
+      pub: commonKeychain,
+      source: 'user',
+      type: 'tss',
+    };
+
+    const mockKmsBackupResponse = {
+      prv: JSON.stringify(backupKeyShare),
+      pub: commonKeychain,
+      source: 'backup',
+      type: 'tss',
+    };
+
     // Reconfigure app with backup KMS URL
-    configStub.restore();
     const dualCfg: AdvancedWalletManagerConfig = {
       ...cfg,
+      keyProviderUrl: kmsUrl,
       backupKmsUrl,
     };
-    configStub = sinon.stub(configModule, 'initConfig').returns(dualCfg);
+    configStub.returns(dualCfg);
     const dualApp = advancedWalletManagerApp(dualCfg);
     const dualAgent = request.agent(dualApp);
 

src/__tests__/api/advancedWalletManager/recoveryMusigEth.test.ts

@@ -22,7 +22,8 @@ describe('recoveryMultisigTransaction', () => {
   const coin = 'hteth';
   const accessToken = 'test-token';
 
-  // sinon stubs
+  // sinon sandbox
+  const sandbox = sinon.createSandbox();
   let configStub: sinon.SinonStub;
 
   before(() => {
@@ -44,7 +45,11 @@ describe('recoveryMultisigTransaction', () => {
       recoveryMode: true,
     };
 
-    configStub = sinon.stub(configModule, 'initConfig').returns(cfg);
+    // Restore any existing stub from other test suites before re-stubbing
+    if (typeof (configModule.initConfig as any).restore === 'function') {
+      (configModule.initConfig as any).restore();
+    }
+    configStub = sandbox.stub(configModule, 'initConfig').returns(cfg);
 
     // app setup
     app = advancedWalletManagerApp(cfg);
@@ -56,7 +61,7 @@ describe('recoveryMultisigTransaction', () => {
   });
 
   after(() => {
-    configStub.restore();
+    sandbox.restore();
   });
 
   it('should generate a successful txHex from unsigned sweep prebuild data', async () => {
@@ -107,17 +112,18 @@ describe('recoveryMultisigTransaction', () => {
   });
 
   it('should route backup key retrieval to backup KMS when configured', async () => {
+    const kmsUrl = 'http://kms.invalid';
     const backupKmsUrl = 'http://backup-kms.invalid';
     const { userPub, backupPub, walletContractAddress, userPrv, backupPrv, txHexResult } = awmData;
     const unsignedSweepPrebuildTx = unsignedSweepRecJSON as unknown as any;
 
     // Reconfigure app with backup KMS URL
-    configStub.restore();
     const dualCfg: AdvancedWalletManagerConfig = {
       ...cfg,
+      keyProviderUrl: kmsUrl,
       backupKmsUrl,
     };
-    configStub = sinon.stub(configModule, 'initConfig').returns(dualCfg);
+    configStub.returns(dualCfg);
     const dualApp = advancedWalletManagerApp(dualCfg);
     const dualAgent = request.agent(dualApp);
 

src/advancedWalletManager/handlers/ecdsaMPCV2Recovery.ts

@@ -52,11 +52,15 @@ export async function ecdsaMPCv2Recovery(
     );
   }
 
-  // setup clients and retreive the keys
-  // TODO: this needs to be segerated if the EBE instance cannot retrieve both keys
-  const keyProvider = new KeyProviderClient(req.config);
-  const { prv: userPrv } = await keyProvider.getKey({ pub, source: 'user' });
-  const { prv: backupPrv } = await keyProvider.getKey({ pub, source: 'backup' });
+  // setup clients and retrieve the keys
+  const userKeyProvider = new KeyProviderClient(req.config);
+  const backupCfg =
+    req.config.backupKmsUrl
+      ? { ...req.config, keyProviderUrl: req.config.backupKmsUrl }
+      : req.config;
+  const backupKeyProvider = new KeyProviderClient(backupCfg);
+  const { prv: userPrv } = await userKeyProvider.getKey({ pub, source: 'user' });
+  const { prv: backupPrv } = await backupKeyProvider.getKey({ pub, source: 'backup' });
 
   // construct tx builder
   const txHash = await getMessageHash(coin, txHex);

src/advancedWalletManager/handlers/utils/utils.ts

@@ -16,7 +16,11 @@ export async function retrieveKeyProviderPrvKey({
   source: string;
   cfg: AdvancedWalletManagerConfig;
 }): Promise<string> {
-  const keyProvider = new KeyProviderClient(cfg);
+  const effectiveCfg =
+    source === 'backup' && cfg.backupKmsUrl
+      ? { ...cfg, keyProviderUrl: cfg.backupKmsUrl }
+      : cfg;
+  const keyProvider = new KeyProviderClient(effectiveCfg);
   // Retrieve the private key from key provider
   let prv: string;
   try {