Merge pull request #6368 from BitGo/WP-5188/store-restore-dkls-dkg-session

feat: store and restore ecdsa dkg sessions

Author: mohammadalfaiyazbitgo

modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dkg.ts

@@ -10,6 +10,13 @@ type BundlerWasmer = typeof import('@silencelaboratories/dkls-wasm-ll-bundler');
 
 type DklsWasm = NodeWasmer | WebWasmer | BundlerWasmer;
 
+export interface DkgSessionData {
+  dkgSessionBytes: Uint8Array;
+  dkgState: DkgState;
+  chainCodeCommitment?: Uint8Array;
+  keyShareBuff?: Buffer;
+}
+
 export class Dkg {
   protected dkgSession: KeygenSession | undefined;
   protected dkgSessionBytes: Uint8Array;
@@ -159,6 +166,7 @@ export class Dkg {
     }
     try {
       const payload = this.dkgSession.createFirstMessage().payload;
+      this.dkgSessionBytes = this.dkgSession.toBytes();
       this._deserializeState();
       return {
         payload: payload,
@@ -276,4 +284,67 @@ export class Dkg {
     }
     return nextRoundDeserializedMessages;
   }
+
+  /**
+   * Get the current session data that can be used to restore the session later
+   * @returns The current session data
+   */
+  getSessionData(): DkgSessionData {
+    const sessionData: DkgSessionData = {
+      dkgSessionBytes: this.dkgSessionBytes,
+      dkgState: this.dkgState,
+    };
+
+    if (this.chainCodeCommitment) {
+      sessionData.chainCodeCommitment = this.chainCodeCommitment;
+    }
+
+    if (this.keyShareBuff) {
+      sessionData.keyShareBuff = this.keyShareBuff;
+    }
+
+    return sessionData;
+  }
+
+  /**
+   * Restore a DKG session from previous session data
+   * Note: This should not be used for Round 1 as that's the initialization phase
+   * @param n Number of parties
+   * @param t Threshold
+   * @param partyIdx Party index
+   * @param sessionData Previous session data
+   * @param seed Optional seed
+   * @param retrofitData Optional retrofit data
+   * @param dklsWasm Optional DKLS wasm instance
+   * @returns A new DKG instance with the restored session
+   */
+  static async restoreSession(
+    n: number,
+    t: number,
+    partyIdx: number,
+    sessionData: DkgSessionData,
+    seed?: Buffer,
+    retrofitData?: RetrofitData,
+    dklsWasm?: BundlerWasmer
+  ): Promise<Dkg> {
+    const dkg = new Dkg(n, t, partyIdx, seed, retrofitData, dklsWasm);
+
+    if (!dkg.dklsWasm) {
+      await dkg.loadDklsWasm();
+    }
+
+    dkg.dkgSessionBytes = sessionData.dkgSessionBytes;
+    dkg.dkgState = sessionData.dkgState;
+
+    if (sessionData.chainCodeCommitment) {
+      dkg.chainCodeCommitment = sessionData.chainCodeCommitment;
+    }
+
+    if (sessionData.keyShareBuff) {
+      dkg.keyShareBuff = sessionData.keyShareBuff;
+    }
+
+    dkg._restoreSession();
+    return dkg;
+  }
 }

modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDkg.ts

@@ -282,4 +282,210 @@ describe('DKLS Dkg 2x3', function () {
     assert.deepEqual(DklsTypes.getCommonKeychain(userKeyShare), DklsTypes.getCommonKeychain(bitgoKeyShare));
     assert.deepEqual(DklsTypes.getCommonKeychain(backupKeyShare), DklsTypes.getCommonKeychain(bitgoKeyShare));
   });
+
+  it('should successfully finish DKG using restored sessions', async function () {
+    const user = new DklsDkg.Dkg(3, 2, 0);
+    const backup = new DklsDkg.Dkg(3, 2, 1);
+    const bitgo = new DklsDkg.Dkg(3, 2, 2);
+
+    // Generate GPG keys for authenticated encryption
+    openpgp.config.rejectCurves = new Set();
+    const userKey = await openpgp.generateKey({
+      userIDs: [{ name: 'user', email: '[email protected]' }],
+      curve: 'secp256k1',
+    });
+    const bitgoKey = await openpgp.generateKey({
+      userIDs: [{ name: 'bitgo', email: '[email protected]' }],
+      curve: 'secp256k1',
+    });
+    const backupKey = await openpgp.generateKey({
+      userIDs: [{ name: 'backup', email: '[email protected]' }],
+      curve: 'secp256k1',
+    });
+
+    const userGpgPubKey: PartyGpgKey = { partyId: 0, gpgKey: userKey.publicKey };
+    const userGpgPrvKey: PartyGpgKey = { partyId: 0, gpgKey: userKey.privateKey };
+    const bitgoGpgPubKey: PartyGpgKey = { partyId: 2, gpgKey: bitgoKey.publicKey };
+    const bitgoGpgPrvKey: PartyGpgKey = { partyId: 2, gpgKey: bitgoKey.privateKey };
+    const backupGpgPubKey: PartyGpgKey = { partyId: 1, gpgKey: backupKey.publicKey };
+    const backupGpgPrvKey: PartyGpgKey = { partyId: 1, gpgKey: backupKey.privateKey };
+
+    // Initialize DKG and get first round messages
+    const userRound1Message = await user.initDkg();
+    const backupRound1Message = await backup.initDkg();
+    const bitgoRound1Message = await bitgo.initDkg();
+
+    // Process round 1 messages to advance to round 2
+    let serializedMessages = serializeMessages({
+      broadcastMessages: [userRound1Message, backupRound1Message],
+      p2pMessages: [],
+    });
+
+    let authEncMessages = await encryptAndAuthOutgoingMessages(
+      serializedMessages,
+      [bitgoGpgPubKey],
+      [userGpgPrvKey, backupGpgPrvKey]
+    );
+
+    const restoredRound1User = await DklsDkg.Dkg.restoreSession(3, 2, 0, user.getSessionData());
+    const restoredRound1Backup = await DklsDkg.Dkg.restoreSession(3, 2, 1, backup.getSessionData());
+    const restoredRound1Bitgo = await DklsDkg.Dkg.restoreSession(3, 2, 2, bitgo.getSessionData());
+
+    // Round 2
+    const userRound2Messages = restoredRound1User.handleIncomingMessages({
+      p2pMessages: [],
+      broadcastMessages: [bitgoRound1Message, backupRound1Message],
+    });
+    const userRound2Data = restoredRound1User.getSessionData();
+
+    const backupRound2Messages = restoredRound1Backup.handleIncomingMessages({
+      p2pMessages: [],
+      broadcastMessages: [userRound1Message, bitgoRound1Message],
+    });
+    const backupRound2Data = restoredRound1Backup.getSessionData();
+
+    const decryptedMessages = await decryptAndVerifyIncomingMessages(
+      authEncMessages,
+      [userGpgPubKey, backupGpgPubKey],
+      [bitgoGpgPrvKey]
+    );
+
+    const deserializedDecryptedMessages = deserializeMessages(decryptedMessages);
+    const bitgoRound2Messages = restoredRound1Bitgo.handleIncomingMessages(deserializedDecryptedMessages);
+    const bitgoRound2Data = restoredRound1Bitgo.getSessionData();
+
+    // Restore sessions for Round 3
+    const restoredRound2User = await DklsDkg.Dkg.restoreSession(3, 2, 0, userRound2Data);
+    const restoredRound2Backup = await DklsDkg.Dkg.restoreSession(3, 2, 1, backupRound2Data);
+    const restoredRound2Bitgo = await DklsDkg.Dkg.restoreSession(3, 2, 2, bitgoRound2Data);
+
+    // Round 3
+    const restoredUserRound3Messages = restoredRound2User.handleIncomingMessages({
+      p2pMessages: backupRound2Messages.p2pMessages
+        .filter((m) => m.to === 0)
+        .concat(bitgoRound2Messages.p2pMessages.filter((m) => m.to === 0)),
+      broadcastMessages: [],
+    });
+    const userRound3Data = restoredRound2User.getSessionData();
+
+    const restoredBackupRound3Messages = restoredRound2Backup.handleIncomingMessages({
+      p2pMessages: bitgoRound2Messages.p2pMessages
+        .filter((m) => m.to === 1)
+        .concat(userRound2Messages.p2pMessages.filter((m) => m.to === 1)),
+      broadcastMessages: [],
+    });
+    const backupRound3Data = restoredRound2Backup.getSessionData();
+
+    // Encrypt messages for bitgo
+    serializedMessages = serializeMessages({
+      broadcastMessages: [],
+      p2pMessages: userRound2Messages.p2pMessages
+        .filter((m) => m.to === 2)
+        .concat(backupRound2Messages.p2pMessages.filter((m) => m.to === 2)),
+    });
+
+    authEncMessages = await encryptAndAuthOutgoingMessages(
+      serializedMessages,
+      [bitgoGpgPubKey],
+      [userGpgPrvKey, backupGpgPrvKey]
+    );
+
+    const restoredBitgoRound3Messages = restoredRound2Bitgo.handleIncomingMessages(
+      deserializeMessages(
+        await decryptAndVerifyIncomingMessages(authEncMessages, [userGpgPubKey, backupGpgPubKey], [bitgoGpgPrvKey])
+      )
+    );
+    const bitgoRound3Data = restoredRound2Bitgo.getSessionData();
+
+    // Restore sessions for Round 4
+    const restoredRound3User = await DklsDkg.Dkg.restoreSession(3, 2, 0, userRound3Data);
+    const restoredRound3Backup = await DklsDkg.Dkg.restoreSession(3, 2, 1, backupRound3Data);
+    const restoredRound3Bitgo = await DklsDkg.Dkg.restoreSession(3, 2, 2, bitgoRound3Data);
+
+    // Round 4
+    const restoredUserRound4Messages = restoredRound3User.handleIncomingMessages({
+      p2pMessages: restoredBackupRound3Messages.p2pMessages
+        .filter((m) => m.to === 0)
+        .concat(restoredBitgoRound3Messages.p2pMessages.filter((m) => m.to === 0)),
+      broadcastMessages: [],
+    });
+    const userRound4Data = restoredRound3User.getSessionData();
+
+    const restoredBackupRound4Messages = restoredRound3Backup.handleIncomingMessages({
+      p2pMessages: restoredBitgoRound3Messages.p2pMessages
+        .filter((m) => m.to === 1)
+        .concat(restoredUserRound3Messages.p2pMessages.filter((m) => m.to === 1)),
+      broadcastMessages: [],
+    });
+    const backupRound4Data = restoredRound3Backup.getSessionData();
+
+    serializedMessages = serializeMessages({
+      broadcastMessages: [],
+      p2pMessages: restoredUserRound3Messages.p2pMessages
+        .filter((m) => m.to === 2)
+        .concat(restoredBackupRound3Messages.p2pMessages.filter((m) => m.to === 2)),
+    });
+
+    authEncMessages = await encryptAndAuthOutgoingMessages(
+      serializedMessages,
+      [bitgoGpgPubKey],
+      [userGpgPrvKey, backupGpgPrvKey]
+    );
+
+    const restoredBitgoRound4Messages = restoredRound3Bitgo.handleIncomingMessages(
+      deserializeMessages(
+        await decryptAndVerifyIncomingMessages(authEncMessages, [userGpgPubKey, backupGpgPubKey], [bitgoGpgPrvKey])
+      )
+    );
+    const bitgoRound4Data = restoredRound3Bitgo.getSessionData();
+
+    // Restore sessions for final messages
+    const restoredRound4User = await DklsDkg.Dkg.restoreSession(3, 2, 0, userRound4Data);
+    const restoredRound4Backup = await DklsDkg.Dkg.restoreSession(3, 2, 1, backupRound4Data);
+    const restoredRound4Bitgo = await DklsDkg.Dkg.restoreSession(3, 2, 2, bitgoRound4Data);
+
+    // Final messages
+    restoredRound4User.handleIncomingMessages({
+      p2pMessages: [],
+      broadcastMessages: restoredBitgoRound4Messages.broadcastMessages.concat(
+        restoredBackupRound4Messages.broadcastMessages
+      ),
+    });
+
+    serializedMessages = serializeMessages({
+      broadcastMessages: restoredBackupRound4Messages.broadcastMessages.concat(
+        restoredUserRound4Messages.broadcastMessages
+      ),
+      p2pMessages: [],
+    });
+
+    authEncMessages = await encryptAndAuthOutgoingMessages(
+      serializedMessages,
+      [bitgoGpgPubKey],
+      [userGpgPrvKey, backupGpgPrvKey]
+    );
+
+    restoredRound4Bitgo.handleIncomingMessages(
+      deserializeMessages(
+        await decryptAndVerifyIncomingMessages(authEncMessages, [userGpgPubKey, backupGpgPubKey], [bitgoGpgPrvKey])
+      )
+    );
+
+    restoredRound4Backup.handleIncomingMessages({
+      p2pMessages: [],
+      broadcastMessages: restoredBitgoRound4Messages.broadcastMessages.concat(
+        restoredUserRound4Messages.broadcastMessages
+      ),
+    });
+
+    // Verify key shares
+    const userKeyShare = restoredRound4User.getKeyShare();
+    const backupKeyShare = restoredRound4Backup.getKeyShare();
+    const bitgoKeyShare = restoredRound4Bitgo.getKeyShare();
+
+    assert.deepEqual(decode(userKeyShare).public_key, decode(bitgoKeyShare).public_key);
+    assert.deepEqual(decode(backupKeyShare).public_key, decode(bitgoKeyShare).public_key);
+    assert.deepEqual(DklsTypes.getCommonKeychain(userKeyShare), DklsTypes.getCommonKeychain(bitgoKeyShare));
+    assert.deepEqual(DklsTypes.getCommonKeychain(backupKeyShare), DklsTypes.getCommonKeychain(bitgoKeyShare));
+  });
 });