From b83566c1936a49028d6779c531ee8cfba98eff66 Mon Sep 17 00:00:00 2001
From: ieow <cwlee0@gmail.com>
Date: Tue, 27 Feb 2024 13:28:18 +0800
Subject: [PATCH] fix: metadata for salt

---
 packages/core/src/core.ts     | 16 +++++++++-------
 packages/core/src/metadata.ts | 11 ++++++++---
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/packages/core/src/core.ts b/packages/core/src/core.ts
index b1bf84c2..864ca64b 100644
--- a/packages/core/src/core.ts
+++ b/packages/core/src/core.ts
@@ -305,8 +305,9 @@ class ThresholdKey implements ITKey {
         if (useTSS) {
           const { factorEncs, factorPubs, tssPolyCommits } = await this._initializeNewTSSKey(this.tssTag, deviceTSSShare, factorPub, deviceTSSIndex);
 
-          const { encryptedSalt } = await this.getSaltAndEncrypted(this.privKey);
+          const { salt, encryptedSalt } = await this.generateSaltAndEncrypted(this.privKey);
           this.metadata.addTSSData({ tssTag: this.tssTag, tssNonce: 0, tssPolyCommits, factorPubs, factorEncs, encryptedSalt });
+          this._accountSalt = salt;
           // await this._setTKeyStoreItem(TSS_MODULE, {
           //   id: "accountSalt",
           //   value: accountSalt,
@@ -630,7 +631,7 @@ class ThresholdKey implements ITKey {
       if (accountSalt) {
         this._accountSalt = accountSalt;
       } else {
-        const { salt, encryptedSalt } = await this.getSaltAndEncrypted(privKey);
+        const { salt, encryptedSalt } = await this.generateSaltAndEncrypted(privKey);
         this.metadata.addTSSData({ tssTag: this.tssTag, encryptedSalt });
         this._accountSalt = salt;
       }
@@ -919,7 +920,7 @@ class ThresholdKey implements ITKey {
           serverEncs: refreshResponse.serverFactorEncs,
         };
       }
-      const { encryptedSalt } = await this.getSaltAndEncrypted(this.privKey);
+      const { salt, encryptedSalt } = await this.generateSaltAndEncrypted(this.privKey);
       this.metadata.addTSSData({
         tssTag: this.tssTag,
         tssNonce: newTssNonce,
@@ -928,6 +929,7 @@ class ThresholdKey implements ITKey {
         factorEncs,
         encryptedSalt,
       });
+      this._accountSalt = salt;
     } catch (error) {
       this.tssTag = oldTag;
       throw error;
@@ -2024,7 +2026,7 @@ class ThresholdKey implements ITKey {
   }
 
   //
-  private async getSaltAndEncrypted(privKey: BN): Promise<{ salt: string; encryptedSalt: EncryptedMessage }> {
+  private async generateSaltAndEncrypted(privKey: BN): Promise<{ salt: string; encryptedSalt: EncryptedMessage }> {
     if (!privKey) {
       throw CoreError.privateKeyUnavailable();
     }
@@ -2036,9 +2038,9 @@ class ThresholdKey implements ITKey {
 
   private async getAccountSalt() {
     if (this._accountSalt) return this._accountSalt;
-    if (this.metadata.encryptedSalt) {
-      const decrypteSalt = await decrypt(this.privKey.toBuffer(), this.metadata.encryptedSalt);
-      return decrypteSalt.toString("hex");
+    if (Object.keys(this.metadata.encryptedSalt).length) {
+      const decryptedSalt = await decrypt(this.privKey.toBuffer(), this.metadata.encryptedSalt as EncryptedMessage);
+      return decryptedSalt.toString("hex");
     }
     return undefined;
   }
diff --git a/packages/core/src/metadata.ts b/packages/core/src/metadata.ts
index 8a8b63d6..11f16c2d 100644
--- a/packages/core/src/metadata.ts
+++ b/packages/core/src/metadata.ts
@@ -68,13 +68,14 @@ class Metadata implements IMetadata {
     };
   };
 
-  encryptedSalt?: EncryptedMessage;
+  encryptedSalt?: EncryptedMessage | Record<string, unknown>;
 
   constructor(input: Point) {
     this.tssPolyCommits = {};
     this.tssNonces = {};
     this.factorPubs = {};
     this.factorEncs = {};
+    this.encryptedSalt = {};
     this.publicPolynomials = {};
     this.publicShares = {};
     this.generalStore = {};
@@ -86,7 +87,8 @@ class Metadata implements IMetadata {
   }
 
   static fromJSON(value: StringifiedType): Metadata {
-    const { pubKey, polyIDList, generalStore, tkeyStore, scopedStore, nonce, tssNonces, tssPolyCommits, factorPubs, factorEncs } = value;
+    const { pubKey, polyIDList, generalStore, tkeyStore, scopedStore, nonce, tssNonces, tssPolyCommits, factorPubs, factorEncs, encryptedSalt } =
+      value;
     const point = Point.fromCompressedPub(pubKey);
     const metadata = new Metadata(point);
     const unserializedPolyIDList: PolyIDAndShares[] = [];
@@ -115,6 +117,8 @@ class Metadata implements IMetadata {
     }
     if (factorEncs) metadata.factorEncs = factorEncs;
 
+    if (encryptedSalt) metadata.encryptedSalt = encryptedSalt;
+
     for (let i = 0; i < polyIDList.length; i += 1) {
       const serializedPolyID: string = polyIDList[i];
       const arrPolyID = serializedPolyID.split("|");
@@ -198,7 +202,7 @@ class Metadata implements IMetadata {
     if (tssPolyCommits) this.tssPolyCommits[tssTag] = tssPolyCommits;
     if (factorPubs) this.factorPubs[tssTag] = factorPubs;
     if (factorEncs) this.factorEncs[tssTag] = factorEncs;
-    if (encryptedSalt && !this.encryptedSalt) this.encryptedSalt = encryptedSalt;
+    if (encryptedSalt && Object.keys(this.encryptedSalt).length === 0) this.encryptedSalt = encryptedSalt;
   }
 
   // appends shares and public polynomial to metadata.
@@ -342,6 +346,7 @@ class Metadata implements IMetadata {
       ...(this.tssPolyCommits && { tssPolyCommits: this.tssPolyCommits }),
       ...(this.factorPubs && { factorPubs: this.factorPubs }),
       ...(this.factorEncs && { factorEncs: this.factorEncs }),
+      ...(this.encryptedSalt && { encryptedSalt: this.encryptedSalt }),
     };
   }
 }