From 3814616c7a25d4160caac838b9330f9e82dc6697 Mon Sep 17 00:00:00 2001
From: Bernardo Garces Chapero <bernardo.chapero@consensys.net>
Date: Tue, 14 Feb 2023 16:06:19 +0000
Subject: [PATCH] feat: refactor background.js to support overrides (#17625)

---
 app/scripts/background.js | 38 ++++++++++++++++++++++++++++++--------
 1 file changed, 30 insertions(+), 8 deletions(-)

diff --git a/app/scripts/background.js b/app/scripts/background.js
index 4e2f31581..6f47041ea 100644
--- a/app/scripts/background.js
+++ b/app/scripts/background.js
@@ -2,6 +2,7 @@
  * @file The entry point for the web extension singleton process.
  */
 
+import EventEmitter from 'events';
 import endOfStream from 'end-of-stream';
 import pump from 'pump';
 import debounce from 'debounce-stream';
@@ -96,6 +97,9 @@ const PHISHING_WARNING_PAGE_TIMEOUT = ONE_SECOND_IN_MILLISECONDS;
 const ACK_KEEP_ALIVE_MESSAGE = 'ACK_KEEP_ALIVE_MESSAGE';
 const WORKER_KEEP_ALIVE_MESSAGE = 'WORKER_KEEP_ALIVE_MESSAGE';
 
+// Event emitter for state persistence
+export const statePersistenceEvents = new EventEmitter();
+
 /**
  * This deferred Promise is used to track whether initialization has finished.
  *
@@ -177,8 +181,6 @@ browser.runtime.onConnectExternal.addListener(async (...args) => {
   connectExternal(...args);
 });
 
-initialize().catch(log.error);
-
 /**
  * @typedef {import('../../shared/constants/transaction').TransactionMeta} TransactionMeta
  */
@@ -331,7 +333,7 @@ async function loadPhishingWarningPage() {
  *
  * @returns {Promise<MetaMaskState>} Last data emitted from previous instance of MetaMask.
  */
-async function loadStateFromPersistence() {
+export async function loadStateFromPersistence() {
   // migrations
   const migrator = new Migrator({ migrations });
   migrator.on('error', console.warn);
@@ -384,8 +386,9 @@ 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)
  */
-function setupController(initState, initLangCode) {
+export function setupController(initState, initLangCode, overrides) {
   //
   // MetaMask Controller
   //
@@ -410,6 +413,7 @@ function setupController(initState, initLangCode) {
       return openMetamaskTabsIDs;
     },
     localStore,
+    overrides,
   });
 
   setupEnsIpfsResolver({
@@ -426,7 +430,10 @@ function setupController(initState, initLangCode) {
   pump(
     storeAsStream(controller.store),
     debounce(1000),
-    createStreamSink((state) => localStore.set(state)),
+    createStreamSink(async (state) => {
+      await localStore.set(state);
+      statePersistenceEvents.emit('state-persisted', state);
+    }),
     (error) => {
       log.error('MetaMask - Persistence pipeline failed', error);
     },
@@ -495,7 +502,8 @@ function setupController(initState, initLangCode) {
     }
 
     if (isMetaMaskInternalProcess) {
-      const portStream = new PortStream(remotePort);
+      const portStream =
+        overrides?.getPortStream?.(remotePort) || new PortStream(remotePort);
       // communication with popup
       controller.isClientOpen = true;
       controller.setupTrustedCommunication(portStream, remotePort.sender);
@@ -553,7 +561,8 @@ function setupController(initState, initLangCode) {
       senderUrl.origin === phishingPageUrl.origin &&
       senderUrl.pathname === phishingPageUrl.pathname
     ) {
-      const portStream = new PortStream(remotePort);
+      const portStream =
+        overrides?.getPortStream?.(remotePort) || new PortStream(remotePort);
       controller.setupPhishingCommunication({
         connectionStream: portStream,
       });
@@ -575,13 +584,18 @@ function setupController(initState, initLangCode) {
 
   // communication with page or other extension
   connectExternal = (remotePort) => {
-    const portStream = new PortStream(remotePort);
+    const portStream =
+      overrides?.getPortStream?.(remotePort) || new PortStream(remotePort);
     controller.setupUntrustedCommunication({
       connectionStream: portStream,
       sender: remotePort.sender,
     });
   };
 
+  if (overrides?.registerConnectListeners) {
+    overrides.registerConnectListeners(connectRemote, connectExternal);
+  }
+
   //
   // User Interface setup
   //
@@ -839,3 +853,11 @@ function setupSentryGetStateGlobal(store) {
     };
   };
 }
+
+function initBackground() {
+  initialize().catch(log.error);
+}
+
+if (!process.env.SKIP_BACKGROUND_INITIALIZATION) {
+  initBackground();
+}