Simplify MV3 initialization (#16559)

* Simplify MV3 initialization

The MV3 initialization logic was complicated and introduced race
difficult-to-reproduce race conditions when dapps connect during
initialization.

It seems that problems were encountered after the UI tried to connect
before the background was initialized. To address this, the
initialization step was _delayed_ until after the first connection.
That first connection was then passed into the initialization function,
and setup properly after initialization had begun.

However, this special treatment is only given for the first connection.
Subsequent connections that still occur during initialization would
fail. This also results in the initialization being needlessly delayed,
which is concerning given that our main performance goal is to speed it
up.

* Setup connect listeners before controller initialization

* Add comments

* Add comment explaining isInitialized step

app/scripts/background.js

@@ -47,7 +47,7 @@ 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';
-import { getPlatform } from './lib/util';
+import { deferredPromise, getPlatform } from './lib/util';
 /* eslint-enable import/first */
 
 const { sentry } = global;
@@ -94,16 +94,17 @@ const ACK_KEEP_ALIVE_MESSAGE = 'ACK_KEEP_ALIVE_MESSAGE';
 const WORKER_KEEP_ALIVE_MESSAGE = 'WORKER_KEEP_ALIVE_MESSAGE';
 
 /**
- * In case of MV3 we attach a "onConnect" event listener as soon as the application is initialised.
- * Reason is that in case of MV3 a delay in doing this was resulting in missing first connect event after service worker is re-activated.
+ * This deferred Promise is used to track whether initialization has finished.
  *
- * @param remotePort
+ * It is very important to ensure that `resolveInitialization` is *always*
+ * called once initialization has completed, and that `rejectInitialization` is
+ * called if initialization fails in an unrecoverable way.
  */
-const initApp = async (remotePort) => {
-  browser.runtime.onConnect.removeListener(initApp);
-  await initialize(remotePort);
-  log.info('MetaMask initialization complete.');
-};
+const {
+  promise: isInitialized,
+  resolve: resolveInitialization,
+  reject: rejectInitialization,
+} = deferredPromise();
 
 /**
  * Sends a message to the dapp(s) content script to signal it can connect to MetaMask background as
@@ -154,13 +155,26 @@ const sendReadyMessageToTabs = async () => {
   }
 };
 
-if (isManifestV3) {
-  browser.runtime.onConnect.addListener(initApp);
-  sendReadyMessageToTabs();
-} else {
-  // initialization flow
-  initialize().catch(log.error);
-}
+// These are set after initialization
+let connectRemote;
+let connectExternal;
+
+browser.runtime.onConnect.addListener(async (...args) => {
+  // Queue up connection attempts here, waiting until after initialization
+  await isInitialized;
+
+  // This is set in `setupController`, which is called as part of initialization
+  connectRemote(...args);
+});
+browser.runtime.onConnectExternal.addListener(async (...args) => {
+  // Queue up connection attempts here, waiting until after initialization
+  await isInitialized;
+
+  // This is set in `setupController`, which is called as part of initialization
+  connectExternal(...args);
+});
+
+initialize().catch(log.error);
 
 /**
  * @typedef {import('../../shared/constants/transaction').TransactionMeta} TransactionMeta
@@ -220,17 +234,22 @@ if (isManifestV3) {
 /**
  * Initializes the MetaMask controller, and sets up all platform configuration.
  *
- * @param {string} remotePort - remote application port connecting to extension.
  * @returns {Promise} Setup complete.
  */
-async function initialize(remotePort) {
-  const initState = await loadStateFromPersistence();
-  const initLangCode = await getFirstPreferredLangCode();
-  setupController(initState, initLangCode, remotePort);
-  if (!isManifestV3) {
-    await loadPhishingWarningPage();
+async function initialize() {
+  try {
+    const initState = await loadStateFromPersistence();
+    const initLangCode = await getFirstPreferredLangCode();
+    setupController(initState, initLangCode);
+    if (!isManifestV3) {
+      await loadPhishingWarningPage();
+    }
+    await sendReadyMessageToTabs();
+    log.info('MetaMask initialization complete.');
+    resolveInitialization();
+  } catch (error) {
+    rejectInitialization(error);
   }
-  log.info('MetaMask initialization complete.');
 }
 
 /**
@@ -362,9 +381,8 @@ 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 {string} remoteSourcePort - remote application port connecting to extension.
  */
-function setupController(initState, initLangCode, remoteSourcePort) {
+function setupController(initState, initLangCode) {
   //
   // MetaMask Controller
   //
@@ -413,16 +431,6 @@ function setupController(initState, initLangCode, remoteSourcePort) {
 
   setupSentryGetStateGlobal(controller);
 
-  //
-  // connect to other contexts
-  //
-  if (isManifestV3 && remoteSourcePort) {
-    connectRemote(remoteSourcePort);
-  }
-
-  browser.runtime.onConnect.addListener(connectRemote);
-  browser.runtime.onConnectExternal.addListener(connectExternal);
-
   const isClientOpenStatus = () => {
     return (
       popupIsOpen ||
@@ -463,7 +471,7 @@ function setupController(initState, initLangCode, remoteSourcePort) {
    *
    * @param {Port} remotePort - The port provided by a new context.
    */
-  function connectRemote(remotePort) {
+  connectRemote = async (remotePort) => {
     const processName = remotePort.name;
 
     if (metamaskBlockedPorts.includes(remotePort.name)) {
@@ -561,16 +569,16 @@ function setupController(initState, initLangCode, remoteSourcePort) {
       }
       connectExternal(remotePort);
     }
-  }
+  };
 
   // communication with page or other extension
-  function connectExternal(remotePort) {
+  connectExternal = (remotePort) => {
     const portStream = new PortStream(remotePort);
     controller.setupUntrustedCommunication({
       connectionStream: portStream,
       sender: remotePort.sender,
     });
-  }
+  };
 
   //
   // User Interface setup

app/scripts/lib/util.js

@@ -191,3 +191,30 @@ export const generateRandomId = () => {
 export const isValidDate = (d) => {
   return d instanceof Date && !isNaN(d);
 };
+
+/**
+ * A deferred Promise.
+ *
+ * A deferred Promise is one that can be resolved or rejected independently of
+ * the Promise construction.
+ *
+ * @typedef {object} DeferredPromise
+ * @property {Promise} promise - The Promise that has been deferred.
+ * @property {() => void} resolve - A function that resolves the Promise.
+ * @property {() => void} reject - A function that rejects the Promise.
+ */
+
+/**
+ * Create a defered Promise.
+ *
+ * @returns {DeferredPromise} A deferred Promise.
+ */
+export function deferredPromise() {
+  let resolve;
+  let reject;
+  const promise = new Promise((innerResolve, innerReject) => {
+    resolve = innerResolve;
+    reject = innerReject;
+  });
+  return { promise, resolve, reject };
+}

app/scripts/lib/util.test.js

@@ -9,7 +9,7 @@ import {
   PLATFORM_CHROME,
   PLATFORM_EDGE,
 } from '../../../shared/constants/app';
-import { getEnvironmentType, getPlatform } from './util';
+import { deferredPromise, getEnvironmentType, getPlatform } from './util';
 
 describe('app utils', () => {
   describe('getEnvironmentType', () => {
@@ -154,4 +154,58 @@ describe('app utils', () => {
       expect(getPlatform()).toStrictEqual(PLATFORM_CHROME);
     });
   });
+
+  describe('deferredPromise', () => {
+    it('should allow rejecting a deferred Promise', async () => {
+      const { promise, reject } = deferredPromise();
+
+      reject(new Error('test'));
+
+      await expect(promise).rejects.toThrow('test');
+    });
+
+    it('should allow resolving a deferred Promise', async () => {
+      const { promise, resolve } = deferredPromise();
+
+      resolve('test');
+
+      await expect(promise).resolves.toBe('test');
+    });
+
+    it('should still be rejected after reject is called twice', async () => {
+      const { promise, reject } = deferredPromise();
+
+      reject(new Error('test'));
+      reject(new Error('different message'));
+
+      await expect(promise).rejects.toThrow('test');
+    });
+
+    it('should still be rejected after resolve is called post-rejection', async () => {
+      const { promise, resolve, reject } = deferredPromise();
+
+      reject(new Error('test'));
+      resolve('different message');
+
+      await expect(promise).rejects.toThrow('test');
+    });
+
+    it('should still be resolved after resolve is called twice', async () => {
+      const { promise, resolve } = deferredPromise();
+
+      resolve('test');
+      resolve('different message');
+
+      await expect(promise).resolves.toBe('test');
+    });
+
+    it('should still be resolved after reject is called post-resolution', async () => {
+      const { promise, resolve, reject } = deferredPromise();
+
+      resolve('test');
+      reject(new Error('different message'));
+
+      await expect(promise).resolves.toBe('test');
+    });
+  });
 });

app/scripts/metamask-controller.test.js

@@ -12,7 +12,7 @@ import createTxMeta from '../../test/lib/createTxMeta';
 import { NETWORK_TYPES } from '../../shared/constants/network';
 import { KEYRING_TYPES } from '../../shared/constants/keyrings';
 import { DEVICE_NAMES } from '../../shared/constants/hardware-wallets';
-import { addHexPrefix } from './lib/util';
+import { addHexPrefix, deferredPromise } from './lib/util';
 
 const Ganache = require('../../test/e2e/ganache');
 
@@ -1238,11 +1238,3 @@ describe('MetaMaskController', function () {
     });
   });
 });
-
-function deferredPromise() {
-  let resolve;
-  const promise = new Promise((_resolve) => {
-    resolve = _resolve;
-  });
-  return { promise, resolve };
-}