Keep memstore contents after service worker restarts (#15913)

* Add all controllers in memstore to store
Add methods to controller to reset memstore
Reset memstore when popup or tab is closed.

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* When profile is loaded, set isFirstTime to true..
After resetting the controllers, set the flag to false.

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Remove console.logs

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* For some reason programmatically computing the store is not working.

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Proper check for browser.storage

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* do a list of rest methods instead of reset controllers.

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Mock controller resetStates and localstore get/set

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Comments about TLC

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* bind this.

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* use globalThis instead of locastore to store first time state.

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Test to check that resetStates is not called a second time

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Set init state in GasFeeController and other controllers so that their
state is persisted accross SW restarts

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Revert localstore changes

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* wrap the reset states changes in MV3 flag

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Remove localstore from metamask-controller

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Always reset state on MMController start in MV2.

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Use relative path for import of isManifestV3

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

* Fix unit test

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

Signed-off-by: Akintayo A. Olusegun <akintayo.segun@gmail.com>

app/scripts/app-init.js

@@ -127,6 +127,10 @@ chrome.runtime.onMessage.addListener(() => {
   return false;
 });
 
+chrome.runtime.onStartup.addListener(() => {
+  globalThis.isFirstTimeProfileLoaded = true;
+});
+
 /*
  * This content script is injected programmatically because
  * MAIN world injection does not work properly via manifest

app/scripts/controllers/ens/index.js

@@ -27,6 +27,11 @@ export default class EnsController {
     }
 
     this.store = new ObservableStore(initState);
+
+    this.resetState = () => {
+      this.store.updateState(initState);
+    };
+
     onNetworkDidChange(() => {
       this.store.putState(initState);
       const chainId = getCurrentChainId();

app/scripts/controllers/swaps.js

@@ -115,6 +115,10 @@ export default class SwapsController {
       swapsState: { ...initialState.swapsState },
     });
 
+    this.resetState = () => {
+      this.store.updateState({ swapsState: { ...initialState.swapsState } });
+    };
+
     this._fetchTradesInfo = fetchTradesInfo;
     this._getCurrentChainId = getCurrentChainId;
     this._getEIP1559GasFeeEstimates = getEIP1559GasFeeEstimates;

app/scripts/controllers/transactions/index.js

@@ -151,6 +151,11 @@ export default class TransactionController extends EventEmitter {
     this.getTokenStandardAndDetails = opts.getTokenStandardAndDetails;
 
     this.memStore = new ObservableStore({});
+
+    this.resetState = () => {
+      this._updateMemstore();
+    };
+
     this.query = new EthQuery(this.provider);
 
     this.txGasUtil = new TxGasUtil(this.provider);

app/scripts/lib/account-tracker.js

@@ -62,6 +62,10 @@ export default class AccountTracker {
     };
     this.store = new ObservableStore(initState);
 
+    this.resetState = () => {
+      this.store.updateState(initState);
+    };
+
     this._provider = opts.provider;
     this._query = pify(new EthQuery(this._provider));
     this._blockTracker = opts.blockTracker;

app/scripts/lib/decrypt-message-manager.js

@@ -41,6 +41,14 @@ export default class DecryptMessageManager extends EventEmitter {
       unapprovedDecryptMsgs: {},
       unapprovedDecryptMsgCount: 0,
     });
+
+    this.resetState = () => {
+      this.memStore.updateState({
+        unapprovedDecryptMsgs: {},
+        unapprovedDecryptMsgCount: 0,
+      });
+    };
+
     this.messages = [];
     this.metricsEvent = opts.metricsEvent;
   }

app/scripts/lib/encryption-public-key-manager.js

@@ -36,6 +36,14 @@ export default class EncryptionPublicKeyManager extends EventEmitter {
       unapprovedEncryptionPublicKeyMsgs: {},
       unapprovedEncryptionPublicKeyMsgCount: 0,
     });
+
+    this.resetState = () => {
+      this.memStore.updateState({
+        unapprovedEncryptionPublicKeyMsgs: {},
+        unapprovedEncryptionPublicKeyMsgCount: 0,
+      });
+    };
+
     this.messages = [];
     this.metricsEvent = opts.metricsEvent;
   }

app/scripts/lib/message-manager.js

@@ -36,6 +36,14 @@ export default class MessageManager extends EventEmitter {
       unapprovedMsgs: {},
       unapprovedMsgCount: 0,
     });
+
+    this.resetState = () => {
+      this.memStore.updateState({
+        unapprovedMsgs: {},
+        unapprovedMsgCount: 0,
+      });
+    };
+
     this.messages = [];
     this.metricsEvent = metricsEvent;
   }

app/scripts/lib/personal-message-manager.js

@@ -43,6 +43,14 @@ export default class PersonalMessageManager extends EventEmitter {
       unapprovedPersonalMsgs: {},
       unapprovedPersonalMsgCount: 0,
     });
+
+    this.resetState = () => {
+      this.memStore.updateState({
+        unapprovedPersonalMsgs: {},
+        unapprovedPersonalMsgCount: 0,
+      });
+    };
+
     this.messages = [];
     this.metricsEvent = metricsEvent;
   }

app/scripts/lib/typed-message-manager.js

@@ -43,6 +43,14 @@ export default class TypedMessageManager extends EventEmitter {
       unapprovedTypedMessages: {},
       unapprovedTypedMessagesCount: 0,
     });
+
+    this.resetState = () => {
+      this.memStore.updateState({
+        unapprovedTypedMessages: {},
+        unapprovedTypedMessagesCount: 0,
+      });
+    };
+
     this.messages = [];
     this.metricsEvent = metricsEvent;
   }

app/scripts/metamask-controller.actions.test.js

@@ -23,6 +23,12 @@ const browserPolyfillMock = {
     },
     getPlatformInfo: async () => 'mac',
   },
+  storage: {
+    local: {
+      get: sinon.stub().resolves({}),
+      set: sinon.stub().resolves(),
+    },
+  },
 };
 
 let loggerMiddlewareMock;

app/scripts/metamask-controller.js

@@ -100,6 +100,7 @@ import {
   getTokenValueParam,
   hexToDecimal,
 } from '../../shared/lib/metamask-controller-utils';
+import { isManifestV3 } from '../../shared/modules/mv3.utils';
 import {
   onMessageReceived,
   checkForMultipleVersionsRunning,
@@ -417,6 +418,7 @@ export default class MetamaskController extends EventEmitter {
       : GAS_API_BASE_URL;
 
     this.gasFeeController = new GasFeeController({
+      state: initState.GasFeeController,
       interval: 10000,
       messenger: gasFeeMessenger,
       clientId: SWAPS_CLIENT_ID,
@@ -481,26 +483,30 @@ export default class MetamaskController extends EventEmitter {
     );
 
     // token exchange rate tracker
-    this.tokenRatesController = new TokenRatesController({
+    this.tokenRatesController = new TokenRatesController(
-      onTokensStateChange: (listener) =>
+      {
-        this.tokensController.subscribe(listener),
+        onTokensStateChange: (listener) =>
-      onCurrencyRateStateChange: (listener) =>
+          this.tokensController.subscribe(listener),
-        this.controllerMessenger.subscribe(
+        onCurrencyRateStateChange: (listener) =>
-          `${this.currencyRateController.name}:stateChange`,
+          this.controllerMessenger.subscribe(
-          listener,
+            `${this.currencyRateController.name}:stateChange`,
-        ),
+            listener,
-      onNetworkStateChange: (cb) =>
+          ),
-        this.networkController.store.subscribe((networkState) => {
+        onNetworkStateChange: (cb) =>
-          const modifiedNetworkState = {
+          this.networkController.store.subscribe((networkState) => {
-            ...networkState,
+            const modifiedNetworkState = {
-            provider: {
+              ...networkState,
-              ...networkState.provider,
+              provider: {
-              chainId: hexToDecimal(networkState.provider.chainId),
+                ...networkState.provider,
-            },
+                chainId: hexToDecimal(networkState.provider.chainId),
-          };
+              },
-          return cb(modifiedNetworkState);
+            };
-        }),
+            return cb(modifiedNetworkState);
-    });
+          }),
+      },
+      undefined,
+      initState.TokenRatesController,
+    );
 
     this.ensController = new EnsController({
       provider: this.provider,
@@ -719,6 +725,7 @@ export default class MetamaskController extends EventEmitter {
     });
 
     this.rateLimitController = new RateLimitController({
+      state: initState.RateLimitController,
       messenger: this.controllerMessenger.getRestricted({
         name: 'RateLimitController',
       }),
@@ -1029,6 +1036,24 @@ export default class MetamaskController extends EventEmitter {
     // ensure isClientOpenAndUnlocked is updated when memState updates
     this.on('update', (memState) => this._onStateUpdate(memState));
 
+    /**
+     * All controllers in Memstore but not in store. They are not persisted.
+     * On chrome profile re-start, they will be re-initialized.
+     */
+    const resetOnRestartStore = {
+      AccountTracker: this.accountTracker.store,
+      TxController: this.txController.memStore,
+      TokenRatesController: this.tokenRatesController,
+      MessageManager: this.messageManager.memStore,
+      PersonalMessageManager: this.personalMessageManager.memStore,
+      DecryptMessageManager: this.decryptMessageManager.memStore,
+      EncryptionPublicKeyManager: this.encryptionPublicKeyManager.memStore,
+      TypesMessageManager: this.typedMessageManager.memStore,
+      SwapsController: this.swapsController.store,
+      EnsController: this.ensController.store,
+      ApprovalController: this.approvalController,
+    };
+
     this.store.updateStructure({
       AppStateController: this.appStateController.store,
       TransactionController: this.txController.store,
@@ -1057,21 +1082,14 @@ export default class MetamaskController extends EventEmitter {
       CronjobController: this.cronjobController,
       NotificationController: this.notificationController,
       ///: END:ONLY_INCLUDE_IN
+      ...resetOnRestartStore,
     });
 
     this.memStore = new ComposableObservableStore({
       config: {
         AppStateController: this.appStateController.store,
         NetworkController: this.networkController.store,
-        AccountTracker: this.accountTracker.store,
-        TxController: this.txController.memStore,
         CachedBalancesController: this.cachedBalancesController.store,
-        TokenRatesController: this.tokenRatesController,
-        MessageManager: this.messageManager.memStore,
-        PersonalMessageManager: this.personalMessageManager.memStore,
-        DecryptMessageManager: this.decryptMessageManager.memStore,
-        EncryptionPublicKeyManager: this.encryptionPublicKeyManager.memStore,
-        TypesMessageManager: this.typedMessageManager.memStore,
         KeyringController: this.keyringController.memStore,
         PreferencesController: this.preferencesController.store,
         MetaMetricsController: this.metaMetricsController.store,
@@ -1085,9 +1103,6 @@ export default class MetamaskController extends EventEmitter {
         PermissionLogController: this.permissionLogController.store,
         SubjectMetadataController: this.subjectMetadataController,
         BackupController: this.backupController,
-        SwapsController: this.swapsController.store,
-        EnsController: this.ensController.store,
-        ApprovalController: this.approvalController,
         AnnouncementController: this.announcementController,
         GasFeeController: this.gasFeeController,
         TokenListController: this.tokenListController,
@@ -1099,11 +1114,36 @@ export default class MetamaskController extends EventEmitter {
         CronjobController: this.cronjobController,
         NotificationController: this.notificationController,
         ///: END:ONLY_INCLUDE_IN
+        ...resetOnRestartStore,
       },
       controllerMessenger: this.controllerMessenger,
     });
     this.memStore.subscribe(this.sendUpdate.bind(this));
 
+    // if this is the first time, clear the state of by calling these methods
+    const resetMethods = [
+      this.accountTracker.resetState,
+      this.txController.resetState,
+      this.messageManager.resetState,
+      this.personalMessageManager.resetState,
+      this.decryptMessageManager.resetState,
+      this.encryptionPublicKeyManager.resetState,
+      this.typedMessageManager.resetState,
+      this.swapsController.resetState,
+      this.ensController.resetState,
+      this.approvalController.clear.bind(this.approvalController),
+      // WE SHOULD ADD TokenListController.resetState here too. But it's not implemented yet.
+    ];
+
+    if (isManifestV3) {
+      if (globalThis.isFirstTimeProfileLoaded === true) {
+        this.resetStates(resetMethods);
+      }
+    } else {
+      // it's always the first time in MV2
+      this.resetStates(resetMethods);
+    }
+
     const password = process.env.CONF?.PASSWORD;
     if (
       password &&
@@ -1137,6 +1177,18 @@ export default class MetamaskController extends EventEmitter {
     checkForMultipleVersionsRunning();
   }
 
+  resetStates(resetMethods) {
+    resetMethods.forEach((resetMethod) => {
+      try {
+        resetMethod();
+      } catch (err) {
+        console.error(err);
+      }
+    });
+
+    globalThis.isFirstTimeProfileLoaded = false;
+  }
+
   ///: BEGIN:ONLY_INCLUDE_IN(flask)
   /**
    * Constructor helper for getting Snap permission specifications.

app/scripts/metamask-controller.test.js

@@ -102,11 +102,14 @@ const CUSTOM_RPC_CHAIN_ID = '0x539';
 
 describe('MetaMaskController', function () {
   let metamaskController;
+
   const sandbox = sinon.createSandbox();
   const noop = () => undefined;
 
   before(async function () {
+    globalThis.isFirstTimeProfileLoaded = true;
     await ganacheServer.start();
+    sinon.spy(MetaMaskController.prototype, 'resetStates');
   });
 
   beforeEach(function () {
@@ -160,6 +163,18 @@ describe('MetaMaskController', function () {
     await ganacheServer.quit();
   });
 
+  describe('should reset states on first time profile load', function () {
+    it('should reset state', function () {
+      assert(metamaskController.resetStates.calledOnce);
+      assert.equal(globalThis.isFirstTimeProfileLoaded, false);
+    });
+
+    it('should not reset states if already set', function () {
+      // global.isFirstTime should also remain false
+      assert.equal(globalThis.isFirstTimeProfileLoaded, false);
+    });
+  });
+
   describe('#getAccounts', function () {
     it('returns first address when dapp calls web3.eth.getAccounts', async function () {
       const password = 'a-fake-password';