Fix firsttimeloaded logic (#18344)

* use session storage, instead of chrome.runtime.onStartup and globalThis, for firstTimeLoaded architecture

* Ensure account tracker accounts remain defined upon service worker restart

* lint fix

* Simplify code

* Only call browser.storage.session in mv3

* Only call browser.storage.session.set after resetStates in mv3

* fix metamask controller reset states unit tests

* fix test

* fix test

* Actually fix tests

* lint fix

app/scripts/app-init.js

@@ -127,10 +127,6 @@ 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/background.js

@@ -268,7 +268,23 @@ async function initialize() {
     await DesktopManager.init(platform.getVersion());
     ///: END:ONLY_INCLUDE_IN
 
-    setupController(initState, initLangCode);
+    let isFirstMetaMaskControllerSetup;
+    if (isManifestV3) {
+      const sessionData = await browser.storage.session.get([
+        'isFirstMetaMaskControllerSetup',
+      ]);
+
+      isFirstMetaMaskControllerSetup =
+        sessionData?.isFirstMetaMaskControllerSetup === undefined;
+      await browser.storage.session.set({ isFirstMetaMaskControllerSetup });
+    }
+
+    setupController(
+      initState,
+      initLangCode,
+      {},
+      isFirstMetaMaskControllerSetup,
+    );
     if (!isManifestV3) {
       await loadPhishingWarningPage();
     }
@@ -410,8 +426,14 @@ export async function loadStateFromPersistence() {
  * @param {object} initState - The initial state to start the controller with, matches the state that is emitted from the controller.
  * @param {string} initLangCode - The region code for the language preferred by the current user.
  * @param {object} overrides - object with callbacks that are allowed to override the setup controller logic (usefull for desktop app)
+ * @param isFirstMetaMaskControllerSetup
  */
-export function setupController(initState, initLangCode, overrides) {
+export function setupController(
+  initState,
+  initLangCode,
+  overrides,
+  isFirstMetaMaskControllerSetup,
+) {
   //
   // MetaMask Controller
   //
@@ -437,6 +459,7 @@ export function setupController(initState, initLangCode, overrides) {
     },
     localStore,
     overrides,
+    isFirstMetaMaskControllerSetup,
   });
 
   setupEnsIpfsResolver({

app/scripts/lib/account-tracker.js

@@ -62,7 +62,7 @@ export default class AccountTracker {
       accounts: {},
       currentBlockGasLimit: '',
     };
-    this.store = new ObservableStore(initState);
+    this.store = new ObservableStore({ ...initState, ...opts.initState });
 
     this.resetState = () => {
       this.store.updateState(initState);

app/scripts/metamask-controller.js

@@ -62,7 +62,6 @@ import {
 } from '@metamask/snaps-controllers';
 ///: END:ONLY_INCLUDE_IN
 
-import browser from 'webextension-polyfill';
 import {
   AssetType,
   TransactionStatus,
@@ -204,6 +203,8 @@ export default class MetamaskController extends EventEmitter {
   constructor(opts) {
     super();
 
+    const { isFirstMetaMaskControllerSetup } = opts;
+
     this.defaultMaxListeners = 20;
 
     this.sendUpdate = debounce(
@@ -638,6 +639,12 @@ export default class MetamaskController extends EventEmitter {
       },
       preferencesController: this.preferencesController,
       onboardingController: this.onboardingController,
+      initState:
+        isManifestV3 &&
+        isFirstMetaMaskControllerSetup === false &&
+        initState.AccountTracker?.accounts
+          ? { accounts: initState.AccountTracker.accounts }
+          : { accounts: {} },
     });
 
     // start and stop polling for balances based on activeControllerConnections
@@ -1369,8 +1376,11 @@ export default class MetamaskController extends EventEmitter {
     ];
 
     if (isManifestV3) {
-      if (globalThis.isFirstTimeProfileLoaded === true) {
+      if (isFirstMetaMaskControllerSetup === true) {
         this.resetStates(resetMethods);
+        this.extension.storage.session.set({
+          isFirstMetaMaskControllerSetup: false,
+        });
       }
     } else {
       // it's always the first time in MV2
@@ -1445,8 +1455,6 @@ export default class MetamaskController extends EventEmitter {
         console.error(err);
       }
     });
-
-    globalThis.isFirstTimeProfileLoaded = false;
   }
 
   ///: BEGIN:ONLY_INCLUDE_IN(flask)
@@ -2679,10 +2687,8 @@ export default class MetamaskController extends EventEmitter {
    */
   async submitEncryptionKey() {
     try {
-      const { loginToken, loginSalt } = await browser.storage.session.get([
+      const { loginToken, loginSalt } =
-        'loginToken',
+        await this.extension.storage.session.get(['loginToken', 'loginSalt']);
-        'loginSalt',
-      ]);
       if (loginToken && loginSalt) {
         const { vault } = this.keyringController.store.getState();
 
@@ -2707,7 +2713,7 @@ export default class MetamaskController extends EventEmitter {
   }
 
   async clearLoginArtifacts() {
-    await browser.storage.session.remove(['loginToken', 'loginSalt']);
+    await this.extension.storage.session.remove(['loginToken', 'loginSalt']);
   }
 
   /**
@@ -4125,7 +4131,7 @@ export default class MetamaskController extends EventEmitter {
     );
 
     if (isManifestV3) {
-      await browser.storage.session.set({ loginToken, loginSalt });
+      await this.extension.storage.session.set({ loginToken, loginSalt });
     }
 
     if (!addresses.length) {

app/scripts/metamask-controller.test.js

@@ -30,6 +30,9 @@ const browserPolyfillMock = {
     },
     getPlatformInfo: async () => 'mac',
   },
+  storage: {
+    session: {},
+  },
 };
 
 let loggerMiddlewareMock;
@@ -79,6 +82,10 @@ const MetaMaskController = proxyquire('./metamask-controller', {
   'ethjs-contract': MockEthContract,
 }).default;
 
+const MetaMaskControllerMV3 = proxyquire('./metamask-controller', {
+  '../../shared/modules/mv3.utils': { isManifestV3: true },
+}).default;
+
 const currentNetworkId = '5';
 const DEFAULT_LABEL = 'Account 1';
 const TEST_SEED =
@@ -168,10 +175,13 @@ describe('MetaMaskController', function () {
   const sandbox = sinon.createSandbox();
   const noop = () => undefined;
 
+  browserPolyfillMock.storage.session.set = sandbox.spy();
+
   before(async function () {
     globalThis.isFirstTimeProfileLoaded = true;
     await ganacheServer.start();
     sinon.spy(MetaMaskController.prototype, 'resetStates');
+    sinon.spy(MetaMaskControllerMV3.prototype, 'resetStates');
   });
 
   beforeEach(function () {
@@ -224,6 +234,7 @@ describe('MetaMaskController', function () {
       },
       browser: browserPolyfillMock,
       infuraProjectId: 'foo',
+      isFirstMetaMaskControllerSetup: true,
     });
 
     // add sinon method spies
@@ -247,14 +258,70 @@ describe('MetaMaskController', function () {
   });
 
   describe('should reset states on first time profile load', function () {
-    it('should reset state', function () {
+    it('in mv2, it should reset state without attempting to call browser storage', function () {
-      assert(metamaskController.resetStates.calledOnce);
+      assert.equal(metamaskController.resetStates.callCount, 1);
-      assert.equal(globalThis.isFirstTimeProfileLoaded, false);
+      assert.equal(browserPolyfillMock.storage.session.set.callCount, 0);
     });
 
-    it('should not reset states if already set', function () {
+    it('in mv3, it should reset state', function () {
-      // global.isFirstTime should also remain false
+      MetaMaskControllerMV3.prototype.resetStates.resetHistory();
-      assert.equal(globalThis.isFirstTimeProfileLoaded, false);
+      const metamaskControllerMV3 = new MetaMaskControllerMV3({
+        showUserConfirmation: noop,
+        encryptor: {
+          encrypt(_, object) {
+            this.object = object;
+            return Promise.resolve('mock-encrypted');
+          },
+          decrypt() {
+            return Promise.resolve(this.object);
+          },
+        },
+        initState: cloneDeep(firstTimeState),
+        initLangCode: 'en_US',
+        platform: {
+          showTransactionNotification: () => undefined,
+          getVersion: () => 'foo',
+        },
+        browser: browserPolyfillMock,
+        infuraProjectId: 'foo',
+        isFirstMetaMaskControllerSetup: true,
+      });
+      assert.equal(metamaskControllerMV3.resetStates.callCount, 1);
+      assert.equal(browserPolyfillMock.storage.session.set.callCount, 1);
+      assert.deepEqual(
+        browserPolyfillMock.storage.session.set.getCall(0).args[0],
+        {
+          isFirstMetaMaskControllerSetup: false,
+        },
+      );
+    });
+
+    it('in mv3, it should not reset states if isFirstMetaMaskControllerSetup is false', function () {
+      MetaMaskControllerMV3.prototype.resetStates.resetHistory();
+      browserPolyfillMock.storage.session.set.resetHistory();
+      const metamaskControllerMV3 = new MetaMaskControllerMV3({
+        showUserConfirmation: noop,
+        encryptor: {
+          encrypt(_, object) {
+            this.object = object;
+            return Promise.resolve('mock-encrypted');
+          },
+          decrypt() {
+            return Promise.resolve(this.object);
+          },
+        },
+        initState: cloneDeep(firstTimeState),
+        initLangCode: 'en_US',
+        platform: {
+          showTransactionNotification: () => undefined,
+          getVersion: () => 'foo',
+        },
+        browser: browserPolyfillMock,
+        infuraProjectId: 'foo',
+        isFirstMetaMaskControllerSetup: false,
+      });
+      assert.equal(metamaskControllerMV3.resetStates.callCount, 0);
+      assert.equal(browserPolyfillMock.storage.session.set.callCount, 0);
     });
   });