Merge pull request #4690 from BitGo/WP-2166-implement-session-restore…

…-and-get-for-dsg

feat(sdk-lib-mpc): implement session get and restore for dsg

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

@@ -49,6 +49,41 @@ export class Dsg {
     }
   }
 
+  /**
+   * Returns the current DSG session as a base64 string.
+   * @returns {string} - base64 string of the current DSG session
+   */
+  getSession(): string {
+    return Buffer.from(this.dsgSessionBytes).toString('base64');
+  }
+
+  /**
+   * Sets the DSG session from a base64 string.
+   * @param {string} session - base64 string of the DSG session
+   */
+  setSession(session: string): void {
+    this.dsgSession = undefined;
+    const sessionBytes = new Uint8Array(Buffer.from(session, 'base64'));
+    const round = decode(sessionBytes).round;
+    switch (true) {
+      case round === 'WaitMsg1':
+        this.dsgState = DsgState.Round1;
+        break;
+      case round === 'WaitMsg2':
+        this.dsgState = DsgState.Round2;
+        break;
+      case round === 'WaitMsg3':
+        this.dsgState = DsgState.Round3;
+        break;
+      case 'WaitMsg4' in round:
+        this.dsgState = DsgState.Round4;
+        break;
+      default:
+        throw Error(`Invalid State: ${round}`);
+    }
+    this.dsgSessionBytes = sessionBytes;
+  }
+
   async init(): Promise<DeserializedBroadcastMessage> {
     if (this.dsgState !== DsgState.Uninitialized) {
       throw Error('DSG session already initialized');
@@ -65,6 +100,8 @@ export class Dsg {
     try {
       const payload = this.dsgSession.createFirstMessage().payload;
       this._deserializeState();
+      this.dsgSessionBytes = this.dsgSession.toBytes();
+      this.dsgSession = undefined;
       return {
         payload: payload,
         from: this.partyIdx,

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

@@ -2,11 +2,20 @@ import { DklsDsg, DklsUtils } from '../../../../src/tss/ecdsa-dkls';
 import * as fs from 'fs';
 import * as crypto from 'crypto';
 import should from 'should';
+import assert from 'assert';
+
 import { Keyshare } from '@silencelaboratories/dkls-wasm-ll-node';
 import { decode } from 'cbor-x';
 import * as mpcv2KeyCardData from './fixtures/mpcv2keycarddata';
 import * as sjcl from 'sjcl';
-import { getDecodedReducedKeyShare, ReducedKeyShare, RetrofitData } from '../../../../src/tss/ecdsa-dkls/types';
+import {
+  DeserializedBroadcastMessage,
+  DeserializedDklsSignature,
+  DeserializedMessages,
+  getDecodedReducedKeyShare,
+  ReducedKeyShare,
+  RetrofitData,
+} from '../../../../src/tss/ecdsa-dkls/types';
 import {
   executeTillRound,
   generate2of2KeyShares,
@@ -255,4 +264,149 @@ describe('DKLS Dsg 2x3', function () {
       Error('Error while creating messages from party 0, round 5: Error: k256 error: signature error')
     );
   });
+
+  it(`should get and set sesssion corretly for all rounds`, async function () {
+    const vector = vectors[0];
+
+    const party2 = new DklsDsg.Dsg(
+      fs.readFileSync(shareFiles[vector.party2]),
+      vector.party2,
+      vector.derivationPath,
+      crypto.createHash('sha256').update(Buffer.from(vector.msgToSign, 'hex')).digest()
+    );
+
+    // round 1
+    let party1Round1Message: DeserializedBroadcastMessage,
+      party2Round1Message: DeserializedBroadcastMessage,
+      party1Round1Session: string;
+    {
+      const party1 = new DklsDsg.Dsg(
+        fs.readFileSync(shareFiles[vector.party1]),
+        vector.party1,
+        vector.derivationPath,
+        crypto.createHash('sha256').update(Buffer.from(vector.msgToSign, 'hex')).digest()
+      );
+      party1Round1Message = await party1.init();
+      party1Round1Session = party1.getSession();
+
+      party2Round1Message = await party2.init();
+    }
+
+    // round 2
+    let party1Round2Messages: DeserializedMessages,
+      party2Round2Messages: DeserializedMessages,
+      party1Round2Session: string;
+    {
+      const party1Round2DSG = new DklsDsg.Dsg(
+        fs.readFileSync(shareFiles[vector.party1]),
+        vector.party1,
+        vector.derivationPath,
+        crypto.createHash('sha256').update(Buffer.from(vector.msgToSign, 'hex')).digest()
+      );
+      party1Round2DSG.setSession(party1Round1Session);
+      party1Round2Messages = party1Round2DSG.handleIncomingMessages({
+        p2pMessages: [],
+        broadcastMessages: [party2Round1Message],
+      });
+      party1Round2Session = party1Round2DSG.getSession();
+
+      party2Round2Messages = party2.handleIncomingMessages({
+        p2pMessages: [],
+        broadcastMessages: [party1Round1Message],
+      });
+    }
+
+    // round 3
+    let party1Round3Messages: DeserializedMessages,
+      party2Round3Messages: DeserializedMessages,
+      party1Round3Session: string;
+    {
+      const party1Round3DSG = new DklsDsg.Dsg(
+        fs.readFileSync(shareFiles[vector.party1]),
+        vector.party1,
+        vector.derivationPath,
+        crypto.createHash('sha256').update(Buffer.from(vector.msgToSign, 'hex')).digest()
+      );
+      party1Round3DSG.setSession(party1Round2Session);
+      party1Round3Messages = party1Round3DSG.handleIncomingMessages({
+        p2pMessages: party2Round2Messages.p2pMessages,
+        broadcastMessages: [],
+      });
+      party1Round3Session = party1Round3DSG.getSession();
+
+      party2Round3Messages = party2.handleIncomingMessages({
+        p2pMessages: party1Round2Messages.p2pMessages,
+        broadcastMessages: [],
+      });
+    }
+
+    // round 4
+    let party1Round4Messages: DeserializedMessages,
+      party2Round4Messages: DeserializedMessages,
+      party1Round4Session: string;
+    {
+      const party1Round4DSG = new DklsDsg.Dsg(
+        fs.readFileSync(shareFiles[vector.party1]),
+        vector.party1,
+        vector.derivationPath,
+        crypto.createHash('sha256').update(Buffer.from(vector.msgToSign, 'hex')).digest()
+      );
+      party1Round4DSG.setSession(party1Round3Session);
+      party1Round4Messages = party1Round4DSG.handleIncomingMessages({
+        p2pMessages: party2Round3Messages.p2pMessages,
+        broadcastMessages: [],
+      });
+      party1Round4Session = party1Round4DSG.getSession();
+
+      party2Round4Messages = party2.handleIncomingMessages({
+        p2pMessages: party1Round3Messages.p2pMessages,
+        broadcastMessages: [],
+      });
+    }
+
+    // round 5
+    let signature: DeserializedDklsSignature;
+    {
+      const party1Round5DSG = new DklsDsg.Dsg(
+        fs.readFileSync(shareFiles[vector.party1]),
+        vector.party1,
+        vector.derivationPath,
+        crypto.createHash('sha256').update(Buffer.from(vector.msgToSign, 'hex')).digest()
+      );
+      party1Round5DSG.setSession(party1Round4Session);
+      party1Round5DSG.handleIncomingMessages({
+        p2pMessages: [],
+        broadcastMessages: party2Round4Messages.broadcastMessages,
+      });
+      party2.handleIncomingMessages({
+        p2pMessages: [],
+        broadcastMessages: party1Round4Messages.broadcastMessages,
+      });
+      party1Round5DSG.signature.should.deepEqual(party2.signature);
+      signature = party1Round5DSG.signature;
+    }
+
+    assert(party1Round4Messages.broadcastMessages[0].signatureR);
+    const combinedSigUsingUtil = DklsUtils.combinePartialSignatures(
+      [party1Round4Messages.broadcastMessages[0].payload, party2Round4Messages.broadcastMessages[0].payload],
+      Buffer.from(party1Round4Messages.broadcastMessages[0].signatureR as Uint8Array).toString('hex')
+    );
+    (
+      (combinedSigUsingUtil.R.every((p) => signature.R.includes(p)) &&
+        signature.R.every((p) => combinedSigUsingUtil.R.includes(p))) ||
+      (signature.S.every((p) => combinedSigUsingUtil.S.includes(p)) &&
+        combinedSigUsingUtil.S.every((p) => signature.S.includes(p)))
+    ).should.equal(true);
+
+    const keyShare: Keyshare = Keyshare.fromBytes(fs.readFileSync(shareFiles[vector.party1]));
+    const convertedSignature = verifyAndConvertDklsSignature(
+      Buffer.from(vector.msgToSign, 'hex'),
+      signature,
+      Buffer.from(keyShare.publicKey).toString('hex') +
+        Buffer.from(decode(keyShare.toBytes()).root_chain_code).toString('hex'),
+      vector.derivationPath
+    );
+    should.exist(convertedSignature);
+    convertedSignature.split(':').length.should.equal(4);
+  });
 });