refactor: Refactor extension storage classes

app/scripts/background.js

@@ -56,11 +56,11 @@ 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 { 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';
@@ -70,7 +70,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';
@@ -79,7 +78,6 @@ import {
   getPlatform,
   shouldEmitDappViewedEvent,
 } from './lib/util';
-import { generateWalletState } from './fixtures/generate-wallet-state';
 import { createOffscreen } from './offscreen';
 
 /* eslint-enable import/first */
@@ -94,12 +92,19 @@ 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({ migrator })
+  : new ExtensionStore({ migrator });
 global.stateHooks.getMostRecentPersistedState = () =>
   localStore.mostRecentRetrievedState;
 
 const { sentry } = global;
-let firstTimeState = { ...rawFirstTimeState };
 
 const metamaskInternalProcessHash = {
   [ENVIRONMENT_TYPE_POPUP]: true,
@@ -120,7 +125,6 @@ let uiIsTriggering = false;
 const openMetamaskTabsIDs = {};
 const requestAccountTabIds = {};
 let controller;
-let versionedData;
 const tabOriginMapping = {};
 
 if (inTest || process.env.METAMASK_DEBUG) {
@@ -607,46 +611,24 @@ 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) {
-    const stateOverrides = await generateWalletState();
-    firstTimeState = { ...firstTimeState, ...stateOverrides };
-  }
-
   // read from disk
   // first from preferred, async API:
-  versionedData =
-    (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');
-  }
+  const preMigrationVersionedData = await localStore.get();
 
   // 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)) {
@@ -664,7 +646,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);
@@ -1301,12 +1283,12 @@ const addAppInstalledEvent = () => {
 
 // On first install, open a new tab with MetaMask
 async function onInstall() {
-  const storeAlreadyExisted = Boolean(await localStore.get());
-  // If the store doesn't exist, then this is the first time running this script,
-  // and is therefore an install
+  const isFirstTimeInstall = await localStore.isFirstTimeInstall();
   if (process.env.IN_TEST) {
     addAppInstalledEvent();
-  } else if (!storeAlreadyExisted && !process.env.METAMASK_DEBUG) {
+  } else if (isFirstTimeInstall && !process.env.METAMASK_DEBUG) {
+    // If isFirstTimeInstall 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,167 @@
+import type Migrator from '../migrator';
+import firstTimeState from '../../first-time-state';
+import { generateWalletState } from '../../fixtures/generate-wallet-state';
+
+/**
+ * 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 meant to be extended by other classes
+ * that implement the methods and properties marked as abstract. There are a
+ * few properties and methods that are not abstract and are implemented here to
+ * be consumed by the extending classes. At the time of writing this class
+ * there are only two extending classes: ReadOnlyNetworkStore and
+ * ExtensionStore. Both of these extending classes are the result of
+ * refactoring the previous storage implementation to TypeScript while
+ * consolidating some logic related to storage that was external to the
+ * implementation of those storage systems. ReadOnlyNetworkStore is a class
+ * that is used while in an End To End or other Test environment where the full
+ * chrome storage API may not be available. ExtensionStore is the class that is
+ * used when the full chrome storage API is available. While Chrome is the
+ * target of this documentation, Firefox also has a mostly identical storage
+ * API that is used interchangeably.
+ *
+ * The classes that extend this system take on the responsibilities listed here
+ * 1. Retrieve the current state from the underlying storage system. If that
+ * state is unavailable, then the storage system should return a default state
+ * in the case that this is the first time the extension has been installed. If
+ * the state is not available due to some form of possible corruption, using
+ * the best methods available to detect such things, then a backup of the vault
+ * should be inserted into a state tree that otherwise resembles a first time
+ * installation. If the backup of the vault is unavailable, then a default
+ * state tree should be used. In any case we should provide clear and concise
+ * communication to the user about what happened and their best recourse for
+ * handling the situation if the extension cannot gracefully recover.
+ *
+ * 2. Set the current state to the underlying storage system. This should be
+ * implemented in such a way that the current metadata is stored in a separate
+ * key that is tracked by the storage system. This metadata should *not* be a
+ * input to the set method. If the underlying storage system allows for partial
+ * state objects it should be sufficient to pass the data key, which is the
+ * full MetaMask state tree. If not, then the metadata should be supplied by
+ * the storage system itself.
+ *
+ * 3. Provide a method for generating a first time state tree. This method is
+ * implemented as a part of this Abstract class and should not be overwritten
+ * unless future work requires specific implementations for different storage
+ * systems. This method should return a state tree that is the default state
+ * tree for a new install.
+ */
+export abstract class BaseStore {
+  /**
+   * isSupported is a boolean that is set to true if the underlying storage
+   * system is supported by the current browser and implementation.
+   */
+  abstract isSupported: 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
+   * used as a way of deduplicating error reports sent to sentry as it is
+   * likely that multiple writes will fail concurrently.
+   */
+  abstract dataPersistenceFailing: boolean;
+
+  /**
+   * mostRecentRetrievedState is a property that holds the most recent state
+   * successfully retrieved from memory. Due to the nature of async read
+   * operations it is beneficial to have a near real-time snapshot of the state
+   * for sending data to sentry as well as other developer tooling.
+   */
+  abstract mostRecentRetrievedState: MetaMaskStorageStructure | null;
+
+  /**
+   * metadata is a property that holds the current metadata object. This object
+   * includes a single key which is 'version' and contains the current version
+   * number of the state tree. This is only incremented via the migrator and in
+   * a well functioning (typical) install should match the latest migration's
+   * version number.
+   */
+  #metadata?: { version: number };
+
+  /**
+   * migrator is a property that holds the migrator instance that is used to
+   * migrate state from one shape to another. This migrator is used to create
+   * the first time state tree.
+   */
+  abstract migrator: Migrator;
+
+  /**
+   * Sets the current metadata. The set method that is implemented in storage
+   * classes only requires an object that is set on the 'data' key. The
+   * metadata key of this class is set on the 'meta' key of the underlying
+   * storage implementation (e.g. chrome.storage.local).
+   */
+  set metadata(metadata: { version: number }) {
+    this.#metadata = metadata;
+  }
+
+  /**
+   * Gets the current metadata object and returns it. The underlying key is
+   * private and implemented in the BaseStore class so that the extending class
+   * can access it through this getter.
+   */
+  get metadata(): { version: number } | undefined {
+    return this.#metadata;
+  }
+
+  /**
+   * Generates the first time state tree. This method is used to generate the
+   * MetaMask state tree in its initial shape using the migrator.
+   *
+   * @returns state - The first time state tree generated by the migrator
+   */
+  async generateFirstTimeState() {
+    let _firstTimeState = { ...firstTimeState };
+    if (process.env.WITH_STATE) {
+      const stateOverrides = await generateWalletState();
+      _firstTimeState = { ..._firstTimeState, ...stateOverrides };
+    }
+
+    return this.migrator.generateInitialState(
+      _firstTimeState,
+    ) as Required<MetaMaskStorageStructure>;
+  }
+
+  abstract set(state: IntermediaryStateType): Promise<void>;
+
+  abstract get(): Promise<MetaMaskStorageStructure>;
+
+  abstract isFirstTimeInstall(): Promise<boolean>;
+}