fix(sdk-lib-mpc): free DKLS signing and keygen sessions (WCN-4341)

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);
+    }
   }
 }

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,18 +273,31 @@ 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.'
         );
       }
       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;

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);
+  });
+});