Change Password flow

.eslintrc.js

@@ -41,7 +41,6 @@ module.exports = {
         '.prettierrc.js',
         'babel.config.js',
         'postcss.config.js',
-        'jest.config.js',
         '*.node.js',
         'webpack.config.js',
         '*.cjs',

package.json

@@ -10,7 +10,8 @@
     "build:production": "node scripts/build.js",
     "type-check": "tsc --noEmit",
     "type-check:watch": "tsc --noEmit --pretty --watch",
-    "test:unit": "node --experimental-vm-modules node_modules/jest/bin/jest.js",
+    "test:unit": "vitest run",
+    "test:unit:watch": "vitest",
     "test:e2e": "playwright test",
     "test:e2e:update": "playwright test --update-snapshots",
     "test:e2e:ui": "playwright test --ui",
@@ -66,7 +67,6 @@
     "@types/canvas-confetti": "^1.6.4",
     "@types/chrome": "^0.0.270",
     "@types/dom-navigation": "^1.0.3",
-    "@types/jest": "^29.5.12",
     "@types/lodash": "^4.17.7",
     "@types/node": "^22.4.0",
     "@types/react": "^18.2.25",
@@ -84,15 +84,14 @@
     "eslint-plugin-react-hooks": "^4.6.0",
     "eslint-plugin-security": "^3.0.1",
     "husky": "^8.0.3",
-    "jest": "^29.7.0",
     "lint-staged": "^13.2.3",
     "npm-run-all": "^4.1.5",
     "parcel": "^2.14.4",
     "prettier": "^2.8.8",
     "process": "^0.11.10",
     "svgo": "^3.3.2",
-    "ts-jest": "^29.1.2",
-    "typescript": "^5.4.3"
+    "typescript": "^5.4.3",
+    "vitest": "^3.1.2"
   },
   "lint-staged": {
     "*.{ts,tsx}": "eslint --cache --fix",
@@ -139,7 +138,7 @@
     "react-qrcode-logo": "^2.9.0",
     "react-router-dom": "^6.18.0",
     "rlp": "^3.0.0",
-    "store-unit": "^1.0.3",
+    "store-unit": "^1.1.0",
     "uuid": "^9.0.0",
     "webextension-polyfill": "^0.10.0",
     "zksync-ethers": "^6.15.2"

src/background/Wallet/Wallet.ts

@@ -190,7 +190,7 @@ export class Wallet {
     this.emitter = createNanoEvents();
 
     this.id = id;
-    this.walletStore = new WalletStore({}, 'wallet');
+    this.walletStore = new WalletStore({});
     this.disposer.add(
       globalPreferences.on('change', (state, prevState) => {
         emitter.emit('globalPreferencesChange', state, prevState);
@@ -239,7 +239,7 @@ export class Wallet {
     if (!this.userCredentials) {
       throw new Error('Cannot save pending wallet: encryptionKey is null');
     }
-    this.walletStore.save(this.id, this.userCredentials.encryptionKey, record);
+    this.walletStore.save(this.id, this.userCredentials, record);
   }
 
   async ready() {

src/background/Wallet/WalletRecord.ts

@@ -1,6 +1,6 @@
 import { decrypt, encrypt } from 'src/modules/crypto';
 import type { Draft } from 'immer';
-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';
@@ -492,7 +492,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

@@ -227,6 +227,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, MISSING_MNEMONIC);
+    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/Wallet/persistence.ts

@@ -1,13 +1,40 @@
 import { produce } from 'immer';
 import { PersistentStore } from 'src/modules/persistent-store';
-import type { Credentials } from '../account/Credentials';
+import { invariant } from 'src/shared/invariant';
+import { getError } from 'src/shared/errors/getError';
+import { ErrorWithEnumerableMessage } from 'src/shared/errors/errors';
+import type { User } from 'src/shared/types/User';
+import type { Credentials, SessionCredentials } from '../account/Credentials';
+import { emitter } from '../events';
+import { Account } from '../account/Account';
 import type { WalletRecord } from './model/types';
 import { WalletRecordModel as Model } from './WalletRecord';
 
 type EncryptedWalletRecord = string;
 
 type WalletStoreState = Record<string, EncryptedWalletRecord | undefined>;
 
+export class InternalBackupError extends ErrorWithEnumerableMessage {
+  didRestore: boolean;
+  constructor(error: Error, { didRestore }: { didRestore: boolean }) {
+    super(error.message);
+    this.name = error.name !== 'Error' ? error.name : 'InternalBackupError';
+    this.didRestore = didRestore;
+  }
+}
+
+type RecordBackup = { user: User; record: string };
+function stringifyBackup({ user, record }: RecordBackup): string {
+  return JSON.stringify({ user, record });
+}
+
+function parseBackup(value: string): RecordBackup {
+  const parsed = JSON.parse(value) as RecordBackup;
+  invariant(parsed.user, 'User not found in backup');
+  invariant(parsed.record, 'Record not found in backup');
+  return parsed;
+}
+
 export class WalletStore extends PersistentStore<WalletStoreState> {
   static key = 'wallet';
   /** Store unencrypted "lastRecord" to avoid unnecessary stringifications */
@@ -41,17 +68,131 @@ export class WalletStore extends PersistentStore<WalletStoreState> {
     return this.lastRecord;
   }
 
-  async save(id: string, encryptionKey: string, record: WalletRecord) {
+  /** Prefer WalletStore['save'] unless necessary */
+  private async encryptAndSave(
+    id: string,
+    credentials: Credentials,
+    record: WalletRecord
+  ) {
+    const encryptedRecord = await Model.encryptRecord(
+      credentials.encryptionKey,
+      record
+    );
+    await this.setState((state) =>
+      produce(state, (draft) => {
+        draft[id] = encryptedRecord;
+      })
+    );
+    this.lastRecord = record;
+  }
+
+  async save(id: string, credentials: Credentials, record: WalletRecord) {
     if (this.lastRecord === record) {
       return;
     }
-    const encryptedRecord = await Model.encryptRecord(encryptionKey, record);
-    this.setState((state) =>
+    await this.encryptAndSave(id, credentials, record);
+  }
+
+  async createBackup(id: string) {
+    /**
+     * Accessing user is a cross-concern, but this is the only way
+     * to make our backup truly atomic and independent:
+     * encrypted record relies on `salt` stored in the user object,
+     * and for a robust backup recovery it's best to store this object
+     * together with the encrypted record
+     */
+    const user = await Account.readCurrentUser();
+    const record = (await this.getSavedState())[id];
+    invariant(record, `Record not found for id: ${id}`);
+    invariant(user && user.id === id, `User not found for id: ${id}`);
+    return this.setState({
+      ...this.state,
+      [`backup:${id}`]: stringifyBackup({ user, record }),
+    });
+  }
+
+  async restoreFromBackup(id: string) {
+    const key = `backup:${id}`;
+    const state = await this.getSavedState();
+    const saved = state[key];
+    invariant(saved, `Backup not found for id: ${id}`);
+    const { user, record } = parseBackup(saved);
+    await Promise.all([
+      this.setState((state) =>
+        produce(state, (draft) => {
+          draft[id] = record;
+          delete draft[key];
+        })
+      ),
+      Account.writeCurrentUser(user),
+    ]);
+  }
+
+  async restoreFromAnyBackup() {
+    const state = await this.getSavedState();
+    const key = Object.keys(state).find((key) => key.startsWith('backup:'));
+    if (key) {
+      await this.restoreFromBackup(key.split(':')[1]);
+    } else {
+      throw new Error('No backups found');
+    }
+  }
+
+  async clearBackup(id: string) {
+    return this.setState((state) =>
       produce(state, (draft) => {
-        draft[id] = encryptedRecord;
+        const key = `backup:${id}`;
+        delete draft[key];
       })
     );
-    this.lastRecord = record;
+  }
+
+  /**
+   * Executes an operation with a backup and an automatic recovery.
+   * Guarantees atomicity by restoring to the previous state if the operation fails.
+   */
+  async withBackup(id: string, operation: () => Promise<unknown>) {
+    await this.createBackup(id);
+    try {
+      await operation();
+      await this.clearBackup(id);
+    } catch (error) {
+      try {
+        await this.restoreFromBackup(id);
+        emitter.emit('globalError', {
+          name: 'internal_error',
+          message: 'Atomic wallet update failed. Restored from backup.',
+        });
+        console.log('Successfully restored wallet record from backup'); // eslint-disable-line no-console
+      } catch {
+        emitter.emit('globalError', {
+          name: 'internal_error',
+          message: 'Atomic wallet update failed. Restore from backup failed.',
+        });
+        throw new InternalBackupError(getError(error), { didRestore: false });
+      }
+      throw error;
+    }
+  }
+
+  async reEncrypt({
+    id,
+    credentials,
+    newCredentials,
+  }: {
+    id: string;
+    credentials: SessionCredentials;
+    newCredentials: SessionCredentials;
+  }) {
+    await this.ready();
+    console.log('reading', { id, credentials, state: this.getState() });
+    const currentRecord = await this.read(id, credentials);
+    invariant(currentRecord, `Record not found for ${id}`);
+    const newRecord = await Model.reEncryptRecord(currentRecord, {
+      credentials,
+      newCredentials,
+    });
+    await this.encryptAndSave(id, newCredentials, newRecord);
   }
 
   deleteMany(keys: string[]) {