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 8066fd89dcb..c213a0b9b4a 100644 --- a/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dkg.ts +++ b/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dkg.ts @@ -77,6 +77,30 @@ export class Dkg { } } + private _freeSession(): void { + const session = this.dkgSession; + this.dkgSession = undefined; + session?.free(); + } + + private _cleanup(resources: { free(): void }[], failed: boolean): void { + let cleanupError: unknown; + let cleanupFailed = false; + for (const resource of resources) { + try { + resource.free(); + } catch (e) { + if (!cleanupFailed) { + cleanupError = e; + cleanupFailed = true; + } + } + } + if (!failed && cleanupFailed) { + throw cleanupError; + } + } + private _createDKLsRetrofitKeyShare() { if (this.retrofitData) { const dklsKeyShare = buildDklsKeyShare({ @@ -138,30 +162,39 @@ export class Dkg { const initDkls = await import('@silencelaboratories/dkls-wasm-ll-web'); await initDkls.default(); } - this._createDKLsRetrofitKeyShare(); if (this.seed && this.seed.length !== 32) { throw Error(`Seed should be 32 bytes, got ${this.seed.length}.`); } const { KeygenSession } = this.getDklsWasm(); - if (this.dklsKeyShareRetrofitObject) { - this.dkgSession = this.seed - ? KeygenSession.initKeyRotation(this.dklsKeyShareRetrofitObject, new Uint8Array(this.seed)) - : KeygenSession.initKeyRotation(this.dklsKeyShareRetrofitObject); - } else { - this.dkgSession = this.seed - ? new KeygenSession(this.n, this.t, this.partyIdx, new Uint8Array(this.seed)) - : new KeygenSession(this.n, this.t, this.partyIdx); - } + let firstMessage: Message | undefined; + let failed = false; try { - const payload = this.dkgSession.createFirstMessage().payload; + this._createDKLsRetrofitKeyShare(); + if (this.dklsKeyShareRetrofitObject) { + this.dkgSession = this.seed + ? KeygenSession.initKeyRotation(this.dklsKeyShareRetrofitObject, new Uint8Array(this.seed)) + : KeygenSession.initKeyRotation(this.dklsKeyShareRetrofitObject); + } else { + this.dkgSession = this.seed + ? new KeygenSession(this.n, this.t, this.partyIdx, new Uint8Array(this.seed)) + : new KeygenSession(this.n, this.t, this.partyIdx); + } + firstMessage = this.dkgSession.createFirstMessage(); + const payload = firstMessage.payload; this.dkgSessionBytes = this.dkgSession.toBytes(); this._deserializeState(); - return { - payload: payload, - from: this.partyIdx, - }; + return { payload, from: this.partyIdx }; } catch (e) { + failed = true; throw Error(`Error while creating the first message from party ${this.partyIdx}: ${e}`); + } finally { + // Key rotation borrows the old share; only keyshare() consumes a session. + const retrofit = this.dklsKeyShareRetrofitObject; + this.dklsKeyShareRetrofitObject = undefined; + this._cleanup( + [{ free: () => this._freeSession() }, ...(firstMessage ? [firstMessage] : []), ...(retrofit ? [retrofit] : [])], + failed + ); } } @@ -201,6 +234,8 @@ export class Dkg { throw Error('Session not initialized'); } const { Message } = this.getDklsWasm(); + let failed = false; + let keyShare: Keyshare | undefined; try { if (this.dkgState === DkgState.Round3) { const commitmentsUnsorted = messagesForIthRound.p2pMessages @@ -228,12 +263,11 @@ export class Dkg { ); } if (this.dkgState === DkgState.Round4) { - this.dkgKeyShare = this.dkgSession.keyshare(); - this.keyShareBuff = Buffer.from(this.dkgKeyShare.toBytes()); - this.dkgKeyShare.free(); - if (this.dklsKeyShareRetrofitObject) { - this.dklsKeyShareRetrofitObject.free(); - } + const session = this.dkgSession; + // The binding consumes the session even when keyshare() throws. + this.dkgSession = undefined; + keyShare = session.keyshare(); + this.keyShareBuff = Buffer.from(keyShare.toBytes()); this.dkgState = DkgState.Complete; return { broadcastMessages: [], p2pMessages: [] }; } else { @@ -265,15 +299,15 @@ export class Dkg { return broadcastReturn; }), }; + this.dkgSessionBytes = this.dkgSession.toBytes(); } catch (e) { + failed = true; throw Error(`Error while creating messages from party ${this.partyIdx}, round ${this.dkgState}: ${e}`); } finally { - nextRoundMessages.forEach((m) => m.free()); - // Session is freed when keyshare is called. - if (this.dkgState !== DkgState.Complete) { - this.dkgSessionBytes = this.dkgSession.toBytes(); - this.dkgSession = undefined; - } + this._cleanup( + [{ free: () => this._freeSession() }, ...nextRoundMessages, ...(keyShare ? [keyShare] : [])], + failed + ); } return nextRoundDeserializedMessages; } @@ -336,11 +370,19 @@ export class Dkg { dkg.keyShareBuff = sessionData.keyShareBuff; } - 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; + let failed = false; + try { + 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; + } catch (e) { + failed = true; + throw e; + } finally { + dkg._cleanup([{ free: () => dkg._freeSession() }], failed); + } } } diff --git a/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dsg.ts b/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dsg.ts index 130778399d0..c9005066690 100644 --- a/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dsg.ts +++ b/modules/sdk-lib-mpc/src/tss/ecdsa-dkls/dsg.ts @@ -39,6 +39,17 @@ export class Dsg { } } + private releaseSession(preserveError = false): void { + const session = this.dsgSession; + this.dsgSession = undefined; + try { + session?.free(); + } catch (e) { + // A native trap can also make destruction fail; retain the original protocol error. + if (!preserveError) throw e; + } + } + private _deserializeState() { if (!this.dsgSession) { throw Error('Session not intialized'); @@ -138,17 +149,34 @@ export class Dsg { throw Error(`Party index: ${this.partyIdx} does not match key share partyId: ${keyShare.partyId} `); } this.dsgSession = new SignSessionOTVariant(keyShare, this.derivationPath); + let failed = false; try { - const payload = this.dsgSession.createFirstMessage().payload; + const firstMessage = this.dsgSession.createFirstMessage(); + let payload: Uint8Array; + let payloadFailed = false; + try { + payload = firstMessage.payload; + } catch (e) { + payloadFailed = true; + throw e; + } finally { + try { + firstMessage.free(); + } catch (e) { + if (!payloadFailed) throw e; + } + } this._deserializeState(); this.dsgSessionBytes = this.dsgSession.toBytes(); - this.dsgSession = undefined; return { payload: payload, from: this.partyIdx, }; } catch (e) { + failed = true; throw Error(`Error while creating the first message from party ${this.partyIdx}: ${e}`); + } finally { + this.releaseSession(failed); } } @@ -166,9 +194,7 @@ export class Dsg { if (this._signature) { throw new Error('Session already ended because combined signature was produced.'); } - if (this.dsgSession) { - this.dsgSession.free(); - } + this.releaseSession(); this.dsgState = DsgState.Uninitialized; } @@ -186,6 +212,7 @@ export class Dsg { throw Error('Session not initialized'); } const { Message } = this.getDklsWasm(); + let failed = false; try { if (this.dsgState === DsgState.Round4) { this.dsgState = DsgState.Complete; @@ -246,6 +273,7 @@ export class Dsg { }), }; } catch (e) { + failed = true; if (e.message.startsWith('Abort the protocol and ban')) { throw Error( 'Signing aborted. Please stop all transaction signing from this wallet and contact support@bitgo.com.' @@ -253,11 +281,23 @@ export class Dsg { } throw Error(`Error while creating messages from party ${this.partyIdx}, round ${this.dsgState}: ${e}`); } finally { - nextRoundMessages.forEach((m) => m.free()); - // Session is freed when combine is called. - if (this.dsgState !== DsgState.Complete) { - this.dsgSessionBytes = this.dsgSession.toBytes(); - this.dsgSession = undefined; + try { + nextRoundMessages.forEach((m) => m.free()); + if (!failed && this.dsgState !== DsgState.Complete) { + this.dsgSessionBytes = this.dsgSession.toBytes(); + } + } catch (e) { + if (!failed) { + failed = true; + throw e; + } + } finally { + // combine consumes the native session, including on failure. + if (this.dsgState === DsgState.Complete) { + this.dsgSession = undefined; + } else { + this.releaseSession(failed); + } } } return nextRoundDeserializedMessages; 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 ae373360021..92ec8acf3cc 100644 --- a/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDkg.ts +++ b/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDkg.ts @@ -1,4 +1,6 @@ import assert from 'assert'; +import sinon from 'sinon'; +import * as wasm from '@silencelaboratories/dkls-wasm-ll-node'; import { DklsDkg, DklsTypes } from '../../../../src/tss/ecdsa-dkls'; import { isRight } from 'fp-ts/Either'; import { @@ -678,3 +680,94 @@ describe('DKLS Dkg 2x3', function () { assert.deepEqual(DklsTypes.getCommonKeychain(backupKeyShare), DklsTypes.getCommonKeychain(bitgoKeyShare)); }); }); + +describe('DKLS key generation ownership', function () { + afterEach(function () { + sinon.restore(); + }); + + it('releases initialization sessions and messages before an external handoff', async function () { + const freeSession = sinon.spy(wasm.KeygenSession.prototype, 'free'); + const freeMessage = sinon.spy(wasm.Message.prototype, 'free'); + const dkg = new DklsDkg.Dkg(3, 2, 0); + await dkg.initDkg(); + assert.strictEqual(freeSession.callCount, 1); + assert.strictEqual(freeMessage.callCount, 1); + assert.strictEqual(dkg['dkgSession'], undefined); + const restored = await DklsDkg.Dkg.restoreSession(3, 2, 0, dkg.getSessionData()); + assert.strictEqual(freeSession.callCount, 2); + assert.strictEqual(restored['dkgSession'], undefined); + assert.deepStrictEqual(restored.getSessionData(), dkg.getSessionData()); + }); + + it('balances session ownership across a complete three-party workflow', async function () { + const freeSession = sinon.spy(wasm.KeygenSession.prototype, 'free'); + const consumeSession = sinon.spy(wasm.KeygenSession.prototype, 'keyshare'); + const parties = await generateDKGKeyShares(); + assert.strictEqual(freeSession.callCount, 12); + assert.strictEqual(consumeSession.callCount, 3); + for (const party of parties) { + assert(party.getKeyShare().length > 0); + assert.strictEqual(party['dkgSession'], undefined); + } + }); + + it('preserves initialization errors when disposal also fails and detaches first', async function () { + const dkg = new DklsDkg.Dkg(3, 2, 0); + const originalFree = wasm.KeygenSession.prototype.free; + sinon.stub(wasm.KeygenSession.prototype, 'createFirstMessage').throws(new Error('first message failure')); + sinon.stub(wasm.KeygenSession.prototype, 'free').callsFake(function (this: wasm.KeygenSession) { + assert.strictEqual(dkg['dkgSession'], undefined); + originalFree.call(this); + throw new Error('cleanup failure'); + }); + await assert.rejects(dkg.initDkg(), /first message failure/); + }); + + it('does not serialize a failed round or mask its error with cleanup', async function () { + const dkg = new DklsDkg.Dkg(3, 2, 0); + await dkg.initDkg(); + const originalFree = wasm.KeygenSession.prototype.free; + const serialize = sinon.spy(wasm.KeygenSession.prototype, 'toBytes'); + sinon.stub(wasm.KeygenSession.prototype, 'handleMessages').throws(new Error('round failure')); + sinon.stub(wasm.KeygenSession.prototype, 'free').callsFake(function (this: wasm.KeygenSession) { + assert.strictEqual(dkg['dkgSession'], undefined); + originalFree.call(this); + throw new Error('cleanup failure'); + }); + assert.throws(() => dkg.handleIncomingMessages({ broadcastMessages: [], p2pMessages: [] }), /round failure/); + assert.strictEqual(serialize.callCount, 0); + }); + + it('releases borrowed retrofit shares when rotation initialization fails', async function () { + const dkg = new DklsDkg.Dkg(3, 2, 0, undefined, { xShare: fixtures.mockDKeyShare.xShare }); + const freeShare = sinon.spy(wasm.Keyshare.prototype, 'free'); + sinon.stub(wasm.KeygenSession, 'initKeyRotation').throws(new Error('rotation failure')); + await assert.rejects(dkg.initDkg(), /rotation failure/); + assert.strictEqual(freeShare.callCount, 1); + assert.strictEqual(dkg['dklsKeyShareRetrofitObject'], undefined); + assert.strictEqual(dkg['dkgSession'], undefined); + }); + + it('releases initialization resources if serialization fails', async function () { + const freeSession = sinon.spy(wasm.KeygenSession.prototype, 'free'); + const freeMessage = sinon.spy(wasm.Message.prototype, 'free'); + sinon.stub(wasm.KeygenSession.prototype, 'toBytes').throws(new Error('serialization failure')); + const dkg = new DklsDkg.Dkg(3, 2, 0); + await assert.rejects(dkg.initDkg(), /serialization failure/); + assert.strictEqual(freeSession.callCount, 1); + assert.strictEqual(freeMessage.callCount, 1); + assert.strictEqual(dkg['dkgSession'], undefined); + }); + + it('does not free a session consumed by a throwing keyshare call', async function () { + const dkg = new DklsDkg.Dkg(3, 2, 0); + await dkg.initDkg(); + dkg['dkgState'] = DklsTypes.DkgState.Round4; + sinon.stub(wasm.KeygenSession.prototype, 'handleMessages').returns([]); + const freeSession = sinon.spy(wasm.KeygenSession.prototype, 'free'); + assert.throws(() => dkg.handleIncomingMessages({ broadcastMessages: [], p2pMessages: [] })); + assert.strictEqual(freeSession.callCount, 0); + assert.strictEqual(dkg['dkgSession'], undefined); + }); +}); diff --git a/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDsg.ts b/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDsg.ts index 95910cd16b6..b2642e53388 100644 --- a/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDsg.ts +++ b/modules/sdk-lib-mpc/test/unit/tss/ecdsa/dklsDsg.ts @@ -3,8 +3,9 @@ import * as fs from 'fs'; import * as crypto from 'crypto'; import should from 'should'; import assert from 'assert'; +import sinon from 'sinon'; -import { Keyshare } from '@silencelaboratories/dkls-wasm-ll-node'; +import { Keyshare, Message, SignSessionOTVariant } from '@silencelaboratories/dkls-wasm-ll-node'; import { decode } from 'cbor-x'; import * as mpcv2KeyCardData from './fixtures/mpcv2keycarddata'; import * as sjcl from '@bitgo/sjcl'; @@ -55,6 +56,87 @@ describe('DKLS Dsg 2x3', function () { `${__dirname}/fixtures/backupShare`, `${__dirname}/fixtures/bitgoShare`, ]; + for (const payloadFails of [true, false]) { + it(`preserves ${payloadFails ? 'payload' : 'cleanup'} failure during first-message disposal`, async function () { + const sandbox = sinon.createSandbox(); + const originalFree = Message.prototype.free; + const sessionFree = sandbox.spy(SignSessionOTVariant.prototype, 'free'); + const messageFree = sandbox.stub(Message.prototype, 'free').callsFake(function (this: Message) { + originalFree.call(this); + throw new Error('message cleanup failure'); + }); + if (payloadFails) { + sandbox.stub(Message.prototype, 'payload').get(() => { + throw new Error('payload read failure'); + }); + } + const party = new DklsDsg.Dsg(fs.readFileSync(shareFiles[0]), 0, 'm', Buffer.alloc(32)); + try { + await assert.rejects(party.init(), payloadFails ? /payload read failure/ : /message cleanup failure/); + assert.strictEqual(messageFree.callCount, 1); + assert.strictEqual(sessionFree.callCount, 1); + assert.strictEqual(party['dsgSession'], undefined); + } finally { + sandbox.restore(); + } + }); + } + + it('preserves the protocol error and detaches the session when cleanup also fails', async function () { + class InspectableDsg extends DklsDsg.Dsg { + hasLiveSession(): boolean { + return this.dsgSession !== undefined; + } + } + const party = new InspectableDsg(fs.readFileSync(shareFiles[0]), 0, 'm', Buffer.alloc(32)); + await party.init(); + const originalHandle = SignSessionOTVariant.prototype.handleMessages; + const originalFree = SignSessionOTVariant.prototype.free; + let freed = false; + SignSessionOTVariant.prototype.handleMessages = function () { + throw new Error('primary protocol failure'); + }; + SignSessionOTVariant.prototype.free = function () { + originalFree.call(this); + freed = true; + throw new Error('secondary cleanup failure'); + }; + try { + assert.throws( + () => party.handleIncomingMessages({ broadcastMessages: [], p2pMessages: [] }), + /primary protocol failure/ + ); + assert.strictEqual(freed, true); + assert.strictEqual(party.hasLiveSession(), false); + } finally { + SignSessionOTVariant.prototype.handleMessages = originalHandle; + SignSessionOTVariant.prototype.free = originalFree; + } + }); + + it('frees serialized sessions after every signing round', async function () { + const originalFree = SignSessionOTVariant.prototype.free; + let freed = 0; + SignSessionOTVariant.prototype.free = function () { + freed++; + return originalFree.call(this); + }; + try { + const parties = [0, 1].map( + (party) => new DklsDsg.Dsg(fs.readFileSync(shareFiles[party]), party, 'm', Buffer.alloc(32)) + ); + const messages = await executeTillRound(4, parties[0], parties[1]); + assert.strictEqual(freed, 8); + parties.forEach((party, index) => { + party.handleIncomingMessages({ p2pMessages: [], broadcastMessages: messages[1 - index].broadcastMessages }); + }); + assert.deepStrictEqual(parties[0].signature, parties[1].signature); + assert.strictEqual(freed, 8); + } finally { + SignSessionOTVariant.prototype.free = originalFree; + } + }); + vectors.forEach(async function (vector) { it(`should create signatures for parties ${vector.party1} and ${vector.party2} with derivation`, async function () { const party1 = new DklsDsg.Dsg(