Merge pull request #8466 from BitGo/worktree-WAL-396-fix-dkg-restore-session

fix(sdk-lib-mpc): re-derive dkgState from WASM bytes in restoreSession

Author: mohammadalfaiyazbitgo

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

@@ -119,7 +119,8 @@ export class Dkg {
         this.dkgState = DkgState.Round3;
         break;
       case 'WaitMsg4':
-        this.dkgState = DkgState.Round4;
+        // keyShareBuff present means keyshare() already ran and freed the session; bytes are frozen at WaitMsg4.
+        this.dkgState = this.keyShareBuff ? DkgState.Complete : DkgState.Round4;
         break;
       case 'Ended':
         this.dkgState = DkgState.Complete;
@@ -339,7 +340,6 @@ export class Dkg {
     }
 
     dkg.dkgSessionBytes = sessionData.dkgSessionBytes;
-    dkg.dkgState = sessionData.dkgState;
 
     if (sessionData.chainCodeCommitment) {
       dkg.chainCodeCommitment = sessionData.chainCodeCommitment;
@@ -350,6 +350,10 @@ export class Dkg {
     }
 
     dkg._restoreSession();
+    // Re-derive state from WASM bytes rather than trusting the caller-supplied dkgState.
+    // This prevents a tampered or corrupted dkgState from causing handleIncomingMessages()
+    // to take the wrong branch (e.g. skipping chain code commitment or calling keyshare() prematurely).
+    dkg._deserializeState();
     return dkg;
   }
 }

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

@@ -283,6 +283,47 @@ describe('DKLS Dkg 2x3', function () {
     assert.deepEqual(DklsTypes.getCommonKeychain(backupKeyShare), DklsTypes.getCommonKeychain(bitgoKeyShare));
   });
 
+  it('restoreSession() should ignore tampered dkgState and re-derive from WASM bytes', async function () {
+    const user = new DklsDkg.Dkg(3, 2, 0);
+
+    // After initDkg() the WASM session encodes WaitMsg1 → DkgState.Round1
+    await user.initDkg();
+
+    const legitimateSessionData = user.getSessionData();
+
+    // Tamper: claim the session is at Round4 when WASM bytes still say Round1
+    const tamperedSessionData = {
+      ...legitimateSessionData,
+      dkgState: DklsTypes.DkgState.Round4,
+    };
+
+    const restoredUser = await DklsDkg.Dkg.restoreSession(3, 2, 0, tamperedSessionData);
+
+    // Must reflect the actual WASM state (Round1), not the tampered Round4
+    assert.strictEqual(
+      restoredUser['dkgState'],
+      DklsTypes.DkgState.Round1,
+      'restoreSession() must re-derive dkgState from WASM bytes and ignore caller-supplied value'
+    );
+  });
+
+  it('restoreSession() should restore a completed DKG session as DkgState.Complete', async function () {
+    const [user] = await generateDKGKeyShares();
+    const completedSessionData = user.getSessionData();
+
+    // dkgSessionBytes holds { round: 'Ended' }; restoreSession() must decode it as Complete
+    // without reconstructing the (already freed) WASM session
+    const restoredUser = await DklsDkg.Dkg.restoreSession(3, 2, 0, completedSessionData);
+
+    assert.strictEqual(
+      restoredUser['dkgState'],
+      DklsTypes.DkgState.Complete,
+      'restoreSession() must decode "Ended" round marker as DkgState.Complete'
+    );
+    // Key share must still be accessible on the restored instance
+    assert.ok(restoredUser.getKeyShare(), 'Key share should be accessible after restoring completed session');
+  });
+
   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);