Change Password flow

Author: everdimension

src/background/Wallet/Wallet.ts

@@ -315,6 +315,22 @@ export class Wallet {
     await this.syncWithWalletStore();
   }
 
+  async assignNewCredentials({
+    params: { credentials, newCredentials },
+  }: PublicMethodParams<{
+    credentials: SessionCredentials;
+    newCredentials: SessionCredentials;
+  }>) {
+    this.ensureRecord(this.record);
+    this.record = await Model.reEncryptRecord(this.record, {
+      credentials,
+      newCredentials,
+    });
+    this.userCredentials = newCredentials;
+    await this.updateWalletStore(this.record);
+    this.setExpirationForSeedPhraseEncryptionKey(1000 * 120);
+  }
+
   async resetCredentials() {
     this.userCredentials = null;
   }

src/background/Wallet/WalletRecord.ts

@@ -1,5 +1,5 @@
 import { decrypt, encrypt } from 'src/modules/crypto';
-import { produce } from 'immer';
+import { createDraft, finishDraft, produce } from 'immer';
 import { nanoid } from 'nanoid';
 import sortBy from 'lodash/sortBy';
 import { toChecksumAddress } from 'src/modules/ethereum/toChecksumAddress';
@@ -477,7 +477,27 @@ export class WalletRecordModel {
       })
     );
 
-    return WalletRecordModel.verifyStateIntegrity(entry as WalletRecord);
+    return WalletRecordModel.verifyStateIntegrity(entry);
+  }
+
+  static async reEncryptRecord(
+    record: WalletRecord,
+    {
+      credentials,
+      newCredentials,
+    }: { credentials: SessionCredentials; newCredentials: SessionCredentials }
+  ) {
+    // Async update flow for Immer: https://immerjs.github.io/immer/async/
+    const draft = createDraft(record);
+    for (const group of draft.walletManager.groups) {
+      if (isMnemonicContainer(group.walletContainer)) {
+        await group.walletContainer.reEncryptWallets({
+          credentials,
+          newCredentials,
+        });
+      }
+    }
+    return finishDraft(draft);
   }
 
   static async getRecoveryPhrase(

src/background/Wallet/model/WalletContainer.ts

@@ -210,6 +210,25 @@ export class MnemonicWalletContainer extends WalletContainerImpl {
     }
     this.wallets.push(wallet);
   }
+
+  async reEncryptWallets({
+    credentials,
+    newCredentials,
+  }: {
+    credentials: SessionCredentials;
+    newCredentials: SessionCredentials;
+  }) {
+    const { mnemonic: encryptedMnemonic } = this.getFirstWallet();
+    invariant(encryptedMnemonic, 'Must be a Mnemonic WalletContainer');
+    const phrase = await decryptMnemonic(encryptedMnemonic.phrase, credentials);
+    const { seedPhraseEncryptionKey } = newCredentials;
+    const updatedPhrase = await encrypt(seedPhraseEncryptionKey, phrase);
+    for (const wallet of this.wallets) {
+      if (wallet.mnemonic) {
+        wallet.mnemonic.phrase = updatedPhrase;
+      }
+    }
+  }
 }
 
 export class PrivateKeyWalletContainer extends WalletContainerImpl {

src/background/account/Account.ts

@@ -9,17 +9,51 @@ import { eraseAndUpdateToLatestVersion } from 'src/shared/core/version';
 import { currentUserKey } from 'src/shared/getCurrentUser';
 import type { PublicUser, User } from 'src/shared/types/User';
 import { payloadId } from '@walletconnect/jsonrpc-utils';
+import { invariant } from 'src/shared/invariant';
 import { Wallet } from '../Wallet/Wallet';
 import { peakSavedWalletState } from '../Wallet/persistence';
 import type { NotificationWindow } from '../NotificationWindow/NotificationWindow';
 import { credentialsKey } from './storage-keys';
+import { isSessionCredentials } from './Credentials';
 
 const TEMPORARY_ID = 'temporary';
 
 async function sha256({ password, salt }: { password: string; salt: string }) {
   return await getSHA256HexDigest(`${salt}:${password}`);
 }
 
+async function deriveUserKeys({
+  user,
+  credentials,
+}: {
+  user: User;
+  credentials: { password: string } | { encryptionKey: string };
+}) {
+  let encryptionKey: string | null = null;
+  let seedPhraseEncryptionKey: string | null = null;
+  let seedPhraseEncryptionKey_deprecated: CryptoKey | null = null;
+  if ('password' in credentials) {
+    const { password } = credentials;
+    const [key1, key2, key3] = await Promise.all([
+      sha256({ salt: user.id, password }),
+      sha256({ salt: user.salt, password }),
+      createCryptoKey(password, user.salt),
+    ]);
+    encryptionKey = key1;
+    seedPhraseEncryptionKey = key2;
+    seedPhraseEncryptionKey_deprecated = key3;
+  } else {
+    encryptionKey = credentials.encryptionKey;
+  }
+
+  return {
+    id: user.id,
+    encryptionKey,
+    seedPhraseEncryptionKey,
+    seedPhraseEncryptionKey_deprecated,
+  };
+}
+
 class EventEmitter<Events extends EventsMap> {
   private emitter = createNanoEvents<Events>();
 
@@ -77,17 +111,27 @@ export class Account extends EventEmitter<AccountEvents> {
     }
   }
 
-  static async createUser(password: string): Promise<User> {
+  static validatePassword(password: string) {
     const validity = validate({ password });
     if (!validity.valid) {
       throw new Error(validity.message);
     }
+  }
+
+  static async createUser(password: string): Promise<User> {
+    Account.validatePassword(password);
     const id = nanoid(36); // use longer id than default (21)
     const salt = createSalt(); // used to encrypt seed phrases
     const record = { id, salt /* passwordHash: hash */ };
     return record;
   }
 
+  /** Updates salt */
+  static async updateUser(user: User): Promise<User> {
+    const salt = createSalt(); // used to encrypt seed phrases
+    return { id: user.id, salt };
+  }
+
   constructor({
     notificationWindow,
   }: {
@@ -100,6 +144,7 @@ export class Account extends EventEmitter<AccountEvents> {
     this.notificationWindow = notificationWindow;
     this.wallet = new Wallet(TEMPORARY_ID, null, this.notificationWindow);
     this.on('authenticated', () => {
+      // TODO: Call Account.writeCurrentUser() here, too?
       if (this.encryptionKey) {
         Account.writeCredentials({ encryptionKey: this.encryptionKey });
       }
@@ -152,39 +197,60 @@ export class Account extends EventEmitter<AccountEvents> {
     await this.setUser(user, { password }, { isNewUser: false });
   }
 
+  async changePassword({
+    currentPassword,
+    newPassword,
+    user: currentUser,
+  }: {
+    user: User;
+    currentPassword: string;
+    newPassword: string;
+  }) {
+    Account.validatePassword(newPassword);
+    await this.login(currentUser, currentPassword);
+    invariant(this.user, 'User must be set');
+    const updatedUser = await Account.updateUser(this.user);
+    const currentCredentials = await deriveUserKeys({
+      user: currentUser,
+      credentials: { password: currentPassword },
+    });
+    const newCredentials = await deriveUserKeys({
+      user: updatedUser,
+      credentials: { password: newPassword },
+    });
+    console.log({ currentCredentials, newCredentials });
+    if (
+      !isSessionCredentials(currentCredentials) ||
+      !isSessionCredentials(newCredentials)
+    ) {
+      throw new Error('Full credentials are expected');
+    }
+    await this.wallet.assignNewCredentials({
+      id: payloadId(),
+      params: { newCredentials, credentials: currentCredentials },
+    });
+    // Update local state only if the above call was successful
+    this.user = updatedUser;
+    this.encryptionKey = newCredentials.encryptionKey;
+    await Account.writeCurrentUser(this.user);
+    this.emit('authenticated');
+  }
+
   async setUser(
     user: User,
-    credentials: { password: string } | { encryptionKey: string },
+    partialCredentials: { password: string } | { encryptionKey: string },
     { isNewUser = false } = {}
   ) {
     this.user = user;
     this.isPendingNewUser = isNewUser;
-    let seedPhraseEncryptionKey: string | null = null;
-    let seedPhraseEncryptionKey_deprecated: CryptoKey | null = null;
-    if ('password' in credentials) {
-      const { password } = credentials;
-      const [key1, key2, key3] = await Promise.all([
-        sha256({ salt: user.id, password }),
-        sha256({ salt: user.salt, password }),
-        createCryptoKey(password, user.salt),
-      ]);
-      this.encryptionKey = key1;
-      seedPhraseEncryptionKey = key2;
-      seedPhraseEncryptionKey_deprecated = key3;
-    } else {
-      this.encryptionKey = credentials.encryptionKey;
-    }
+    const credentials = await deriveUserKeys({
+      user,
+      credentials: partialCredentials,
+    });
+    this.encryptionKey = credentials.encryptionKey;
     await this.wallet.updateCredentials({
       id: payloadId(),
-      params: {
-        credentials: {
-          id: user.id,
-          encryptionKey: this.encryptionKey,
-          seedPhraseEncryptionKey,
-          seedPhraseEncryptionKey_deprecated,
-        },
-        isNewUser,
-      },
+      params: { credentials, isNewUser },
     });
     if (!this.isPendingNewUser) {
       this.emit('authenticated');
@@ -272,16 +338,25 @@ export class AccountPublicRPC {
     return null;
   }
 
+  async verifyUser(user: PublicUser) {
+    const currentUser = await Account.readCurrentUser();
+    if (!currentUser || currentUser.id !== user.id) {
+      throw new Error(`User ${user.id} not found`);
+    }
+    return currentUser;
+  }
+
   async login({
     params: { user, password },
   }: PublicMethodParams<{
     user: PublicUser;
     password: string;
   }>): Promise<PublicUser | null> {
-    const currentUser = await Account.readCurrentUser();
-    if (!currentUser || currentUser.id !== user.id) {
-      throw new Error(`User ${user.id} not found`);
-    }
+    const currentUser = await this.verifyUser(user);
+    // const currentUser = await Account.readCurrentUser();
+    // if (!currentUser || currentUser.id !== user.id) {
+    //   throw new Error(`User ${user.id} not found`);
+    // }
     const canAuthorize = await this.account.verifyPassword(
       currentUser,
       password
@@ -294,6 +369,21 @@ export class AccountPublicRPC {
     }
   }
 
+  async changePassword({
+    params: { user, currentPassword, newPassword },
+  }: PublicMethodParams<{
+    user: PublicUser;
+    currentPassword: string;
+    newPassword: string;
+  }>) {
+    const currentUser = await this.verifyUser(user);
+    await this.account.changePassword({
+      user: currentUser,
+      currentPassword,
+      newPassword,
+    });
+  }
+
   async hasActivePasswordSession() {
     return this.account.hasActivePasswordSession();
   }