Merge pull request #83 from Web3Auth/feat/importRecoverTssKey

himanshuchawla009 · web-flow · commit 84a8bd71cd69 · 2024-01-05T11:42:48.000+05:30
feat: import recover tss key
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -6,6 +6,6 @@
   "eslint.format.enable": true,
   "eslint.lintTask.enable": true,
   "editor.codeActionsOnSave": {
-      "source.fixAll": true
+    "source.fixAll": "explicit"
   }
 }
diff --git a/src/interfaces.ts b/src/interfaces.ts
@@ -77,7 +77,7 @@ export type MPCKeyDetails = {
   tssPubKey?: TkeyPoint;
 };
 
-export type OauthLoginParams = SubVerifierDetailsParams | AggregateVerifierLoginParams;
+export type OauthLoginParams = (SubVerifierDetailsParams | AggregateVerifierLoginParams) & { importTssKey?: string };
 export type UserInfo = TorusVerifierResponse & LoginWindowResponse;
 
 export interface EnableMFAParams {
@@ -132,6 +132,11 @@ export interface IdTokenLoginParams {
    * Any additional parameter (key value pair) you'd like to pass to the login function.
    */
   additionalParams?: ExtraParams;
+
+  /**
+   * Key to import key into Tss during first time login.
+   */
+  importTssKey?: string;
 }
 
 export interface Web3AuthState {
diff --git a/src/mpcCoreKit.ts b/src/mpcCoreKit.ts
@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/member-ordering */
 import { BNString, encrypt, getPubKeyPoint, Point as TkeyPoint, SHARE_DELETED, ShareStore, StringifiedType } from "@tkey-mpc/common-types";
-import ThresholdKey, { CoreError } from "@tkey-mpc/core";
+import ThresholdKey, { CoreError, lagrangeInterpolation } from "@tkey-mpc/core";
 import { TorusServiceProvider } from "@tkey-mpc/service-provider-torus";
 import { ShareSerializationModule } from "@tkey-mpc/share-serialization";
 import { TorusStorageLayer } from "@tkey-mpc/storage-layer-torus";
@@ -207,6 +207,39 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
     return this.options.uxMode === UX_MODE.REDIRECT;
   }
 
+  // RecoverTssKey only valid for user that enable MFA where user has 2 type shares :
+  // TssShareType.DEVICE and TssShareType.RECOVERY
+  // if the factors key provided is the same type recovery will not works
+  public async _UNSAFE_recoverTssKey(factorKey: string[]) {
+    this.checkReady();
+    const factorKeyBN = new BN(factorKey[0], "hex");
+    const shareStore0 = await this.getFactorKeyMetadata(factorKeyBN);
+    await this.tKey.initialize({ withShare: shareStore0 });
+
+    this.tkey.privKey = new BN(factorKey[1], "hex");
+
+    const tssShares: BN[] = [];
+    const tssIndexes: number[] = [];
+    const tssIndexesBN: BN[] = [];
+    for (let i = 0; i < factorKey.length; i++) {
+      const factorKeyBNInput = new BN(factorKey[i], "hex");
+      const { tssIndex, tssShare } = await this.tKey.getTSSShare(factorKeyBNInput);
+      if (tssIndexes.includes(tssIndex)) {
+        // reset instance before throw error
+        await this.init();
+        throw new Error("Duplicate TSS Index");
+      }
+      tssIndexes.push(tssIndex);
+      tssIndexesBN.push(new BN(tssIndex));
+      tssShares.push(tssShare);
+    }
+
+    const finalKey = lagrangeInterpolation(tssShares, tssIndexesBN);
+    // reset instance after recovery completed
+    await this.init();
+    return finalKey.toString("hex", 64);
+  }
+
   public async init(params: InitParams = { handleRedirectResult: true }): Promise<void> {
     this.resetState();
     if (params.rehydrate === undefined) params.rehydrate = true;
@@ -274,7 +307,7 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
   public async loginWithOauth(params: OauthLoginParams): Promise<void> {
     this.checkReady();
     if (this.isNodejsOrRN(this.options.uxMode)) throw new Error(`Oauth login is NOT supported in ${this.options.uxMode}`);
-
+    const { importTssKey } = params;
     const tkeyServiceProvider = this.tKey.serviceProvider as TorusServiceProvider;
     try {
       // oAuth login.
@@ -307,7 +340,7 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
         });
       }
 
-      await this.setupTkey();
+      await this.setupTkey(importTssKey);
     } catch (err: unknown) {
       log.error("login error", err);
       if (err instanceof CoreError) {
@@ -319,7 +352,7 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
 
   public async loginWithJWT(idTokenLoginParams: IdTokenLoginParams): Promise<void> {
     this.checkReady();
-
+    const { importTssKey } = idTokenLoginParams;
     const { verifier, verifierId, idToken } = idTokenLoginParams;
     try {
       // oAuth login.
@@ -357,7 +390,7 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
         signatures: this._getSignatures(loginResponse.sessionData.sessionTokenData),
       });
 
-      await this.setupTkey();
+      await this.setupTkey(importTssKey);
     } catch (err: unknown) {
       log.error("login error", err);
       if (err instanceof CoreError) {
@@ -758,7 +791,7 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
     if (!this.state.signatures) throw new Error("signatures not present");
 
     const tssKeyBN = new BN(tssKey, "hex");
-    this.tKey.importTssKey({ tag: this.tKey.tssTag, importKey: tssKeyBN, factorPub, newTSSIndex }, { authSignatures: this.state.signatures });
+    await this.tKey.importTssKey({ tag: this.tKey.tssTag, importKey: tssKeyBN, factorPub, newTSSIndex }, { authSignatures: this.state.signatures });
   }
 
   public async _UNSAFE_exportTssKey(): Promise<string> {
@@ -780,7 +813,7 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
     return tssNonce;
   }
 
-  private async setupTkey(): Promise<void> {
+  private async setupTkey(importTssKey?: string): Promise<void> {
     if (!this.state.oAuthKey) {
       throw new Error("user not logged in");
     }
@@ -797,10 +830,15 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
       } else {
         factorKey = getHashedPrivateKey(this.state.oAuthKey, this.options.hashedFactorNonce);
       }
-      const deviceTSSShare = new BN(generatePrivate());
       const deviceTSSIndex = TssShareType.DEVICE;
       const factorPub = getPubKeyPoint(factorKey);
-      await this.tKey.initialize({ useTSS: true, factorPub, deviceTSSShare, deviceTSSIndex });
+      if (!importTssKey) {
+        const deviceTSSShare = new BN(generatePrivate());
+        await this.tKey.initialize({ useTSS: true, factorPub, deviceTSSShare, deviceTSSIndex });
+      } else {
+        await this.tKey.initialize();
+        await this.importTssKey(importTssKey, factorPub, deviceTSSIndex);
+      }
 
       // Finalize initialization.
       await this.tKey.reconstructKey();
@@ -814,6 +852,7 @@ export class Web3AuthMPCCoreKit implements ICoreKit {
         await this.addFactorDescription(factorKey, FactorKeyTypeShareDescription.HashedShare);
       }
     } else {
+      if (importTssKey) throw new Error("Cannot import tss key for existing user");
       await this.tKey.initialize({ neverInitializeNewKey: true });
       const hashedFactorKey = getHashedPrivateKey(this.state.oAuthKey, this.options.hashedFactorNonce);
       if ((await this.checkIfFactorKeyValid(hashedFactorKey)) && !this.options.disableHashedFactorKey) {
diff --git a/tests/importRecovery.spec.ts b/tests/importRecovery.spec.ts
@@ -0,0 +1,113 @@
+/* eslint-disable mocha/handle-done-callback */
+import assert from "node:assert";
+import test from "node:test";
+
+import * as TssLib from "@toruslabs/tss-lib-node";
+import { log } from "@web3auth/base";
+
+import { BrowserStorage, IdTokenLoginParams, TssShareType, WEB3AUTH_NETWORK, Web3AuthMPCCoreKit } from "../src";
+import { criticalResetAccount, mockLogin, newCoreKitLogInInstance } from "./setup";
+
+type ImportKeyTestVariable = {
+  manualSync?: boolean;
+  email: string;
+  importKeyEmail: string;
+};
+export const ImportTest = async (testVariable: ImportKeyTestVariable) => {
+  test(`import recover tss key : ${testVariable.manualSync}`, async function (t) {
+    t.before(async () => {
+      const instance = await newCoreKitLogInInstance({
+        network: WEB3AUTH_NETWORK.DEVNET,
+        manualSync: testVariable.manualSync,
+        email: testVariable.email,
+      });
+      await criticalResetAccount(instance);
+      await instance.logout();
+    });
+
+    await t.test("#recover Tss key using 2 factors key, import tss key to new oauth login", async function () {
+      const { idToken, parsedToken } = await mockLogin(testVariable.email);
+
+      const idTokenLoginParams = {
+        verifier: "torus-test-health",
+        verifierId: parsedToken.email,
+        idToken,
+      } as IdTokenLoginParams;
+
+      const coreKitInstance = new Web3AuthMPCCoreKit({
+        web3AuthClientId: "torus-key-test",
+        web3AuthNetwork: WEB3AUTH_NETWORK.DEVNET,
+        baseUrl: "http://localhost:3000",
+        uxMode: "nodejs",
+        tssLib: TssLib,
+        storageKey: "memory",
+        manualSync: testVariable.manualSync,
+      });
+
+      await coreKitInstance.init();
+      await coreKitInstance.loginWithJWT(idTokenLoginParams);
+
+      const factorKeyDevice = await coreKitInstance.createFactor({
+        shareType: TssShareType.DEVICE,
+      });
+
+      const factorKeyRecovery = await coreKitInstance.createFactor({
+        shareType: TssShareType.RECOVERY,
+      });
+
+      const exportedTssKey1 = await coreKitInstance._UNSAFE_exportTssKey();
+      // recover key
+      // reinitalize corekit
+      await coreKitInstance.logout();
+      BrowserStorage.getInstance("memory").resetStore();
+
+      const recoveredTssKey = await coreKitInstance._UNSAFE_recoverTssKey([factorKeyDevice, factorKeyRecovery]);
+      assert.strictEqual(recoveredTssKey, exportedTssKey1);
+
+      await criticalResetAccount(coreKitInstance);
+      BrowserStorage.getInstance("memory").resetStore();
+
+      // reinitialize corekit
+      const newEmail = testVariable.importKeyEmail;
+      const newLogin = await mockLogin(newEmail);
+
+      const newIdTokenLoginParams = {
+        verifier: "torus-test-health",
+        verifierId: newLogin.parsedToken.email,
+        idToken: newLogin.idToken,
+        importTssKey: recoveredTssKey,
+      } as IdTokenLoginParams;
+
+      const coreKitInstance2 = new Web3AuthMPCCoreKit({
+        web3AuthClientId: "torus-key-test",
+        web3AuthNetwork: WEB3AUTH_NETWORK.DEVNET,
+        baseUrl: "http://localhost:3000",
+        uxMode: "nodejs",
+        tssLib: TssLib,
+        storageKey: "memory",
+      });
+
+      await coreKitInstance2.init();
+      await coreKitInstance2.loginWithJWT(newIdTokenLoginParams);
+
+      const exportedTssKey = await coreKitInstance2._UNSAFE_exportTssKey();
+      criticalResetAccount(coreKitInstance2);
+      BrowserStorage.getInstance("memory").resetStore();
+
+      assert.strictEqual(exportedTssKey, recoveredTssKey);
+    });
+
+    t.afterEach(function () {
+      return log.info("finished running recovery test");
+    });
+    t.after(function () {
+      return log.info("finished running recovery tests");
+    });
+  });
+};
+
+const variable: ImportKeyTestVariable[] = [{ manualSync: false, email: "emailexport", importKeyEmail: "emailimport" }];
+
+variable.forEach(async (testVariable) => {
+  await ImportTest(testVariable);
+});
diff --git a/tests/login.spec.ts b/tests/login.spec.ts
@@ -8,7 +8,7 @@ import * as TssLib from "@toruslabs/tss-lib-node";
 import BN from "bn.js";
 import { ec as EC } from "elliptic";
 
-import { COREKIT_STATUS, WEB3AUTH_NETWORK, WEB3AUTH_NETWORK_TYPE, Web3AuthMPCCoreKit } from "../src";
+import { BrowserStorage, COREKIT_STATUS, WEB3AUTH_NETWORK, WEB3AUTH_NETWORK_TYPE, Web3AuthMPCCoreKit } from "../src";
 import { criticalResetAccount, mockLogin } from "./setup";
 
 type TestVariable = {
@@ -52,6 +52,7 @@ variable.forEach((testVariable) => {
   test(`#Login Test with JWT + logout :  ${testNameSuffix}`, async (t) => {
     t.before(async function () {
       if (coreKitInstance.status === COREKIT_STATUS.INITIALIZED) await criticalResetAccount(coreKitInstance);
+      BrowserStorage.getInstance("memory").resetStore();
     });
 
     t.after(async function () {
diff --git a/tests/mpcCoreKit.ts b/tests/mpcCoreKit.ts
diff --git a/tests/signing.spec.ts b/tests/signing.spec.ts

Original file line number	Diff line number	Diff line change
`@@ -6,6 +6,6 @@`
`6`	`6`	`"eslint.format.enable": true,`
`7`	`7`	`"eslint.lintTask.enable": true,`
`8`	`8`	`"editor.codeActionsOnSave": {`
`9`		`- "source.fixAll": true`
	`9`	`+ "source.fixAll": "explicit"`
`10`	`10`	`}`
`11`	`11`	`}`