feat(sdk-core): getUserAndBackupSession + createKeychains retrofit wiring

modules/sdk-core/src/bitgo/utils/tss/eddsa/eddsaMPCv2.ts

@@ -16,7 +16,7 @@ import { EddsaMPCv2KeyGenCallbacks } from '../../../wallet/iWallets';
 import { ed25519 } from '@noble/curves/ed25519';
 import { EddsaMPSDkg, EddsaMPSDsg, MPSComms, MPSTypes, MPSUtil } from '@bitgo/sdk-lib-mpc';
 import { KeychainsTriplet } from '../../../baseCoin';
-import { AddKeychainOptions, Keychain, KeyType, WebauthnKeyEncryptionInfo } from '../../../keychain';
+import { AddKeychainOptions, DecryptedRetrofitPayload, Keychain, KeyType, WebauthnKeyEncryptionInfo } from '../../../keychain';
 import { envRequiresBitgoPubGpgKeyConfig, isBitgoEddsaMpcv2PubKey } from '../../../tss/bitgoPubKeys';
 import { getBitgoSignatureShare, getTxRequest, sendSignatureShareV2, sendTxRequest } from '../../../tss/common';
 import { decodeWithCodec } from '../../codecs';
@@ -63,6 +63,7 @@ export class EddsaMPCv2Utils extends BaseEddsaUtils {
     passphrase: string;
     enterprise: string;
     originalPasscodeEncryptionCode?: string;
+    retrofit?: DecryptedRetrofitPayload;
     webauthnInfo?: WebauthnKeyEncryptionInfo;
     encryptionVersion?: EncryptionVersion;
     // Wallet Safes v1 (@experimental): tags the resulting user/backup/bitgo root keys with this safe.
@@ -93,8 +94,7 @@ export class EddsaMPCv2Utils extends BaseEddsaUtils {
     const bitgoPk = await MPSComms.extractEd25519PublicKey(bitgoKeyObj);
 
     // Create DKG sessions for user (party 0) and backup (party 1)
-    const userDkg = new EddsaMPSDkg.DKG(3, 2, MPCv2PartiesEnum.USER);
-    const backupDkg = new EddsaMPSDkg.DKG(3, 2, MPCv2PartiesEnum.BACKUP);
+    const { userDkg, backupDkg } = await this.getUserAndBackupSession(params.retrofit);
 
     // #region round 1
     await userDkg.initDkg(userSk, [backupPk, bitgoPk]);
@@ -116,6 +116,7 @@ export class EddsaMPCv2Utils extends BaseEddsaUtils {
         backupGpgPublicKey,
         userMsg1: userSignedMsg1,
         backupMsg1: backupSignedMsg1,
+        ...(params.retrofit?.walletId ? { walletId: params.retrofit.walletId } : {}),
       },
       params.safeId
     );
@@ -459,15 +460,15 @@ export class EddsaMPCv2Utils extends BaseEddsaUtils {
 
   async sendKeyGenerationRound1(
     enterprise: string,
-    payload: EddsaMPCv2KeyGenRound1Request,
+    payload: EddsaMPCv2KeyGenRound1Request & { walletId?: string },
     safeId?: string
   ): Promise<EddsaMPCv2KeyGenRound1Response> {
     return this.sendKeyGenerationRound1BySender(KeyGenSenderForEnterprise(this.bitgo, enterprise, safeId), payload);
   }
 
   async sendKeyGenerationRound1BySender(
     senderFn: EddsaMPCv2KeyGenSendFn<EddsaMPCv2KeyGenRound1Response>,
-    payload: EddsaMPCv2KeyGenRound1Request
+    payload: EddsaMPCv2KeyGenRound1Request & { walletId?: string }
   ): Promise<EddsaMPCv2KeyGenRound1Response> {
     return senderFn(MPCv2KeyGenStateEnum['MPCv2-R1'], payload);
   }
@@ -1070,6 +1071,26 @@ export class EddsaMPCv2Utils extends BaseEddsaUtils {
 
   // #region retrofit
 
+  private async getUserAndBackupSession(retrofit?: DecryptedRetrofitPayload): Promise<{
+    userDkg: EddsaMPSDkg.DKG;
+    backupDkg: EddsaMPSDkg.DKG;
+  }> {
+    if (retrofit) {
+      const { userRetrofitData, backupRetrofitData } = await this.getMpcV2RetrofitDataFromMpcV1Keys({
+        mpcv1UserKeyShare: retrofit.decryptedUserKey,
+        mpcv1BackupKeyShare: retrofit.decryptedBackupKey,
+      });
+      return {
+        userDkg: new EddsaMPSDkg.DKG(3, 2, MPCv2PartiesEnum.USER, userRetrofitData),
+        backupDkg: new EddsaMPSDkg.DKG(3, 2, MPCv2PartiesEnum.BACKUP, backupRetrofitData),
+      };
+    }
+    return {
+      userDkg: new EddsaMPSDkg.DKG(3, 2, MPCv2PartiesEnum.USER),
+      backupDkg: new EddsaMPSDkg.DKG(3, 2, MPCv2PartiesEnum.BACKUP),
+    };
+  }
+
   async getMpcV2RetrofitDataFromMpcV1Keys(params: { mpcv1UserKeyShare: string; mpcv1BackupKeyShare: string }): Promise<{
     userRetrofitData: MPSTypes.EddsaRetrofitData;
     backupRetrofitData: MPSTypes.EddsaRetrofitData;

modules/sdk-core/test/unit/bitgo/utils/tss/eddsa/eddsaMPCv2.ts

@@ -2467,3 +2467,149 @@ describe('EddsaMPCv2Utils.getMpcV2RetrofitDataFromMpcV1Keys', () => {
     );
   });
 });
+
+describe('EddsaMPCv2Utils.getUserAndBackupSession', () => {
+  let utils: EddsaMPCv2Utils;
+  let userSigningMaterial: Record<string, unknown>;
+  let backupSigningMaterial: Record<string, unknown>;
+
+  before(async () => {
+    const MPC = await getInitializedMpcInstance();
+    const user = MPC.keyShare(1, 2, 3);
+    const backup = MPC.keyShare(2, 2, 3);
+    const bitgo = MPC.keyShare(3, 2, 3);
+    userSigningMaterial = {
+      uShare: user.uShare,
+      bitgoYShare: bitgo.yShares[1],
+      backupYShare: backup.yShares[1],
+    };
+    backupSigningMaterial = {
+      uShare: backup.uShare,
+      bitgoYShare: bitgo.yShares[2],
+      userYShare: user.yShares[2],
+    };
+  });
+
+  beforeEach(() => {
+    const mockBitGo = {} as unknown as BitGoBase;
+    const mockCoin = {} as unknown as IBaseCoin;
+    utils = new EddsaMPCv2Utils(mockBitGo, mockCoin);
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it('returns plain DKG sessions when retrofit is undefined', async () => {
+    const { userDkg, backupDkg } = await (utils as any).getUserAndBackupSession(undefined);
+    assert.ok(userDkg, 'user DKG should be created');
+    assert.ok(backupDkg, 'backup DKG should be created');
+  });
+
+  it('returns retrofit-seeded DKG sessions when retrofit payload is supplied', async () => {
+    const retrofit = {
+      decryptedUserKey: JSON.stringify(userSigningMaterial),
+      decryptedBackupKey: JSON.stringify(backupSigningMaterial),
+      walletId: 'wallet-123',
+    };
+    const { userDkg, backupDkg } = await (utils as any).getUserAndBackupSession(retrofit);
+    assert.ok(userDkg, 'user DKG should be created with retrofit data');
+    assert.ok(backupDkg, 'backup DKG should be created with retrofit data');
+  });
+});
+
+describe('EddsaMPCv2Utils.createKeychains with retrofit wiring', () => {
+  let utils: EddsaMPCv2Utils;
+  let userSigningMaterial: Record<string, unknown>;
+  let backupSigningMaterial: Record<string, unknown>;
+  let bitgoGpgPublicKeyArmored: string;
+  const enterprise = 'enterprise-id';
+  const sessionId = 'session-001';
+  const walletId = 'wallet-retrofit-123';
+
+  before(async () => {
+    const MPC = await getInitializedMpcInstance();
+    const user = MPC.keyShare(1, 2, 3);
+    const backup = MPC.keyShare(2, 2, 3);
+    const bitgo = MPC.keyShare(3, 2, 3);
+    userSigningMaterial = {
+      uShare: user.uShare,
+      bitgoYShare: bitgo.yShares[1],
+      backupYShare: backup.yShares[1],
+    };
+    backupSigningMaterial = {
+      uShare: backup.uShare,
+      bitgoYShare: bitgo.yShares[2],
+      userYShare: user.yShares[2],
+    };
+    // Generate a real Ed25519 GPG key to stand in for the BitGo GPG key
+    const bitgoGpgKeyPair = await generateGPGKeyPair('ed25519');
+    bitgoGpgPublicKeyArmored = bitgoGpgKeyPair.publicKey;
+  });
+
+  beforeEach(() => {
+    const mockBitGo = {
+      getEnv: sinon.stub().returns('dev'),
+      encrypt: sinon.stub().resolves('encrypted'),
+    } as any;
+    const mockKeychains = {
+      add: sinon
+        .stub()
+        .callsFake((params: any) =>
+          Promise.resolve({ id: `${params.source}-key-id`, commonKeychain: 'a'.repeat(128), isMPCv2: true })
+        ),
+    };
+    const mockCoin = {
+      keychains: sinon.stub().returns(mockKeychains),
+    } as any;
+
+    utils = new EddsaMPCv2Utils(mockBitGo, mockCoin);
+    sinon.stub(utils, 'getBitgoGpgPubkeyBasedOnFeatureFlags' as any).resolves({ eddsaMpcv2PublicKey: null });
+    // Use a real armored GPG public key so pgp.readKey() succeeds inside createKeychains
+    (utils as any).bitgoEddsaMpcv2PublicGpgKey = { armor: () => bitgoGpgPublicKeyArmored };
+    sinon.stub(utils as any, 'addBitgoKeychain').resolves({ id: 'bitgo-key-id', commonKeychain: 'a'.repeat(128) });
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it('spreads walletId into round-1 payload when retrofit is provided', async () => {
+    const capturedPayloads: any[] = [];
+    sinon.stub(utils, 'sendKeyGenerationRound1').callsFake(async (_enterprise: string, payload: any) => {
+      capturedPayloads.push(payload);
+      // Return a bad bitgoMsg1 to short-circuit the ceremony after R1 capture
+      return { sessionId: sessionId as any, bitgoMsg1: { message: '', signature: '' } as any };
+    });
+
+    const retrofit = {
+      decryptedUserKey: JSON.stringify(userSigningMaterial),
+      decryptedBackupKey: JSON.stringify(backupSigningMaterial),
+      walletId,
+    };
+
+    await assert.rejects(
+      () => utils.createKeychains({ passphrase: 'test', enterprise, retrofit }),
+      () => true
+    );
+
+    assert.strictEqual(capturedPayloads.length, 1, 'sendKeyGenerationRound1 should be called once');
+    assert.strictEqual(capturedPayloads[0].walletId, walletId, 'walletId must be present in round-1 payload');
+  });
+
+  it('omits walletId from round-1 payload when retrofit is absent', async () => {
+    const capturedPayloads: any[] = [];
+    sinon.stub(utils, 'sendKeyGenerationRound1').callsFake(async (_enterprise: string, payload: any) => {
+      capturedPayloads.push(payload);
+      return { sessionId: sessionId as any, bitgoMsg1: { message: '', signature: '' } as any };
+    });
+
+    await assert.rejects(
+      () => utils.createKeychains({ passphrase: 'test', enterprise }),
+      () => true
+    );
+
+    assert.strictEqual(capturedPayloads.length, 1, 'sendKeyGenerationRound1 should be called once');
+    assert.strictEqual(capturedPayloads[0].walletId, undefined, 'walletId must be absent when no retrofit');
+  });
+});