Make changePassword atomic

Implement backup and restore mechanics

Author: everdimension

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() {
@@ -294,25 +294,6 @@ 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;
-
-    const { encryptionKey } = this.userCredentials;
-    await this.walletStore.encryptAndSave(this.id, encryptionKey, this.record);
-
-    this.setExpirationForSeedPhraseEncryptionKey(1000 * 120);
-  }
-
   async resetCredentials() {
     this.userCredentials = null;
   }

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 */
@@ -42,25 +69,130 @@ export class WalletStore extends PersistentStore<WalletStoreState> {
   }
 
   /** Prefer WalletStore['save'] unless necessary */
-  async encryptAndSave(
+  private async encryptAndSave(
     id: string,
-    encryptionKey: string,
+    credentials: Credentials,
     record: WalletRecord
   ) {
-    const encryptedRecord = await Model.encryptRecord(encryptionKey, record);
-    this.setState((state) =>
+    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, encryptionKey: string, record: WalletRecord) {
+  async save(id: string, credentials: Credentials, record: WalletRecord) {
     if (this.lastRecord === record) {
       return;
     }
-    await this.encryptAndSave(id, encryptionKey, record);
+    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) => {
+        const key = `backup:${id}`;
+        delete draft[key];
+      })
+    );
+  }
+
+  /**
+   * 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[]) {

src/background/account/Account.ts

@@ -98,8 +98,9 @@ export class Account extends EventEmitter<AccountEvents> {
   private notificationWindow: NotificationWindow;
 
   isPendingNewUser: boolean;
+  isWriteLocked = false;
 
-  private static async writeCurrentUser(user: User) {
+  static async writeCurrentUser(user: User) {
     await BrowserStorage.set(currentUserKey, user);
   }
 
@@ -213,13 +214,17 @@ export class Account extends EventEmitter<AccountEvents> {
   }
 
   async initialize() {
-    // Try to automatically login if credentials are found in storage
-    const credentials = await Account.readCredentials();
-    const user = await Account.readCurrentUser();
-    if (user && credentials) {
-      const valid = await this.verifyCredentials(user, credentials);
-      if (valid) {
-        await this.setUser(user, credentials);
+    try {
+      await this.wallet.walletStore.restoreFromAnyBackup();
+    } catch {
+      // Try to automatically login if credentials are found in storage
+      const credentials = await Account.readCredentials();
+      const user = await Account.readCurrentUser();
+      if (user && credentials) {
+        const valid = await this.verifyCredentials(user, credentials);
+        if (valid) {
+          await this.setUser(user, credentials);
+        }
       }
     }
   }
@@ -295,15 +300,25 @@ export class Account extends EventEmitter<AccountEvents> {
     ) {
       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');
+    await this.logout();
+    const { walletStore } = this.wallet;
+    try {
+      this.isWriteLocked = true;
+      await walletStore.withBackup(currentCredentials.id, async () => {
+        await walletStore.reEncrypt({
+          id: currentCredentials.id,
+          credentials: currentCredentials,
+          newCredentials,
+        });
+        await Account.writeCurrentUser(updatedUser);
+        await this.login(updatedUser, newPassword);
+      });
+    } catch (error) {
+      await this.login(currentUser, currentPassword);
+      throw error;
+    } finally {
+      this.isWriteLocked = false;
+    }
   }
 
   async setUser(
@@ -422,6 +437,7 @@ export class AccountPublicRPC {
     user: PublicUser;
     password: string;
   }>): Promise<PublicUser | null> {
+    invariant(!this.account.isWriteLocked, 'Atomic operation in progress');
     const currentUser = await this.verifyUser(user);
     const canAuthorize = await this.account.verifyPassword(
       currentUser,

src/background/events.ts

@@ -37,7 +37,7 @@ export const emitter = createNanoEvents<{
   chainChanged: (chain: Chain, origin: string) => void;
   'ui:chainSelected': (chain: Chain) => void;
   globalError: (data: {
-    name: 'network_error' | 'signing_error';
+    name: 'network_error' | 'signing_error' | 'internal_error';
     message: string;
   }) => void;
   switchChainError: (chainId: ChainId, origin: string, error: unknown) => void;

src/modules/persistent-store-base.ts

@@ -0,0 +1,76 @@
+import { Store } from 'store-unit';
+
+type Options<S> = {
+  retrieve: (key: string) => Promise<S | undefined>;
+  save: (key: string, value: S) => Promise<void>;
+};
+
+export class PersistentStoreBase<T> extends Store<T> {
+  protected key: string;
+  protected isReady: boolean;
+  private readyPromise: Promise<void>;
+
+  public options: Options<T>;
+  public defaultOptions: Options<T> = {
+    retrieve: () => {
+      throw new Error('PersistentStore:retrieve Not Implemented');
+    },
+    save: () => {
+      throw new Error('PersistentStore:save Not Implemented');
+    },
+  };
+
+  constructor(initialState: T, key: string, options: Partial<Options<T>> = {}) {
+    super(initialState);
+    this.key = key;
+    this.options = { ...this.defaultOptions, ...options };
+    this.isReady = false;
+    this.readyPromise = this.restore();
+    this.on('change', (state) => {
+      // only write to disk after restore so that we're not making
+      // a redundant write on each initialization
+      if (this.isReady) {
+        this.options.save(this.key, state);
+      }
+    });
+  }
+
+  getState() {
+    if (!this.isReady) {
+      throw new Error('Do not access getState() before checking ready()');
+    }
+    return super.getState();
+  }
+
+  setState(...args: Parameters<Store<T>['setState']>) {
+    if (!this.isReady) {
+      if (process.env.NODE_ENV === 'development') {
+        // Throw only in dev mode in case some production flow depends on
+        // setting state sooner. Before this refactoring it was possible, so
+        // it shouldn't really break anything
+        throw new Error(
+          'You are trying to write to a PersistentStore before {ready}'
+        );
+      }
+    }
+    super.setState(...args);
+    return this.options.save(this.key, this.state);
+  }
+
+  async restore() {
+    const saved = await this.options.retrieve(this.key);
+    if (saved) {
+      super.setState(saved);
+    }
+    this.isReady = true;
+  }
+
+  async ready(): Promise<void> {
+    return this.isReady ? Promise.resolve() : this.readyPromise;
+  }
+
+  async getSavedState() {
+    await this.ready();
+    return this.getState();
+  }
+}