refactor: Refactor state classes to prepare for state corruption backup mitigation

app/scripts/background.js

@@ -49,11 +49,12 @@ import {
 import { getCurrentChainId } from '../../shared/modules/selectors/networks';
 import { addNonceToCsp } from '../../shared/modules/add-nonce-to-csp';
 import { checkURLForProviderInjection } from '../../shared/modules/provider-injection';
+import { PersistanceManager } from './lib/Stores/PersistanceManager';
+import ExtensionStore from './lib/Stores/ExtensionStore';
+import ReadOnlyNetworkStore from './lib/Stores/ReadOnlyNetworkStore';
 import migrations from './migrations';
 import Migrator from './lib/migrator';
 import ExtensionPlatform from './platforms/extension';
-import LocalStore from './lib/local-store';
-import ReadOnlyNetworkStore from './lib/network-store';
 import { SENTRY_BACKGROUND_STATE } from './constants/sentry-state';
 
 import createStreamSink from './lib/createStreamSink';
@@ -63,7 +64,6 @@ import NotificationManager, {
 import MetamaskController, {
   METAMASK_CONTROLLER_EVENTS,
 } from './metamask-controller';
-import rawFirstTimeState from './first-time-state';
 import getFirstPreferredLangCode from './lib/get-first-preferred-lang-code';
 import getObjStructure from './lib/getObjStructure';
 import setupEnsIpfsResolver from './lib/ens-ipfs/setup';
@@ -72,8 +72,9 @@ import {
   getPlatform,
   shouldEmitDappViewedEvent,
 } from './lib/util';
-import { generateWalletState } from './fixtures/generate-wallet-state';
 import { createOffscreen } from './offscreen';
+import { generateWalletState } from './fixtures/generate-wallet-state';
+import rawFirstTimeState from './first-time-state';
 
 /* eslint-enable import/first */
 
@@ -87,7 +88,15 @@ const BADGE_MAX_COUNT = 9;
 
 // Setup global hook for improved Sentry state snapshots during initialization
 const inTest = process.env.IN_TEST;
-const localStore = inTest ? new ReadOnlyNetworkStore() : new LocalStore();
+const migrator = new Migrator({
+  migrations,
+  defaultVersion: process.env.WITH_STATE
+    ? FIXTURE_STATE_METADATA_VERSION
+    : null,
+});
+
+const _localStore = inTest ? new ReadOnlyNetworkStore() : new ExtensionStore();
+const localStore = new PersistanceManager({ localStore: _localStore });
 global.stateHooks.getMostRecentPersistedState = () =>
   localStore.mostRecentRetrievedState;
 
@@ -113,7 +122,6 @@ let uiIsTriggering = false;
 const openMetamaskTabsIDs = {};
 const requestAccountTabIds = {};
 let controller;
-let versionedData;
 const tabOriginMapping = {};
 
 if (inTest || process.env.METAMASK_DEBUG) {
@@ -599,12 +607,6 @@ async function loadPhishingWarningPage() {
  */
 export async function loadStateFromPersistence() {
   // migrations
-  const migrator = new Migrator({
-    migrations,
-    defaultVersion: process.env.WITH_STATE
-      ? FIXTURE_STATE_METADATA_VERSION
-      : null,
-  });
   migrator.on('error', console.warn);
 
   if (process.env.WITH_STATE) {
@@ -614,31 +616,21 @@ export async function loadStateFromPersistence() {
 
   // read from disk
   // first from preferred, async API:
-  versionedData =
+  const preMigrationVersionedData =
     (await localStore.get()) || migrator.generateInitialState(firstTimeState);
 
-  // check if somehow state is empty
-  // this should never happen but new error reporting suggests that it has
-  // for a small number of users
-  // https://github.com/metamask/metamask-extension/issues/3919
-  if (versionedData && !versionedData.data) {
-    // unable to recover, clear state
-    versionedData = migrator.generateInitialState(firstTimeState);
-    sentry.captureMessage('MetaMask - Empty vault found - unable to recover');
-  }
-
   // report migration errors to sentry
   migrator.on('error', (err) => {
     // get vault structure without secrets
-    const vaultStructure = getObjStructure(versionedData);
+    const vaultStructure = getObjStructure(preMigrationVersionedData);
     sentry.captureException(err, {
       // "extra" key is required by Sentry
       extra: { vaultStructure },
     });
   });
 
   // migrate data
-  versionedData = await migrator.migrateData(versionedData);
+  const versionedData = await migrator.migrateData(preMigrationVersionedData);
   if (!versionedData) {
     throw new Error('MetaMask - migrator returned undefined');
   } else if (!isObject(versionedData.meta)) {
@@ -656,7 +648,7 @@ export async function loadStateFromPersistence() {
     );
   }
   // this initializes the meta/version data as a class variable to be used for future writes
-  localStore.setMetadata(versionedData.meta);
+  localStore.metadata = versionedData.meta;
 
   // write to disk
   localStore.set(versionedData.data);
@@ -1310,6 +1302,8 @@ async function onInstall() {
   if (process.env.IN_TEST) {
     addAppInstalledEvent();
   } else if (!storeAlreadyExisted && !process.env.METAMASK_DEBUG) {
+    // If storeAlreadyExisted is true then this is a fresh installation
+    // and an app installed event should be tracked.
     addAppInstalledEvent();
     platform.openExtensionInBrowser();
   }

app/scripts/lib/Stores/BaseStore.ts

@@ -0,0 +1,67 @@
+/**
+ * This type is a temporary type that is used to represent the state tree of
+ * MetaMask. This type is used in the BaseStore class and its extending classes
+ * and should ultimately be replaced by the fully typed State Tree once that is
+ * available for consumption. We should likely optimize the state tree by
+ * storing the individual controllers in their own keys in the state tree. This
+ * would allow for partial updates at the controller state level, without
+ * modifying the entire data key.
+ */
+export type IntermediaryStateType = Record<string, unknown>;
+
+/**
+ * This type represents the 'meta' key on the state object. This key is used to
+ * store the current version of the state tree as set in the various migrations
+ * ran by the migrator. This key is used to determine if the state tree should
+ * be updated when the extension is loaded, by comparing the version to the
+ * target versions of the migrations.
+ */
+export type MetaData = { version: number };
+
+/**
+ * This type represents the structure of the storage object that is saved in
+ * extension storage. This object has two keys, 'data' and 'meta'. The 'data'
+ * key is the entire state tree of MetaMask and the meta key contains an object
+ * with a single key 'version' that is the current version of the state tree.
+ */
+export type MetaMaskStorageStructure = {
+  data?: IntermediaryStateType;
+  meta?: MetaData;
+};
+
+/**
+ * When loading state from storage, if the state is not available, then the
+ * extension storage api, at least in the case of chrome, returns an empty
+ * object. This type represents that empty object to be used in error handling
+ * and state initialization.
+ */
+export type EmptyState = Omit<MetaMaskStorageStructure, 'data' | 'meta'>;
+
+/**
+ * The BaseStore class is an abstract class designed to be extended by other
+ * classes that implement the abstract methods `set` and `get`. This class
+ * provides the foundation for different storage implementations, enabling
+ * them to adhere to a consistent interface for retrieving and setting
+ * application state.
+ *
+ * Responsibilities of extending classes:
+ * 1. **Retrieve State:**
+ * - Implement a `get` method that retrieves the current state from the
+ * underlying storage system. This method should handle scenarios where
+ * the state is unavailable by providing a fallback, such as returning
+ * null` or an appropriate default value.
+ *
+ * 2. **Set State:**
+ * - Implement a `set` method that updates the state in the underlying
+ * storage system. This method should handle necessary validation or
+ * error handling to ensure the state is persisted correctly.
+ *
+ * This class does not provide any concrete implementation for these methods,
+ * leaving the specifics to the extending classes based on the storage
+ * mechanism they represent.
+ */
+export abstract class BaseStore {
+  abstract set(state: IntermediaryStateType): Promise<void>;
+
+  abstract get(): Promise<MetaMaskStorageStructure | null>;
+}

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

@@ -0,0 +1,122 @@
+import browser from 'webextension-polyfill';
+import { checkForLastError } from '../../../../shared/modules/browser-runtime.utils';
+import ExtensionStore from './ExtensionStore';
+
+const MOCK_STATE = { data: {}, meta: { version: 1 } };
+
+jest.mock('webextension-polyfill', () => ({
+  runtime: { lastError: null },
+  storage: { local: true },
+}));
+
+jest.mock('../../../../shared/modules/browser-runtime.utils', () => ({
+  checkForLastError: jest.fn(),
+}));
+
+const setup = (
+  options: { localMock?: { get?: unknown; set?: unknown } | false } = {},
+) => {
+  if (typeof options.localMock === 'undefined') {
+    browser.storage.local =
+      jest.fn() as unknown as browser.Storage.LocalStorageArea;
+  } else if (options.localMock === false) {
+    browser.storage.local =
+      undefined as unknown as browser.Storage.LocalStorageArea;
+  } else {
+    browser.storage.local =
+      options.localMock as unknown as browser.Storage.LocalStorageArea;
+  }
+  return new ExtensionStore();
+};
+describe('ExtensionStore', () => {
+  beforeEach(() => {
+    global.sentry = {
+      captureException: jest.fn(),
+    };
+  });
+
+  afterEach(() => {
+    jest.resetModules();
+    jest.clearAllMocks();
+    browser.storage.local =
+      undefined as unknown as browser.Storage.LocalStorageArea;
+  });
+  describe('constructor', () => {
+    it('should set isSupported property to false when browser does not support local storage', () => {
+      const localStore = setup({ localMock: false });
+
+      expect(localStore.isSupported).toBe(false);
+    });
+
+    it('should set isSupported property to true when browser supports local storage', () => {
+      const localStore = setup();
+      expect(localStore.isSupported).toBe(true);
+    });
+  });
+
+  describe('set', () => {
+    it('should throw an error if called in a browser that does not support local storage', async () => {
+      const localStore = setup({ localMock: false });
+      await expect(() => localStore.set(MOCK_STATE)).rejects.toThrow(
+        'Metamask- cannot persist state to local store as this browser does not support this action',
+      );
+    });
+
+    it('should not throw if passed a valid argument and metadata has been set', async () => {
+      const setMock = jest.fn();
+
+      const localStore = setup({ localMock: { set: setMock } });
+      await expect(async function () {
+        localStore.set(MOCK_STATE);
+      }).not.toThrow();
+    });
+
+    it('should call the browser storage.local.set method', async () => {
+      const setMock = jest.fn();
+      const localStore = setup({ localMock: { set: setMock } });
+
+      await localStore.set(MOCK_STATE);
+
+      expect(setMock).toHaveBeenCalledWith(MOCK_STATE);
+    });
+
+    it('throws an error if checkForLastError returns an error after setting', async () => {
+      const setMock = jest.fn();
+      setMock.mockResolvedValueOnce(undefined);
+      const fakeError = new Error('Some set error');
+      (checkForLastError as jest.Mock).mockReturnValueOnce(fakeError);
+
+      const localStore = setup({ localMock: { set: setMock } });
+      await expect(
+        localStore.set({ data: { abc: 123 }, meta: { version: 10 } }),
+      ).rejects.toThrow('Some set error');
+    });
+  });
+
+  describe('get', () => {
+    it('should return null if called in a browser that does not support local storage', async () => {
+      const localStore = setup({ localMock: false });
+      const result = await localStore.get();
+      expect(result).toBe(null);
+    });
+
+    it('should call the browser storage.local.get method', async () => {
+      const getMock = jest.fn();
+      const localStore = setup({ localMock: { get: getMock } });
+
+      await localStore.get();
+
+      expect(getMock).toHaveBeenCalledWith(null);
+    });
+
+    it('throws an error if checkForLastError returns an error', async () => {
+      const getMock = jest.fn();
+      getMock.mockResolvedValueOnce({ test: true });
+      const fakeError = new Error('Some browser error');
+      (checkForLastError as jest.Mock).mockReturnValueOnce(fakeError);
+
+      const localStore = setup({ localMock: { get: getMock } });
+      await expect(localStore.get()).rejects.toThrow('Some browser error');
+    });
+  });
+});

app/scripts/lib/Stores/ExtensionStore.ts

@@ -0,0 +1,70 @@
+import browser from 'webextension-polyfill';
+import log from 'loglevel';
+import { checkForLastError } from '../../../../shared/modules/browser-runtime.utils';
+import {
+  type IntermediaryStateType,
+  MetaMaskStorageStructure,
+  BaseStore,
+} from './BaseStore';
+
+/**
+ * An implementation of the MetaMask Extension BaseStore system that uses the
+ * browser.storage.local API to persist and retrieve state.
+ */
+export default class ExtensionStore extends BaseStore {
+  isSupported: boolean;
+
+  constructor() {
+    super();
+    this.isSupported = Boolean(browser.storage.local);
+    if (!this.isSupported) {
+      log.error('Storage local API not available.');
+    }
+  }
+
+  /**
+   * Returns all of the keys currently saved
+   *
+   * @returns the key-value map from local storage
+   */
+  async get(): Promise<MetaMaskStorageStructure | null> {
+    if (!this.isSupported) {
+      log.error('Storage local API not available.');
+      return null;
+    }
+    const { local } = browser.storage;
+    const result = await local.get(null);
+    const err = checkForLastError();
+    if (err) {
+      throw err;
+    }
+    return result;
+  }
+
+  /**
+   * Sets the key in local state
+   *
+   * @param obj - The key to set
+   * @param obj.data - The MetaMask State tree
+   * @param obj.meta - The metadata object
+   * @param obj.meta.version - The version of the state tree determined by the
+   * migration
+   * @returns a promise resolving to undefined.
+   */
+  async set(obj: {
+    data: IntermediaryStateType;
+    meta: { version: number };
+  }): Promise<void> {
+    if (!this.isSupported) {
+      throw new Error(
+        'Metamask- cannot persist state to local store as this browser does not support this action',
+      );
+    }
+    const { local } = browser.storage;
+    await local.set(obj);
+    const err = checkForLastError();
+    if (err) {
+      throw err;
+    }
+  }
+}