diff --git a/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dkg.ts b/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dkg.ts index 11ee916648..f20d71bdb9 100644 --- a/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dkg.ts +++ b/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; } } diff --git a/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDkg.ts b/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDkg.ts index 99561f3a53..252fcf38a7 100644 --- a/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDkg.ts +++ b/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);