Persist state in metaRPCHandler so that we are sure state is persisted before sending back response to actions (#15978)

* persist state in metaRPCHandler so that we are sure state is persisted before sending back response to actions

app/scripts/background.js

@@ -7,9 +7,8 @@ import pump from 'pump';
 import debounce from 'debounce-stream';
 import log from 'loglevel';
 import browser from 'webextension-polyfill';
-import { storeAsStream, storeTransformStream } from '@metamask/obs-store';
+import { storeAsStream } from '@metamask/obs-store';
 import PortStream from 'extension-port-stream';
-import { captureException } from '@sentry/browser';
 
 import { ethErrors } from 'eth-rpc-errors';
 import {
@@ -289,16 +288,11 @@ async function loadStateFromPersistence() {
   if (!versionedData) {
     throw new Error('MetaMask - migrator returned undefined');
   }
+  // this initializes the meta/version data as a class variable to be used for future writes
+  localStore.setMetadata(versionedData.meta);
 
   // write to disk
-  if (localStore.isSupported) {
-    localStore.set(versionedData);
-  } else {
-    // throw in setTimeout so as to not block boot
-    setTimeout(() => {
-      throw new Error('MetaMask - Localstore not supported');
-    });
-  }
+  localStore.set(versionedData.data);
 
   // return just the data
   return versionedData.data;
@@ -339,6 +333,7 @@ function setupController(initState, initLangCode, remoteSourcePort) {
     getOpenMetamaskTabsIds: () => {
       return openMetamaskTabsIDs;
     },
+    localStore,
   });
 
   setupEnsIpfsResolver({
@@ -355,8 +350,7 @@ function setupController(initState, initLangCode, remoteSourcePort) {
   pump(
     storeAsStream(controller.store),
     debounce(1000),
-    storeTransformStream(versionifyData),
-    createStreamSink(persistData),
+    createStreamSink((state) => localStore.set(state)),
     (error) => {
       log.error('MetaMask - Persistence pipeline failed', error);
     },
@@ -364,43 +358,6 @@ function setupController(initState, initLangCode, remoteSourcePort) {
 
   setupSentryGetStateGlobal(controller);
 
-  /**
-   * Assigns the given state to the versioned object (with metadata), and returns that.
-   *
-   * @param {object} state - The state object as emitted by the MetaMaskController.
-   * @returns {VersionedData} The state object wrapped in an object that includes a metadata key.
-   */
-  function versionifyData(state) {
-    versionedData.data = state;
-    return versionedData;
-  }
-
-  let dataPersistenceFailing = false;
-
-  async function persistData(state) {
-    if (!state) {
-      throw new Error('MetaMask - updated state is missing');
-    }
-    if (!state.data) {
-      throw new Error('MetaMask - updated state does not have data');
-    }
-    if (localStore.isSupported) {
-      try {
-        await localStore.set(state);
-        if (dataPersistenceFailing) {
-          dataPersistenceFailing = false;
-        }
-      } catch (err) {
-        // log error so we dont break the pipeline
-        if (!dataPersistenceFailing) {
-          dataPersistenceFailing = true;
-          captureException(err);
-        }
-        log.error('error setting state in local store:', err);
-      }
-    }
-  }
-
   //
   // connect to other contexts
   //

app/scripts/lib/createMetaRPCHandler.js

@@ -1,6 +1,7 @@
 import { ethErrors, serializeError } from 'eth-rpc-errors';
+import { isManifestV3 } from '../../../shared/modules/mv3.utils';
 
-const createMetaRPCHandler = (api, outStream) => {
+const createMetaRPCHandler = (api, outStream, store, localStoreApiWrapper) => {
   return async (data) => {
     if (outStream._writableState.ended) {
       return;
@@ -22,6 +23,10 @@ const createMetaRPCHandler = (api, outStream) => {
       result = await api[data.method](...data.params);
     } catch (err) {
       error = err;
+    } finally {
+      if (isManifestV3 && store && data.method !== 'getState') {
+        localStoreApiWrapper.set(store.getState());
+      }
     }
 
     if (outStream._writableState.ended) {

app/scripts/lib/local-store.js

@@ -1,5 +1,6 @@
 import browser from 'webextension-polyfill';
 import log from 'loglevel';
+import { captureException } from '@sentry/browser';
 import { checkForError } from './util';
 
 /**
@@ -11,6 +12,45 @@ export default class ExtensionStore {
     if (!this.isSupported) {
       log.error('Storage local API not available.');
     }
+    // we use this flag to avoid flooding sentry with a ton of errors:
+    // once data persistence fails once and it flips true we don't send further
+    // data persistence errors to sentry
+    this.dataPersistenceFailing = false;
+  }
+
+  setMetadata(initMetaData) {
+    this.metadata = initMetaData;
+  }
+
+  async set(state) {
+    if (!this.isSupported) {
+      throw new Error(
+        'Metamask- cannot persist state to local store as this browser does not support this action',
+      );
+    }
+    if (!state) {
+      throw new Error('MetaMask - updated state is missing');
+    }
+    if (!this.metadata) {
+      throw new Error(
+        'MetaMask - metadata must be set on instance of ExtensionStore before calling "set"',
+      );
+    }
+    try {
+      // we format the data for storage as an object with the "data" key for the controller state object
+      // and the "meta" key for a metadata object containing a version number that tracks how the data shape
+      // has changed using migrations to adapt to backwards incompatible changes
+      await this._set({ data: state, meta: this.metadata });
+      if (this.dataPersistenceFailing) {
+        this.dataPersistenceFailing = false;
+      }
+    } catch (err) {
+      if (!this.dataPersistenceFailing) {
+        this.dataPersistenceFailing = true;
+        captureException(err);
+      }
+      log.error('error setting state in local store:', err);
+    }
   }
 
   /**
@@ -31,16 +71,6 @@ export default class ExtensionStore {
     return result;
   }
 
-  /**
-   * Sets the key in local state
-   *
-   * @param {object} state - The state to set
-   * @returns {Promise<void>}
-   */
-  async set(state) {
-    return this._set(state);
-  }
-
   /**
    * Returns all of the keys currently saved
    *

app/scripts/lib/local-store.test.js

@@ -0,0 +1,79 @@
+import browser from 'webextension-polyfill';
+import LocalStore from './local-store';
+
+jest.mock('webextension-polyfill', () => ({
+  storage: { local: true },
+}));
+
+const setup = ({ isSupported }) => {
+  browser.storage.local = isSupported;
+  return new LocalStore();
+};
+describe('LocalStore', () => {
+  afterEach(() => {
+    jest.resetModules();
+  });
+  describe('contructor', () => {
+    it('should set isSupported property to false when browser does not support local storage', () => {
+      const localStore = setup({ isSupported: false });
+
+      expect(localStore.isSupported).toBe(false);
+    });
+
+    it('should set isSupported property to true when browser supports local storage', () => {
+      const localStore = setup({ isSupported: true });
+      expect(localStore.isSupported).toBe(true);
+    });
+  });
+
+  describe('setMetadata', () => {
+    it('should set the metadata property on LocalStore', () => {
+      const metadata = { version: 74 };
+      const localStore = setup({ isSupported: true });
+      localStore.setMetadata(metadata);
+
+      expect(localStore.metadata).toStrictEqual(metadata);
+    });
+  });
+
+  describe('set', () => {
+    it('should throw an error if called in a browser that does not support local storage', async () => {
+      const localStore = setup({ isSupported: false });
+      await expect(() => localStore.set()).rejects.toThrow(
+        'Metamask- cannot persist state to local store as this browser does not support this action',
+      );
+    });
+
+    it('should throw an error if not passed a truthy value as an argument', async () => {
+      const localStore = setup({ isSupported: true });
+      await expect(() => localStore.set()).rejects.toThrow(
+        'MetaMask - updated state is missing',
+      );
+    });
+
+    it('should throw an error if passed a valid argument but metadata has not yet been set', async () => {
+      const localStore = setup({ isSupported: true });
+      await expect(() =>
+        localStore.set({ appState: { test: true } }),
+      ).rejects.toThrow(
+        'MetaMask - metadata must be set on instance of ExtensionStore before calling "set"',
+      );
+    });
+
+    it('should not throw if passed a valid argument and metadata has been set', async () => {
+      const localStore = setup({ isSupported: true });
+      localStore.setMetadata({ version: 74 });
+      await expect(async function () {
+        localStore.set({ appState: { test: true } });
+      }).not.toThrow();
+    });
+  });
+
+  describe('get', () => {
+    it('should return undefined if called in a browser that does not support local storage', async () => {
+      const localStore = setup({ isSupported: false });
+      const result = await localStore.get();
+      expect(result).toStrictEqual(undefined);
+    });
+  });
+});

app/scripts/lib/network-store.js

@@ -50,16 +50,37 @@ export default class ReadOnlyNetworkStore {
     return this._state;
   }
 
+  /**
+   * Set metadata/version state
+   *
+   * @param {object} metadata - The metadata/version data to set
+   */
+  setMetadata(metadata) {
+    this.metadata = metadata;
+  }
+
   /**
    * Set state
    *
    * @param {object} state - The state to set
-   * @returns {Promise<void>}
    */
   async set(state) {
+    if (!this.isSupported) {
+      throw new Error(
+        'Metamask- cannot persist state to local store as this browser does not support this action',
+      );
+    }
+    if (!state) {
+      throw new Error('MetaMask - updated state is missing');
+    }
+    if (!this.metadata) {
+      throw new Error(
+        'MetaMask - metadata must be set on instance of ExtensionStore before calling "set"',
+      );
+    }
     if (!this._initialized) {
       await this._initializing;
     }
-    this._state = state;
+    this._state = { data: state, meta: this._metadata };
   }
 }

app/scripts/metamask-controller.js

@@ -201,6 +201,9 @@ export default class MetamaskController extends EventEmitter {
 
     this.controllerMessenger = new ControllerMessenger();
 
+    // instance of a class that wraps the extension's storage local API.
+    this.localStoreApiWrapper = opts.localStore;
+
     // observable state store
     this.store = new ComposableObservableStore({
       state: initState,
@@ -3480,7 +3483,15 @@ export default class MetamaskController extends EventEmitter {
     this.emit('controllerConnectionChanged', this.activeControllerConnections);
 
     // set up postStream transport
-    outStream.on('data', createMetaRPCHandler(api, outStream));
+    outStream.on(
+      'data',
+      createMetaRPCHandler(
+        api,
+        outStream,
+        this.store,
+        this.localStoreApiWrapper,
+      ),
+    );
     const handleUpdate = (update) => {
       if (outStream._writableState.ended) {
         return;