MV3: Update service worker restart logic and keep-alive logic for dapp support (#16075)

* dapp: add debug statements

* dapp: add retry logic [debug]

* dapp: keep SW alive on rpc request

* Revert "dapp: add debug statements"

This reverts commit ea21786f7f66c712eea02405cd68fe925d227ffa.

* dapp: try to set up ext streams asap on reset

* dapp: apply keep alive logic to phishingPageStream

* dapp:put keep-alive logic behind isManifestV3 flag

* Re-activate streams after a period of service worker in-activity

* dapp: rm extra function

* dapp: update phishing onDisconnect

* dapp: fix eslint missing global chrome

* add EXTENSION_MESSAGES const

* use EXTENSION_MESSAGES more generic comment

* update comment

* dapp: clean timeout and interval

* Fix DAPP action replay

* execute DAPP action replay for only MV3

* fix

* fix

* fix

* comment out DAPP action replay code

* fix

* fix

* fix

* scripts/background: use browser polyfill

* Revert "scripts/background: use browser polyfill"

This reverts commit 2ab6234d11b3b11e10dd993d454eeaad63bfc886.

* scripts/background: use browser polyfill

* script/background: check lastError

* dapp: use EXTENSION_MESSAGES

* scripts/background: send ready msg to all tabs

* dapp: update onMessage handler comment and name

* dapp: return values onMessage listener
see: https://developer.mozilla.org/en-US/docs/Mozilla/Add-ons/WebExtensions/API/runtime/onMessage#addlistener_syntax

* dapp: mv onMessage listener

* dapp: add onMessage setupPhishingExtStreams

* dapp: rn reset -> destroy streams

* dapp: rn reset -> destroy for phishing streams

* dapp: clean comment

* dapp: rm unused comments
planning to be readded in follow-up PR: #16250

* dapp: onMessage return Promise|undefined

* dapp:clean: add missing undefined return type

* dapp: use new checkForErrorAndLog for Chrome API
handy stackoverflow: https://stackoverflow.com/a/28432087/4053142

* dapp:fix: return tabs.query result

* dapp:eslint: return undefined
fix Expected to return a value at the end of arrow
function.eslintconsistent-return

* background: do not query tabs w/out url

* background: rm Could not establish... catch
- no longer needed after improved tabs query

* dapp:clean: rm unused checkForError... for now...

* dapp: prevent setupExtensionStreams called twice
- calling connect will trigger disconnect and may cause issues
- only setup streams if they are not connected

* dapp: handle onDisconnect lastError
- throwing errors from contentscript will break the dapp, so only warn
- not handling lastError when it's found will also break the dapp

* background: update tabs.query url comment

* background: update tabs.query url comment 2

* dapp: fix SW restart for multi dapp support
- ref: https://stackoverflow.com/a/54686484/4053142

* dapp:clean: rm extra "." from console.warn

* clean: comments for dapp and background

* Adding catch block (#16454)

* fix: FireFox provider injection

* lib/util: fix invalid checkForErrorAndWarn export

* bg: add explanation for tabs.sendMessage catch

* dapp: add browser-runtime.utils

* runtime.utils: add checkForLastErrorAndLog

Co-authored-by: Jyoti Puri <jyotipuri@gmail.com>
Co-authored-by: legobeat <109787230+legobeat@users.noreply.github.com>

app/scripts/background.js

@@ -15,6 +15,7 @@ import {
   ENVIRONMENT_TYPE_POPUP,
   ENVIRONMENT_TYPE_NOTIFICATION,
   ENVIRONMENT_TYPE_FULLSCREEN,
+  EXTENSION_MESSAGES,
   PLATFORM_FIREFOX,
 } from '../../shared/constants/app';
 import { SECOND } from '../../shared/constants/time';
@@ -25,6 +26,7 @@ import {
   EVENT_NAMES,
   TRAITS,
 } from '../../shared/constants/metametrics';
+import { checkForLastErrorAndLog } from '../../shared/modules/browser-runtime.utils';
 import { isManifestV3 } from '../../shared/modules/mv3.utils';
 import { maskObject } from '../../shared/modules/object.utils';
 import migrations from './migrations';
@@ -94,16 +96,67 @@ 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.
+ *
+ * @param remotePort
  */
-
 const initApp = async (remotePort) => {
   browser.runtime.onConnect.removeListener(initApp);
   await initialize(remotePort);
   log.info('MetaMask initialization complete.');
 };
 
+/**
+ * Sends a message to the dapp(s) content script to signal it can connect to MetaMask background as
+ * the backend is not active. It is required to re-connect dapps after service worker re-activates.
+ * For non-dapp pages, the message will be sent and ignored.
+ */
+const sendReadyMessageToTabs = async () => {
+  const tabs = await browser.tabs
+    .query({
+      /**
+       * Only query tabs that our extension can run in. To do this, we query for all URLs that our
+       * extension can inject scripts in, which is by using the "<all_urls>" value and __without__
+       * the "tabs" manifest permission. If we included the "tabs" permission, this would also fetch
+       * URLs that we'd not be able to inject in, e.g. chrome://pages, chrome://extension, which
+       * is not what we'd want.
+       *
+       * You might be wondering, how does the "url" param work without the "tabs" permission?
+       *
+       * @see {@link https://bugs.chromium.org/p/chromium/issues/detail?id=661311#c1}
+       *  "If the extension has access to inject scripts into Tab, then we can return the url
+       *   of Tab (because the extension could just inject a script to message the location.href)."
+       */
+      url: '<all_urls>',
+      windowType: 'normal',
+    })
+    .then((result) => {
+      checkForLastErrorAndLog();
+      return result;
+    })
+    .catch(() => {
+      checkForLastErrorAndLog();
+    });
+
+  /** @todo we should only sendMessage to dapp tabs, not all tabs. */
+  for (const tab of tabs) {
+    browser.tabs
+      .sendMessage(tab.id, {
+        name: EXTENSION_MESSAGES.READY,
+      })
+      .then(() => {
+        checkForLastErrorAndLog();
+      })
+      .catch(() => {
+        // An error may happen if the contentscript is blocked from loading,
+        // and thus there is no runtime.onMessage handler to listen to the message.
+        checkForLastErrorAndLog();
+      });
+  }
+};
+
 if (isManifestV3) {
   browser.runtime.onConnect.addListener(initApp);
+  sendReadyMessageToTabs();
 } else {
   // initialization flow
   initialize().catch(log.error);

app/scripts/contentscript.js

@@ -5,6 +5,11 @@ import browser from 'webextension-polyfill';
 import PortStream from 'extension-port-stream';
 import { obj as createThoughStream } from 'through2';
 
+import { EXTENSION_MESSAGES, MESSAGE_TYPE } from '../../shared/constants/app';
+import {
+  checkForLastError,
+  checkForLastErrorAndWarn,
+} from '../../shared/modules/browser-runtime.utils';
 import { isManifestV3 } from '../../shared/modules/mv3.utils';
 import shouldInjectProvider from '../../shared/modules/provider-injection';
 
@@ -44,9 +49,6 @@ let legacyExtMux,
   legacyPagePublicConfigChannel,
   notificationTransformStream;
 
-const WORKER_KEEP_ALIVE_INTERVAL = 1000;
-const WORKER_KEEP_ALIVE_MESSAGE = 'WORKER_KEEP_ALIVE_MESSAGE';
-
 const phishingPageUrl = new URL(process.env.PHISHING_WARNING_PAGE_URL);
 
 let phishingExtChannel,
@@ -82,6 +84,51 @@ function injectScript(content) {
   }
 }
 
+/**
+ * SERVICE WORKER LOGIC
+ */
+
+const WORKER_KEEP_ALIVE_INTERVAL = 1000;
+const WORKER_KEEP_ALIVE_MESSAGE = 'WORKER_KEEP_ALIVE_MESSAGE';
+const TIME_45_MIN_IN_MS = 45 * 60 * 1000;
+
+/**
+ * Don't run the keep-worker-alive logic for JSON-RPC methods called on initial load.
+ * This is to prevent the service worker from being kept alive when accounts are not
+ * connected to the dapp or when the user is not interacting with the extension.
+ * The keep-alive logic should not work for non-dapp pages.
+ */
+const IGNORE_INIT_METHODS_FOR_KEEP_ALIVE = [
+  MESSAGE_TYPE.GET_PROVIDER_STATE,
+  MESSAGE_TYPE.SEND_METADATA,
+];
+
+let keepAliveInterval;
+let keepAliveTimer;
+
+/**
+ * Running this method will ensure the service worker is kept alive for 45 minutes.
+ * The first message is sent immediately and subsequent messages are sent at an
+ * interval of WORKER_KEEP_ALIVE_INTERVAL.
+ */
+const runWorkerKeepAliveInterval = () => {
+  clearTimeout(keepAliveTimer);
+
+  keepAliveTimer = setTimeout(() => {
+    clearInterval(keepAliveInterval);
+  }, TIME_45_MIN_IN_MS);
+
+  clearInterval(keepAliveInterval);
+
+  browser.runtime.sendMessage({ name: WORKER_KEEP_ALIVE_MESSAGE });
+
+  keepAliveInterval = setInterval(() => {
+    if (browser.runtime.id) {
+      browser.runtime.sendMessage({ name: WORKER_KEEP_ALIVE_MESSAGE });
+    }
+  }, WORKER_KEEP_ALIVE_INTERVAL);
+};
+
 /**
  * PHISHING STREAM LOGIC
  */
@@ -93,6 +140,14 @@ function setupPhishingPageStreams() {
     target: PHISHING_WARNING_PAGE,
   });
 
+  if (isManifestV3) {
+    phishingPageStream.on('data', ({ data: { method } }) => {
+      if (!IGNORE_INIT_METHODS_FOR_KEEP_ALIVE.includes(method)) {
+        runWorkerKeepAliveInterval();
+      }
+    });
+  }
+
   // create and connect channel muxers
   // so we can handle the channels individually
   phishingPageMux = new ObjectMultiplex();
@@ -142,6 +197,9 @@ const setupPhishingExtStreams = () => {
       error,
     ),
   );
+
+  // eslint-disable-next-line no-use-before-define
+  phishingExtPort.onDisconnect.addListener(onDisconnectDestroyPhishingStreams);
 };
 
 /** Destroys all of the phishing extension streams */
@@ -153,19 +211,42 @@ const destroyPhishingExtStreams = () => {
 
   phishingExtChannel.removeAllListeners();
   phishingExtChannel.destroy();
+
+  phishingExtStream = null;
 };
 
 /**
- * Resets the extension stream with new streams to channel with the phishing page streams,
- * and creates a new event listener to the reestablished extension port.
+ * This listener destroys the phishing extension streams when the extension port is disconnected,
+ * so that streams may be re-established later the phishing extension port is reconnected.
  */
-const resetPhishingStreamAndListeners = () => {
-  phishingExtPort.onDisconnect.removeListener(resetPhishingStreamAndListeners);
+const onDisconnectDestroyPhishingStreams = () => {
+  checkForLastErrorAndWarn();
+
+  phishingExtPort.onDisconnect.removeListener(
+    onDisconnectDestroyPhishingStreams,
+  );
 
   destroyPhishingExtStreams();
-  setupPhishingExtStreams();
+};
 
-  phishingExtPort.onDisconnect.addListener(resetPhishingStreamAndListeners);
+/**
+ * When the extension background is loaded it sends the EXTENSION_MESSAGES.READY message to the browser tabs.
+ * This listener/callback receives the message to set up the streams after service worker in-activity.
+ *
+ * @param {object} msg
+ * @param {string} msg.name - custom property and name to identify the message received
+ * @returns {Promise|undefined}
+ */
+const onMessageSetUpPhishingStreams = (msg) => {
+  if (msg.name === EXTENSION_MESSAGES.READY) {
+    if (!phishingExtStream) {
+      setupPhishingExtStreams();
+    }
+    return Promise.resolve(
+      `MetaMask: handled "${EXTENSION_MESSAGES.READY}" for phishing streams`,
+    );
+  }
+  return undefined;
 };
 
 /**
@@ -177,7 +258,7 @@ const initPhishingStreams = () => {
   setupPhishingPageStreams();
   setupPhishingExtStreams();
 
-  phishingExtPort.onDisconnect.addListener(resetPhishingStreamAndListeners);
+  browser.runtime.onMessage.addListener(onMessageSetUpPhishingStreams);
 };
 
 /**
@@ -191,6 +272,14 @@ const setupPageStreams = () => {
     target: INPAGE,
   });
 
+  if (isManifestV3) {
+    pageStream.on('data', ({ data: { method } }) => {
+      if (!IGNORE_INIT_METHODS_FOR_KEEP_ALIVE.includes(method)) {
+        runWorkerKeepAliveInterval();
+      }
+    });
+  }
+
   // create and connect channel muxers
   // so we can handle the channels individually
   pageMux = new ObjectMultiplex();
@@ -231,7 +320,8 @@ const setupExtensionStreams = () => {
   extensionPhishingStream = extensionMux.createStream('phishing');
   extensionPhishingStream.once('data', redirectToPhishingWarning);
 
-  notifyInpageOfExtensionStreamConnect();
+  // eslint-disable-next-line no-use-before-define
+  extensionPort.onDisconnect.addListener(onDisconnectDestroyStreams);
 };
 
 /** Destroys all of the extension streams */
@@ -243,10 +333,13 @@ const destroyExtensionStreams = () => {
 
   extensionChannel.removeAllListeners();
   extensionChannel.destroy();
+
+  extensionStream = null;
 };
 
 /**
  * LEGACY STREAM LOGIC
+ * TODO:LegacyProvider: Delete
  */
 
 // TODO:LegacyProvider: Delete
@@ -256,6 +349,14 @@ const setupLegacyPageStreams = () => {
     target: LEGACY_INPAGE,
   });
 
+  if (isManifestV3) {
+    legacyPageStream.on('data', ({ data: { method } }) => {
+      if (!IGNORE_INIT_METHODS_FOR_KEEP_ALIVE.includes(method)) {
+        runWorkerKeepAliveInterval();
+      }
+    });
+  }
+
   legacyPageMux = new ObjectMultiplex();
   legacyPageMux.setMaxListeners(25);
 
@@ -331,19 +432,47 @@ const destroyLegacyExtensionStreams = () => {
 };
 
 /**
- * Resets the extension stream with new streams to channel with the in page streams,
- * and creates a new event listener to the reestablished extension port.
+ * When the extension background is loaded it sends the EXTENSION_MESSAGES.READY message to the browser tabs.
+ * This listener/callback receives the message to set up the streams after service worker in-activity.
+ *
+ * @param {object} msg
+ * @param {string} msg.name - custom property and name to identify the message received
+ * @returns {Promise|undefined}
  */
-const resetStreamAndListeners = () => {
-  extensionPort.onDisconnect.removeListener(resetStreamAndListeners);
-
-  destroyExtensionStreams();
+const onMessageSetUpExtensionStreams = (msg) => {
+  if (msg.name === EXTENSION_MESSAGES.READY) {
+    if (!extensionStream) {
       setupExtensionStreams();
+      setupLegacyExtensionStreams();
+    }
+    return Promise.resolve(`MetaMask: handled ${EXTENSION_MESSAGES.READY}`);
+  }
+  return undefined;
+};
 
+/**
+ * This listener destroys the extension streams when the extension port is disconnected,
+ * so that streams may be re-established later when the extension port is reconnected.
+ */
+const onDisconnectDestroyStreams = () => {
+  const err = checkForLastError();
+
+  extensionPort.onDisconnect.removeListener(onDisconnectDestroyStreams);
+
+  destroyExtensionStreams();
   destroyLegacyExtensionStreams();
-  setupLegacyExtensionStreams();
 
-  extensionPort.onDisconnect.addListener(resetStreamAndListeners);
+  /**
+   * If an error is found, reset the streams. When running two or more dapps, resetting the service
+   * worker may cause the error, "Error: Could not establish connection. Receiving end does not
+   * exist.", due to a race-condition. The disconnect event may be called by runtime.connect which
+   * may cause issues. We suspect that this is a chromium bug as this event should only be called
+   * once the port and connections are ready. Delay time is arbitrary.
+   */
+  if (err) {
+    console.warn(`${err} Resetting the streams.`);
+    setTimeout(setupExtensionStreams, 1000);
+  }
 };
 
 /**
@@ -353,13 +482,12 @@ const resetStreamAndListeners = () => {
  */
 const initStreams = () => {
   setupPageStreams();
-  setupExtensionStreams();
-
-  // TODO:LegacyProvider: Delete
   setupLegacyPageStreams();
+
+  setupExtensionStreams();
   setupLegacyExtensionStreams();
 
-  extensionPort.onDisconnect.addListener(resetStreamAndListeners);
+  browser.runtime.onMessage.addListener(onMessageSetUpExtensionStreams);
 };
 
 // TODO:LegacyProvider: Delete
@@ -389,26 +517,6 @@ function logStreamDisconnectWarning(remoteLabel, error) {
   );
 }
 
-/**
- * The function send message to inpage to notify it of extension stream connection
- */
-function notifyInpageOfExtensionStreamConnect() {
-  window.postMessage(
-    {
-      target: INPAGE, // the post-message-stream "target"
-      data: {
-        // this object gets passed to obj-multiplex
-        name: PROVIDER, // the obj-multiplex channel name
-        data: {
-          jsonrpc: '2.0',
-          method: 'METAMASK_EXTENSION_STREAM_CONNECT',
-        },
-      },
-    },
-    window.location.origin,
-  );
-}
-
 /**
  * This function must ONLY be called in pump destruction/close callbacks.
  * Notifies the inpage context that streams have failed, via window.postMessage.
@@ -446,12 +554,6 @@ function redirectToPhishingWarning(data = {}) {
   window.location.href = `${baseUrl}#${querystring}`;
 }
 
-const initKeepWorkerAlive = () => {
-  setInterval(() => {
-    browser.runtime.sendMessage({ name: WORKER_KEEP_ALIVE_MESSAGE });
-  }, WORKER_KEEP_ALIVE_INTERVAL);
-};
-
 const start = () => {
   const isDetectedPhishingSite =
     window.location.origin === phishingPageUrl.origin &&
@@ -463,9 +565,7 @@ const start = () => {
   }
 
   if (shouldInjectProvider()) {
-    if (isManifestV3) {
-      initKeepWorkerAlive();
-    } else {
+    if (!isManifestV3) {
       injectScript(inpageBundle);
     }
     initStreams();

app/scripts/lib/util.js

@@ -96,6 +96,7 @@ function BnMultiplyByFraction(targetBN, numerator, denominator) {
  * Returns an Error if extension.runtime.lastError is present
  * this is a workaround for the non-standard error object that's used
  *
+ * @deprecated use checkForLastError in shared/modules/browser-runtime.utils.js
  * @returns {Error|undefined}
  */
 function checkForError() {

shared/constants/app.ts

@@ -57,6 +57,13 @@ export const MESSAGE_TYPE = {
   ///: END:ONLY_INCLUDE_IN
 } as const;
 
+/**
+ * Custom messages to send and be received by the extension
+ */
+export const EXTENSION_MESSAGES = {
+  READY: 'METAMASK_EXTENSION_READY',
+} as const;
+
 /**
  * The different kinds of subjects that MetaMask may interact with, including
  * third parties and itself (e.g. when the background communicated with the UI).

shared/modules/browser-runtime.utils.js

@@ -0,0 +1,55 @@
+/**
+ * Utility Functions to support browser.runtime JavaScript API
+ */
+
+import browser from 'webextension-polyfill';
+import log from 'loglevel';
+
+/**
+ * Returns an Error if extension.runtime.lastError is present
+ * this is a workaround for the non-standard error object that's used
+ *
+ * According to the docs, we are expected to check lastError in runtime API callbacks:
+ * "
+ * If you call an asynchronous function that may set lastError, you are expected to
+ * check for the error when you handle the result of the function. If lastError has been
+ * set and you don't check it within the callback function, then an error will be raised.
+ * "
+ *
+ * @see {@link https://developer.mozilla.org/en-US/docs/Mozilla/Add-ons/WebExtensions/API/runtime/lastError}
+ * @returns {Error|undefined}
+ */
+export function checkForLastError() {
+  const { lastError } = browser.runtime;
+  if (!lastError) {
+    return undefined;
+  }
+  // if it quacks like an Error, its an Error
+  if (lastError.stack && lastError.message) {
+    return lastError;
+  }
+  // repair incomplete error object (eg chromium v77)
+  return new Error(lastError.message);
+}
+
+/** @returns {Error|undefined} */
+export function checkForLastErrorAndLog() {
+  const error = checkForLastError();
+
+  if (error) {
+    log.error(error);
+  }
+
+  return error;
+}
+
+/** @returns {Error|undefined} */
+export function checkForLastErrorAndWarn() {
+  const error = checkForLastError();
+
+  if (error) {
+    console.warn(error);
+  }
+
+  return error;
+}