MV3: Re-activate service worker and reconnect UI streams (#14781)

app/scripts/app-init.js

@@ -1,8 +1,9 @@
+/* global chrome */
 // This file is used only for manifest version 3
 
 // Represents if importAllScripts has been run
 // eslint-disable-next-line
-let scriptsLoaded = false;
+let scriptsLoadInitiated = false;
 
 // Variable testMode is set to true when preparing test build.
 // This helps in changing service worker execution in test environment.
@@ -35,10 +36,10 @@ function tryImport(...fileNames) {
 
 function importAllScripts() {
   // Bail if we've already imported scripts
-  if (scriptsLoaded) {
+  if (scriptsLoadInitiated) {
     return;
   }
-
+  scriptsLoadInitiated = true;
   const files = [];
 
   // In testMode individual files are imported, this is to help capture load time stats
@@ -69,9 +70,6 @@ function importAllScripts() {
     loadFile('./runtime-cjs.js');
   }
 
-  // Mark scripts as loaded
-  scriptsLoaded = true;
-
   const fileList = [
     // The list of files is injected at build time by replacing comment below with comma separated strings of file names
     // https://github.com/MetaMask/metamask-extension/blob/496d9d81c3367931031edc11402552690c771acf/development/build/scripts.js#L406
@@ -111,14 +109,20 @@ function importAllScripts() {
   }
 }
 
+// Ref: https://stackoverflow.com/questions/66406672/chrome-extension-mv3-modularize-service-worker-js-file
 // eslint-disable-next-line no-undef
 self.addEventListener('install', importAllScripts);
 
 /*
- * Message event listener below loads script if they are no longer available.
+ * A keepalive message listener to prevent Service Worker getting shut down due to inactivity.
+ * UI sends the message periodically, in a setInterval.
+ * Chrome will revive the service worker if it was shut down, whenever a new message is sent, but only if a listener was defined here.
+ *
  * chrome below needs to be replaced by cross-browser object,
  * but there is issue in importing webextension-polyfill into service worker.
  * chrome does seems to work in at-least all chromium based browsers
  */
-// eslint-disable-next-line no-undef
-chrome.runtime.onMessage.addListener(importAllScripts);
+chrome.runtime.onMessage.addListener(() => {
+  importAllScripts();
+  return false;
+});

app/scripts/ui.js

@@ -11,7 +11,7 @@ import Eth from 'ethjs';
 import EthQuery from 'eth-query';
 import StreamProvider from 'web3-stream-provider';
 import log from 'loglevel';
-import launchMetaMaskUi from '../../ui';
+import launchMetaMaskUi, { updateBackgroundConnection } from '../../ui';
 import {
   ENVIRONMENT_TYPE_FULLSCREEN,
   ENVIRONMENT_TYPE_POPUP,
@@ -24,34 +24,37 @@ import { setupMultiplex } from './lib/stream-utils';
 import { getEnvironmentType } from './lib/util';
 import metaRPCClientFactory from './lib/metaRPCClientFactory';
 
-start().catch(log.error);
-
-async function start() {
-  async function displayCriticalError(container, err, metamaskState) {
-    const html = await getErrorHtml(SUPPORT_LINK, metamaskState);
-
-    container.innerHTML = html;
+const container = document.getElementById('app-content');
 
-    const button = document.getElementById('critical-error-button');
+const WORKER_KEEP_ALIVE_INTERVAL = 1000;
+const WORKER_KEEP_ALIVE_MESSAGE = 'WORKER_KEEP_ALIVE_MESSAGE';
 
-    button.addEventListener('click', (_) => {
-      browser.runtime.reload();
-    });
+/*
+ * As long as UI is open it will keep sending messages to service worker
+ * In service worker as this message is received
+ * if service worker is inactive it is reactivated and script re-loaded
+ * Time has been kept to 1000ms but can be reduced for even faster re-activation of service worker
+ */
+if (isManifestV3()) {
+  setInterval(() => {
+    browser.runtime.sendMessage({ name: WORKER_KEEP_ALIVE_MESSAGE });
+  }, WORKER_KEEP_ALIVE_INTERVAL);
+}
 
-    log.error(err.stack);
-    throw err;
-  }
+start().catch(log.error);
 
+async function start() {
   // create platform global
   global.platform = new ExtensionPlatform();
 
   // identify window type (popup, notification)
   const windowType = getEnvironmentType();
 
-  // setup stream to background
-  const extensionPort = browser.runtime.connect({ name: windowType });
+  let isUIInitialised = false;
 
-  const connectionStream = new PortStream(extensionPort);
+  // setup stream to background
+  let extensionPort = browser.runtime.connect({ name: windowType });
+  let connectionStream = new PortStream(extensionPort);
 
   const activeTab = await queryCurrentActiveTab(windowType);
 
@@ -60,23 +63,52 @@ async function start() {
    * Code below ensures that UI is rendered only after background is ready.
    */
   if (isManifestV3()) {
-    extensionPort.onMessage.addListener((message) => {
+    /*
+     * In case of MV3 the issue of blank screen was very frequent, it is caused by UI initialising before background is ready to send state.
+     * Code below ensures that UI is rendered only after CONNECTION_READY message is received thus background is ready.
+     * In case the UI is already rendered, only update the streams.
+     */
+    const messageListener = (message) => {
       if (message?.name === 'CONNECTION_READY') {
-        initializeUiWithTab(activeTab);
+        if (isUIInitialised) {
+          updateUiStreams();
+        } else {
+          initializeUiWithTab(activeTab);
+        }
       }
-    });
+    };
+
+    // resetExtensionStreamAndListeners takes care to remove listeners from closed streams
+    // it also creates new streams and attach event listeners to them
+    const resetExtensionStreamAndListeners = () => {
+      extensionPort.onMessage.removeListener(messageListener);
+      extensionPort.onDisconnect.removeListener(
+        resetExtensionStreamAndListeners,
+      );
+      // message below will try to activate service worker
+      // in MV3 is likely that reason of stream closing is service worker going in-active
+      browser.runtime.sendMessage({ name: WORKER_KEEP_ALIVE_MESSAGE });
+
+      extensionPort = browser.runtime.connect({ name: windowType });
+      connectionStream = new PortStream(extensionPort);
+      extensionPort.onMessage.addListener(messageListener);
+      extensionPort.onDisconnect.addListener(resetExtensionStreamAndListeners);
+    };
+
+    extensionPort.onMessage.addListener(messageListener);
+    extensionPort.onDisconnect.addListener(resetExtensionStreamAndListeners);
   } else {
     initializeUiWithTab(activeTab);
   }
 
   function initializeUiWithTab(tab) {
-    const container = document.getElementById('app-content');
-    initializeUi(tab, container, connectionStream, (err, store) => {
+    initializeUi(tab, connectionStream, (err, store) => {
       if (err) {
         // if there's an error, store will be = metamaskState
-        displayCriticalError(container, err, store);
+        displayCriticalError(err, store);
         return;
       }
+      isUIInitialised = true;
 
       const state = store.getState();
       const { metamask: { completedOnboarding } = {} } = state;
@@ -86,6 +118,18 @@ async function start() {
       }
     });
   }
+
+  // Function to update new backgroundConnection in the UI
+  function updateUiStreams() {
+    connectToAccountManager(connectionStream, (err, backgroundConnection) => {
+      if (err) {
+        displayCriticalError(err);
+        return;
+      }
+
+      updateBackgroundConnection(backgroundConnection);
+    });
+  }
 }
 
 async function queryCurrentActiveTab(windowType) {
@@ -112,7 +156,7 @@ async function queryCurrentActiveTab(windowType) {
   });
 }
 
-function initializeUi(activeTab, container, connectionStream, cb) {
+function initializeUi(activeTab, connectionStream, cb) {
   connectToAccountManager(connectionStream, (err, backgroundConnection) => {
     if (err) {
       cb(err, null);
@@ -130,6 +174,21 @@ function initializeUi(activeTab, container, connectionStream, cb) {
   });
 }
 
+async function displayCriticalError(err, metamaskState) {
+  const html = await getErrorHtml(SUPPORT_LINK, metamaskState);
+
+  container.innerHTML = html;
+
+  const button = document.getElementById('critical-error-button');
+
+  button.addEventListener('click', (_) => {
+    browser.runtime.reload();
+  });
+
+  log.error(err.stack);
+  throw err;
+}
+
 /**
  * Establishes a connection to the background and a Web3 provider
  *

ui/index.js

@@ -31,9 +31,30 @@ import txHelper from './helpers/utils/tx-helper';
 
 log.setLevel(global.METAMASK_DEBUG ? 'debug' : 'warn');
 
+let reduxStore;
+
+/**
+ * Method to update backgroundConnection object use by UI
+ *
+ * @param backgroundConnection - connection object to background
+ */
+export const updateBackgroundConnection = (backgroundConnection) => {
+  actions._setBackgroundConnection(backgroundConnection);
+  backgroundConnection.onNotification((data) => {
+    if (data.method === 'sendUpdate') {
+      reduxStore.dispatch(actions.updateMetamaskState(data.params[0]));
+    } else {
+      throw new Error(
+        `Internal JSON-RPC Notification Not Handled:\n\n ${JSON.stringify(
+          data,
+        )}`,
+      );
+    }
+  });
+};
+
 export default function launchMetamaskUi(opts, cb) {
   const { backgroundConnection } = opts;
-  actions._setBackgroundConnection(backgroundConnection);
   // check if we are unlocked first
   backgroundConnection.getState(function (err, metamaskState) {
     if (err) {
@@ -117,6 +138,7 @@ async function startApp(metamaskState, backgroundConnection, opts) {
   }
 
   const store = configureStore(draftInitialState);
+  reduxStore = store;
 
   // if unconfirmed txs, start on txConf page
   const unapprovedTxsAll = txHelper(
@@ -138,17 +160,7 @@ async function startApp(metamaskState, backgroundConnection, opts) {
     );
   }
 
-  backgroundConnection.onNotification((data) => {
-    if (data.method === 'sendUpdate') {
-      store.dispatch(actions.updateMetamaskState(data.params[0]));
-    } else {
-      throw new Error(
-        `Internal JSON-RPC Notification Not Handled:\n\n ${JSON.stringify(
-          data,
-        )}`,
-      );
-    }
-  });
+  updateBackgroundConnection(backgroundConnection);
 
   // global metamask api - used by tooling
   global.metamask = {