MV3: contentscript.js - re-activate streams when Service Worker terminates and then resets (#15494)

* contentscript: reactivate streams pt.1

* contentscript: reactivate streams pt. 2

* scripts/inpage.js: use const

* clean: rm unused @param

* contentscript: reactivate streams pt. 3
- reorganize functions

* resetSteamAndListeners -> resetStreamAndListeners

* contentscript: skip Legacy Provider while WIP

* contentscript: rm comment sentence

* initPageStreams -> initStreams

* setupPageStreams -> setupStreams

* contentscript: only destroy extension streams WIP

* remove pageMuxChannel listeners (It works!)
- credits to @gudhatt for this missing piece

Co-authored-by: Mark Stacey <markjstacey@gmail.com>

* contentscript: cleanup

* contentscript: support legacy streams
w/ service workers

* contentscript: support phishing streams
w/ service workers

* contentscript: muxChannel -> channel

* contentscript: phishingExension -> phishingExt

* contentscript: add section comments

* contentscript: revert condensing comment

* contentscript: pagePhishing -> phishingPage

* contentscript: update comments

* contentscript: fix phishingExtPort

* contentscript: stream -> streams

* contentscript: update SW keep alive logic
should only keep alive when dapp is open / contentscript page

* rm console.log

Co-authored-by: Mark Stacey <markjstacey@gmail.com>
Co-authored-by: Jyoti Puri <jyotipuri@gmail.com>

app/scripts/contentscript.js

@@ -36,19 +36,33 @@ const LEGACY_INPAGE = 'inpage';
 const LEGACY_PROVIDER = 'provider';
 const LEGACY_PUBLIC_CONFIG = 'publicConfig';
 
+let legacyExtMux,
+  legacyExtChannel,
+  legacyExtPublicConfigChannel,
+  legacyPageMux,
+  legacyPageMuxLegacyProviderChannel,
+  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);
 
-if (
+let phishingExtChannel,
-  window.location.origin === phishingPageUrl.origin &&
+  phishingExtMux,
-  window.location.pathname === phishingPageUrl.pathname
+  phishingExtPort,
-) {
+  phishingExtStream,
-  setupPhishingStream();
+  phishingPageChannel,
-} else if (shouldInjectProvider()) {
+  phishingPageMux;
-  if (!isManifestV3) {
+
-    injectScript(inpageBundle);
+let extensionMux,
-  }
+  extensionChannel,
-  setupStreams();
+  extensionPort,
-}
+  extensionPhishingStream,
+  extensionStream,
+  pageMux,
+  pageChannel;
 
 /**
  * Injects a script tag into the current document
@@ -68,26 +82,41 @@ function injectScript(content) {
   }
 }
 
-async function setupPhishingStream() {
+/**
+ * PHISHING STREAM LOGIC
+ */
+
+function setupPhishingPageStreams() {
   // the transport-specific streams for communication between inpage and background
-  const pageStream = new WindowPostMessageStream({
+  const phishingPageStream = new WindowPostMessageStream({
     name: CONTENT_SCRIPT,
     target: PHISHING_WARNING_PAGE,
   });
-  const extensionPort = browser.runtime.connect({ name: CONTENT_SCRIPT });
-  const extensionStream = new PortStream(extensionPort);
 
   // create and connect channel muxers
   // so we can handle the channels individually
-  const pageMux = new ObjectMultiplex();
+  phishingPageMux = new ObjectMultiplex();
-  pageMux.setMaxListeners(25);
+  phishingPageMux.setMaxListeners(25);
-  const extensionMux = new ObjectMultiplex();
-  extensionMux.setMaxListeners(25);
 
-  pump(pageMux, pageStream, pageMux, (err) =>
+  pump(phishingPageMux, phishingPageStream, phishingPageMux, (err) =>
     logStreamDisconnectWarning('MetaMask Inpage Multiplex', err),
   );
-  pump(extensionMux, extensionStream, extensionMux, (err) => {
+
+  phishingPageChannel = phishingPageMux.createStream(PHISHING_SAFELIST);
+}
+
+const setupPhishingExtStreams = () => {
+  phishingExtPort = browser.runtime.connect({
+    name: CONTENT_SCRIPT,
+  });
+  phishingExtStream = new PortStream(phishingExtPort);
+
+  // create and connect channel muxers
+  // so we can handle the channels individually
+  phishingExtMux = new ObjectMultiplex();
+  phishingExtMux.setMaxListeners(25);
+
+  pump(phishingExtMux, phishingExtStream, phishingExtMux, (err) => {
     logStreamDisconnectWarning('MetaMask Background Multiplex', err);
     window.postMessage(
       {
@@ -106,112 +135,230 @@ async function setupPhishingStream() {
   });
 
   // forward communication across inpage-background for these channels only
-  forwardTrafficBetweenMuxes(PHISHING_SAFELIST, pageMux, extensionMux);
+  phishingExtChannel = phishingExtMux.createStream(PHISHING_SAFELIST);
-}
+  pump(phishingPageChannel, phishingExtChannel, phishingPageChannel, (error) =>
+    console.debug(
+      `MetaMask: Muxed traffic for channel "${PHISHING_SAFELIST}" failed.`,
+      error,
+    ),
+  );
+};
+
+/** Destroys all of the phishing extension streams */
+const destroyPhishingExtStreams = () => {
+  phishingPageChannel.removeAllListeners();
+
+  phishingExtMux.removeAllListeners();
+  phishingExtMux.destroy();
+
+  phishingExtChannel.removeAllListeners();
+  phishingExtChannel.destroy();
+};
 
 /**
- * Sets up two-way communication streams between the
+ * Resets the extension stream with new streams to channel with the phishing page streams,
- * browser extension and local per-page browser context.
+ * and creates a new event listener to the reestablished extension port.
- *
  */
-async function setupStreams() {
+const resetPhishingStreamAndListeners = () => {
+  phishingExtPort.onDisconnect.removeListener(resetPhishingStreamAndListeners);
+
+  destroyPhishingExtStreams();
+  setupPhishingExtStreams();
+
+  phishingExtPort.onDisconnect.addListener(resetPhishingStreamAndListeners);
+};
+
+/**
+ * Initializes two-way communication streams between the browser extension and
+ * the phishing page context. This function also creates an event listener to
+ * reset the streams if the service worker resets.
+ */
+const initPhishingStreams = () => {
+  setupPhishingPageStreams();
+  setupPhishingExtStreams();
+
+  phishingExtPort.onDisconnect.addListener(resetPhishingStreamAndListeners);
+};
+
+/**
+ * INPAGE - EXTENSION STREAM LOGIC
+ */
+
+const setupPageStreams = () => {
   // the transport-specific streams for communication between inpage and background
   const pageStream = new WindowPostMessageStream({
     name: CONTENT_SCRIPT,
     target: INPAGE,
   });
-  const extensionPort = browser.runtime.connect({ name: CONTENT_SCRIPT });
-  const extensionStream = new PortStream(extensionPort);
 
   // create and connect channel muxers
   // so we can handle the channels individually
-  const pageMux = new ObjectMultiplex();
+  pageMux = new ObjectMultiplex();
   pageMux.setMaxListeners(25);
-  const extensionMux = new ObjectMultiplex();
-  extensionMux.setMaxListeners(25);
-  extensionMux.ignoreStream(LEGACY_PUBLIC_CONFIG); // TODO:LegacyProvider: Delete
 
   pump(pageMux, pageStream, pageMux, (err) =>
     logStreamDisconnectWarning('MetaMask Inpage Multiplex', err),
   );
+
+  pageChannel = pageMux.createStream(PROVIDER);
+};
+
+const setupExtensionStreams = () => {
+  extensionPort = browser.runtime.connect({ name: CONTENT_SCRIPT });
+  extensionStream = new PortStream(extensionPort);
+
+  // create and connect channel muxers
+  // so we can handle the channels individually
+  extensionMux = new ObjectMultiplex();
+  extensionMux.setMaxListeners(25);
+  extensionMux.ignoreStream(LEGACY_PUBLIC_CONFIG); // TODO:LegacyProvider: Delete
+
   pump(extensionMux, extensionStream, extensionMux, (err) => {
     logStreamDisconnectWarning('MetaMask Background Multiplex', err);
     notifyInpageOfStreamFailure();
   });
 
   // forward communication across inpage-background for these channels only
-  forwardTrafficBetweenMuxes(PROVIDER, pageMux, extensionMux);
+  extensionChannel = extensionMux.createStream(PROVIDER);
+  pump(pageChannel, extensionChannel, pageChannel, (error) =>
+    console.debug(
+      `MetaMask: Muxed traffic for channel "${PROVIDER}" failed.`,
+      error,
+    ),
+  );
 
   // connect "phishing" channel to warning system
-  const phishingStream = extensionMux.createStream('phishing');
+  extensionPhishingStream = extensionMux.createStream('phishing');
-  phishingStream.once('data', redirectToPhishingWarning);
+  extensionPhishingStream.once('data', redirectToPhishingWarning);
+};
+
+/** Destroys all of the extension streams */
+const destroyExtensionStreams = () => {
+  pageChannel.removeAllListeners();
+
+  extensionMux.removeAllListeners();
+  extensionMux.destroy();
+
+  extensionChannel.removeAllListeners();
+  extensionChannel.destroy();
+};
+
+/**
+ * LEGACY STREAM LOGIC
+ */
 
 // TODO:LegacyProvider: Delete
-  // handle legacy provider
+const setupLegacyPageStreams = () => {
   const legacyPageStream = new WindowPostMessageStream({
     name: LEGACY_CONTENT_SCRIPT,
     target: LEGACY_INPAGE,
   });
 
-  const legacyPageMux = new ObjectMultiplex();
+  legacyPageMux = new ObjectMultiplex();
   legacyPageMux.setMaxListeners(25);
-  const legacyExtensionMux = new ObjectMultiplex();
-  legacyExtensionMux.setMaxListeners(25);
 
   pump(legacyPageMux, legacyPageStream, legacyPageMux, (err) =>
     logStreamDisconnectWarning('MetaMask Legacy Inpage Multiplex', err),
   );
+
+  legacyPageMuxLegacyProviderChannel =
+    legacyPageMux.createStream(LEGACY_PROVIDER);
+  legacyPagePublicConfigChannel =
+    legacyPageMux.createStream(LEGACY_PUBLIC_CONFIG);
+};
+
+// TODO:LegacyProvider: Delete
+const setupLegacyExtensionStreams = () => {
+  legacyExtMux = new ObjectMultiplex();
+  legacyExtMux.setMaxListeners(25);
+
+  notificationTransformStream = getNotificationTransformStream();
   pump(
-    legacyExtensionMux,
+    legacyExtMux,
     extensionStream,
-    getNotificationTransformStream(),
+    notificationTransformStream,
-    legacyExtensionMux,
+    legacyExtMux,
     (err) => {
       logStreamDisconnectWarning('MetaMask Background Legacy Multiplex', err);
       notifyInpageOfStreamFailure();
     },
   );
 
-  forwardNamedTrafficBetweenMuxes(
+  legacyExtChannel = legacyExtMux.createStream(PROVIDER);
-    LEGACY_PROVIDER,
+  pump(
-    PROVIDER,
+    legacyPageMuxLegacyProviderChannel,
-    legacyPageMux,
+    legacyExtChannel,
-    legacyExtensionMux,
+    legacyPageMuxLegacyProviderChannel,
-  );
+    (error) =>
-  forwardTrafficBetweenMuxes(
-    LEGACY_PUBLIC_CONFIG,
-    legacyPageMux,
-    legacyExtensionMux,
-  );
-}
-
-function forwardTrafficBetweenMuxes(channelName, muxA, muxB) {
-  const channelA = muxA.createStream(channelName);
-  const channelB = muxB.createStream(channelName);
-  pump(channelA, channelB, channelA, (error) =>
       console.debug(
-      `MetaMask: Muxed traffic for channel "${channelName}" failed.`,
+        `MetaMask: Muxed traffic between channels "${LEGACY_PROVIDER}" and "${PROVIDER}" failed.`,
         error,
       ),
   );
-}
+
+  legacyExtPublicConfigChannel =
+    legacyExtMux.createStream(LEGACY_PUBLIC_CONFIG);
+  pump(
+    legacyPagePublicConfigChannel,
+    legacyExtPublicConfigChannel,
+    legacyPagePublicConfigChannel,
+    (error) =>
+      console.debug(
+        `MetaMask: Muxed traffic for channel "${LEGACY_PUBLIC_CONFIG}" failed.`,
+        error,
+      ),
+  );
+};
+
+/**
+ * Destroys all of the legacy extension streams
+ * TODO:LegacyProvider: Delete
+ */
+const destroyLegacyExtensionStreams = () => {
+  legacyPageMuxLegacyProviderChannel.removeAllListeners();
+  legacyPagePublicConfigChannel.removeAllListeners();
+
+  legacyExtMux.removeAllListeners();
+  legacyExtMux.destroy();
+
+  legacyExtChannel.removeAllListeners();
+  legacyExtChannel.destroy();
+
+  legacyExtPublicConfigChannel.removeAllListeners();
+  legacyExtPublicConfigChannel.destroy();
+};
+
+/**
+ * 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.
+ */
+const resetStreamAndListeners = () => {
+  extensionPort.onDisconnect.removeListener(resetStreamAndListeners);
+
+  destroyExtensionStreams();
+  setupExtensionStreams();
+
+  destroyLegacyExtensionStreams();
+  setupLegacyExtensionStreams();
+
+  extensionPort.onDisconnect.addListener(resetStreamAndListeners);
+};
+
+/**
+ * Initializes two-way communication streams between the browser extension and
+ * the local per-page browser context. This function also creates an event listener to
+ * reset the streams if the service worker resets.
+ */
+const initStreams = () => {
+  setupPageStreams();
+  setupExtensionStreams();
 
   // TODO:LegacyProvider: Delete
-function forwardNamedTrafficBetweenMuxes(
+  setupLegacyPageStreams();
-  channelAName,
+  setupLegacyExtensionStreams();
-  channelBName,
+
-  muxA,
+  extensionPort.onDisconnect.addListener(resetStreamAndListeners);
-  muxB,
+};
-) {
-  const channelA = muxA.createStream(channelAName);
-  const channelB = muxB.createStream(channelBName);
-  pump(channelA, channelB, channelA, (error) =>
-    console.debug(
-      `MetaMask: Muxed traffic between channels "${channelAName}" and "${channelBName}" failed.`,
-      error,
-    ),
-  );
-}
 
 // TODO:LegacyProvider: Delete
 function getNotificationTransformStream() {
@@ -276,3 +423,31 @@ function redirectToPhishingWarning(data = {}) {
   const querystring = new URLSearchParams({ hostname, href, newIssueUrl });
   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 &&
+    window.location.pathname === phishingPageUrl.pathname;
+
+  if (isDetectedPhishingSite) {
+    initPhishingStreams();
+    return;
+  }
+
+  if (shouldInjectProvider()) {
+    if (isManifestV3) {
+      initKeepWorkerAlive();
+    } else {
+      injectScript(inpageBundle);
+    }
+    initStreams();
+  }
+};
+
+start();

app/scripts/inpage.js

@@ -36,6 +36,10 @@ import { WindowPostMessageStream } from '@metamask/post-message-stream';
 import { initializeProvider } from '@metamask/providers/dist/initializeInpageProvider';
 import shouldInjectProvider from '../../shared/modules/provider-injection';
 
+// contexts
+const CONTENT_SCRIPT = 'metamask-contentscript';
+const INPAGE = 'metamask-inpage';
+
 restoreContextAfterImports();
 
 log.setDefaultLevel(process.env.METAMASK_DEBUG ? 'debug' : 'warn');
@@ -47,8 +51,8 @@ log.setDefaultLevel(process.env.METAMASK_DEBUG ? 'debug' : 'warn');
 if (shouldInjectProvider()) {
   // setup background connection
   const metamaskStream = new WindowPostMessageStream({
-    name: 'metamask-inpage',
+    name: INPAGE,
-    target: 'metamask-contentscript',
+    target: CONTENT_SCRIPT,
   });
 
   initializeProvider({