Restore vault when state corruption occurs

app/scripts/background.js

@@ -617,6 +617,14 @@ export async function loadStateFromPersistence() {
   // first from preferred, async API:
   const preMigrationVersionedData = await localStore.get();
 
+  const currentVersionNumber = versionedData.meta?.version;
+
+  // Check if the current version number is out of sync with the last known
+  // successfully applied migration.
+  if (localStore.doesMigrationNumberHaveMismatch(currentVersionNumber)) {
+    sentry?.captureMessage('MetaMask - Migration Version Mismatch');
+  }
+
   // report migration errors to sentry
   migrator.on('error', (err) => {
     // get vault structure without secrets
@@ -1275,6 +1283,7 @@ const addAppInstalledEvent = () => {
     });
     return;
   }
+
   setTimeout(() => {
     // If the controller is not set yet, we wait and try to add the "App Installed" event again.
     addAppInstalledEvent();

app/scripts/lib/Stores/BaseStore.ts

@@ -1,6 +1,7 @@
 import type Migrator from '../migrator';
 import firstTimeState from '../../first-time-state';
 import { generateWalletState } from '../../fixtures/generate-wallet-state';
+import localStorageWithOffscreen from './localStorageWithOffscreen';
 
 /**
  * This type is a temporary type that is used to represent the state tree of
@@ -90,6 +91,15 @@ export abstract class BaseStore {
    */
   abstract isSupported: boolean;
 
+  /**
+   * stateCorruptionDetected is a boolean that is set to true if the storage
+   * system attempts to retrieve state and it is missing, but the localStorage
+   * key 'MMStateExisted' is present. This is an indication that the state was
+   * successfully retrieved at some point in the past and thus likely
+   * some form of corruption has occurred.
+   */
+  abstract stateCorruptionDetected: boolean;
+
   /**
    * dataPersistenceFailing is a boolean that is set to true if the storage
    * system attempts to write state and the write operation fails. This is only
@@ -129,6 +139,12 @@ export abstract class BaseStore {
    * storage implementation (e.g. chrome.storage.local).
    */
   set metadata(metadata: { version: number }) {
+    // In an effort to track irregularities in migration and data storage, we
+    // store the last good migration that ran without crashing the app at this
+    // point. We can then compare this number to the current migration version
+    // higher up in this file to determine if something happened during loading
+    // state from storage that caused the migration number to be out of sync.
+    this.lastGoodMigrationVersion = metadata.version;
     this.#metadata = metadata;
   }
 
@@ -141,6 +157,93 @@ export abstract class BaseStore {
     return this.#metadata;
   }
 
+  /**
+   * In exactly one instance the UI needs to be able to modify the state of one
+   * of our localStorage keys "USER_OPTED_IN_TO_RESTORE". This is a helper that
+   * does not require instantiation of the class to set that value.
+   */
+  static optIntoRestoreOnRestart() {
+    localStorageWithOffscreen.setItem('USER_OPTED_IN_TO_RESTORE', 'true');
+  }
+
+  /**
+   * Return whether the user has acknowledged a state corruption issue and
+   * has opted into restoring either their vault, or the extension to its
+   * default state if the vault was not backed up.
+   */
+  get hasUserOptedIntoRestart(): boolean {
+    return localStorageWithOffscreen.getItem('USER_OPTED_IN_TO_RESTORE') === 'true';
+  }
+
+  /**
+   * Set whether the user has acknowledged a state corruption issue and opted
+   * into restoring their state tree to the best possible outcome. This will
+   * remove the presence of the key if 'value' is false. Otherwise if 'value'
+   * is true, then the key will be set to 'true'.
+   */
+  set hasUserOptedIntoRestart(value: boolean) {
+    if (value === true) {
+      localStorageWithOffscreen.setItem('USER_OPTED_IN_TO_RESTORE', 'true');
+    } else {
+      localStorageWithOffscreen.removeItem('USER_OPTED_IN_TO_RESTORE');
+    }
+  }
+
+  /**
+   * After successfully migrating the state tree this method is called to
+   * persist the last known successfully ran migration number to localStorage.
+   * This is useful in detecting
+   */
+  set lastGoodMigrationVersion(version: number) {
+    localStorageWithOffscreen.setItem('lastGoodMigrationVersion', version.toString());
+  }
+
+  get lastGoodMigrationVersion() {
+    return parseInt(
+      localStorageWithOffscreen.getItem('lastGoodMigrationVersion') ?? '0',
+      10,
+    );
+  }
+
+  /**
+   * Compares a version number to the last known good migration number stored
+   * in localStorage. This method is used to determine if the state tree has
+   * been corrupted in some way that caused the migration number to be out of
+   * sync with the last known good migration number.
+   *
+   * @param currentMigrationNumber - The current migration number that is being
+   * compared against. This number is typically the version number of the state
+   * after initially loading from storage.
+   * @returns whether the last good migration number is different from the
+   * current one.
+   */
+  doesMigrationNumberHaveMismatch(currentMigrationNumber: number) {
+    if (localStorageWithOffscreen.getItem('lastGoodMigrationVersion') === null) {
+      return false;
+    }
+    return this.lastGoodMigrationVersion !== currentMigrationNumber;
+  }
+
+  /**
+   * Records a timestamp in localStorage for the last time the state was read
+   * from extension memory successfully. This is used to determine if the state
+   * has ever existed in extension memory, which is useful for determining if
+   * corruption has occurred.
+   */
+  recordStateExistence() {
+    localStorageWithOffscreen.setItem('MMStateExisted', Date.now().toString());
+  }
+
+  /**
+   * Checks if a timestamp exists in localStorage for the last time the state
+   * was read from extension memory successfully. This is used to determine if
+   * the state has ever existed in extension memory, which is useful for
+   * determining if corruption has occurred.
+   */
+  get hasStateExisted() {
+    return localStorageWithOffscreen.getItem('MMStateExisted') !== null;
+  }
+
   /**
    * Generates the first time state tree. This method is used to generate the
    * MetaMask state tree in its initial shape using the migrator.

app/scripts/lib/Stores/ExtensionStore.test.ts

@@ -3,11 +3,29 @@ import Migrator from '../migrator';
 import { ExtensionStore } from './ExtensionStore';
 import { IntermediaryStateType } from './BaseStore';
 
+const EXPECTED_DEFAULT_DATA = {
+  config: {},
+};
+const EXPECTED_DATA_IF_CORRUPTION = {
+  ...EXPECTED_DEFAULT_DATA,
+  PreferencesController: {
+    initializationFlags: {
+      corruptionDetected: true,
+      vaultBackedUp: false,
+    },
+  },
+};
+
 const DEFAULT_INITIAL_STATE = {
-  data: { config: {} },
+  data: EXPECTED_DEFAULT_DATA,
   meta: { version: 0 },
 };
 
+const EXPECTED_STATE_IF_CORRUPTION = {
+  ...DEFAULT_INITIAL_STATE,
+  data: EXPECTED_DATA_IF_CORRUPTION,
+};
+
 jest.mock('webextension-polyfill', () => ({
   runtime: { lastError: null },
   storage: { local: true },
@@ -147,7 +165,7 @@ describe('ExtensionStore', () => {
 
       const result = await localStore.get();
 
-      expect(result).toStrictEqual(DEFAULT_INITIAL_STATE);
+      expect(result).toStrictEqual(EXPECTED_STATE_IF_CORRUPTION);
 
       expect(localStore.mostRecentRetrievedState).toStrictEqual(null);
       expect(localStore.stateCorruptionDetected).toStrictEqual(true);
@@ -162,7 +180,7 @@ describe('ExtensionStore', () => {
 
       const result = await localStore.get();
 
-      expect(result).toStrictEqual(DEFAULT_INITIAL_STATE);
+      expect(result).toStrictEqual(EXPECTED_STATE_IF_CORRUPTION);
 
       expect(localStore.mostRecentRetrievedState).toStrictEqual(null);
       expect(localStore.stateCorruptionDetected).toStrictEqual(true);
@@ -179,7 +197,7 @@ describe('ExtensionStore', () => {
 
       const result = await localStore.get();
 
-      expect(result).toStrictEqual(DEFAULT_INITIAL_STATE);
+      expect(result).toStrictEqual(EXPECTED_STATE_IF_CORRUPTION);
 
       expect(localStore.mostRecentRetrievedState).toStrictEqual(null);
       expect(localStore.stateCorruptionDetected).toStrictEqual(true);
@@ -196,7 +214,7 @@ describe('ExtensionStore', () => {
 
       const result = await localStore.get();
 
-      expect(result).toStrictEqual(DEFAULT_INITIAL_STATE);
+      expect(result).toStrictEqual(EXPECTED_STATE_IF_CORRUPTION);
 
       expect(localStore.mostRecentRetrievedState).toStrictEqual(null);
       expect(localStore.stateCorruptionDetected).toStrictEqual(true);