chore: Refactor offscreen creation logic (#25302)

## **Description**

Refactor initialization logic to defer creation of the offscreen
document until the `MetaMaskController` is initialized. This adds a
`offscreenPromise` to the controller that can be awaited for
functionality that requires the offscreen document to be created.

Additionally this PR adds a message that the offscreen document will
send once initial execution of the offscreen page has finished. This is
awaited in the `offscreenPromise`.

We await `offscreenPromise` before unlocking the keyrings as some
keyrings rely on the offscreen document to process requests, e.g.
hardware wallets.

There may be room for more improvements here though, that I have not
tackled in this PR. As the hardware wallet logic doesn't seem to wait
for iframes to fully load, so there is a chance of some missed messages.

I have tested that hardware wallet support, at least for Ledger, is
still working following the changes in this PR.

[![Open in GitHub
Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/MetaMask/metamask-extension/pull/25302?quickstart=1)

Author: FrederikBolding

app/scripts/app-init.js

@@ -185,27 +185,3 @@ const registerInPageContentScript = async () => {
 };
 
 registerInPageContentScript();
-
-/**
- * Creates an offscreen document that can be used to load additional scripts
- * and iframes that can communicate with the extension through the chrome
- * runtime API. Only one offscreen document may exist, so any iframes required
- * by extension can be embedded in the offscreen.html file. See the offscreen
- * folder for more details.
- */
-async function createOffscreen() {
-  if (!chrome.offscreen || (await chrome.offscreen.hasDocument())) {
-    return;
-  }
-
-  await chrome.offscreen.createDocument({
-    url: './offscreen.html',
-    reasons: ['IFRAME_SCRIPTING'],
-    justification:
-      'Used for Hardware Wallet and Snaps scripts to communicate with the extension.',
-  });
-
-  console.debug('Offscreen iframe loaded');
-}
-
-createOffscreen();

app/scripts/background.js

@@ -67,6 +67,7 @@ import {
   shouldEmitDappViewedEvent,
 } from './lib/util';
 import { generateSkipOnboardingState } from './skip-onboarding';
+import { createOffscreen } from './offscreen';
 
 /* eslint-enable import/first */
 
@@ -261,6 +262,8 @@ function saveTimestamp() {
  */
 async function initialize() {
   try {
+    const offscreenPromise = isManifestV3 ? createOffscreen() : null;
+
     const initData = await loadStateFromPersistence();
 
     const initState = initData.data;
@@ -293,6 +296,7 @@ async function initialize() {
       {},
       isFirstMetaMaskControllerSetup,
       initData.meta,
+      offscreenPromise,
     );
     if (!isManifestV3) {
       await loadPhishingWarningPage();
@@ -507,13 +511,15 @@ function emitDappViewedMetricEvent(
  * @param {object} overrides - object with callbacks that are allowed to override the setup controller logic
  * @param isFirstMetaMaskControllerSetup
  * @param {object} stateMetadata - Metadata about the initial state and migrations, including the most recent migration version
+ * @param {Promise<void>} offscreenPromise - A promise that resolves when the offscreen document has finished initialization.
  */
 export function setupController(
   initState,
   initLangCode,
   overrides,
   isFirstMetaMaskControllerSetup,
   stateMetadata,
+  offscreenPromise,
 ) {
   //
   // MetaMask Controller
@@ -542,6 +548,7 @@ export function setupController(
     isFirstMetaMaskControllerSetup,
     currentMigrationVersion: stateMetadata.version,
     featureFlags: {},
+    offscreenPromise,
   });
 
   setupEnsIpfsResolver({

app/scripts/metamask-controller.js

@@ -374,6 +374,8 @@ export default class MetamaskController extends EventEmitter {
     // the only thing that uses controller connections are open metamask UI instances
     this.activeControllerConnections = 0;
 
+    this.offscreenPromise = opts.offscreenPromise ?? Promise.resolve();
+
     this.getRequestAccountTabIds = opts.getRequestAccountTabIds;
     this.getOpenMetamaskTabsIds = opts.getOpenMetamaskTabsIds;
 
@@ -4069,6 +4071,10 @@ export default class MetamaskController extends EventEmitter {
    */
   async submitPassword(password) {
     const { completedOnboarding } = this.onboardingController.store.getState();
+
+    // Before attempting to unlock the keyrings, we need the offscreen to have loaded.
+    await this.offscreenPromise;
+
     await this.keyringController.submitPassword(password);
 
     ///: BEGIN:ONLY_INCLUDE_IF(build-mmi)

app/scripts/offscreen.js

@@ -0,0 +1,44 @@
+import { OffscreenCommunicationTarget } from '../../shared/constants/offscreen-communication';
+
+/**
+ * Creates an offscreen document that can be used to load additional scripts
+ * and iframes that can communicate with the extension through the chrome
+ * runtime API. Only one offscreen document may exist, so any iframes required
+ * by extension can be embedded in the offscreen.html file. See the offscreen
+ * folder for more details.
+ */
+export async function createOffscreen() {
+  const { chrome } = globalThis;
+  if (!chrome.offscreen || (await chrome.offscreen.hasDocument())) {
+    return;
+  }
+
+  const loadPromise = new Promise((resolve) => {
+    const messageListener = (msg) => {
+      if (
+        msg.target === OffscreenCommunicationTarget.extensionMain &&
+        msg.isBooted
+      ) {
+        chrome.runtime.onMessage.removeListener(messageListener);
+        resolve();
+      }
+    };
+    chrome.runtime.onMessage.addListener(messageListener);
+  });
+
+  await chrome.offscreen.createDocument({
+    url: './offscreen.html',
+    reasons: ['IFRAME_SCRIPTING'],
+    justification:
+      'Used for Hardware Wallet and Snaps scripts to communicate with the extension.',
+  });
+
+  // In case we are in a bad state where the offscreen document is not loading, timeout and let execution continue.
+  const timeoutPromise = new Promise((resolve) => {
+    setTimeout(resolve, 5000);
+  });
+
+  await Promise.race([loadPromise, timeoutPromise]);
+
+  console.debug('Offscreen iframe loaded');
+}

offscreen/scripts/offscreen.ts

@@ -1,5 +1,6 @@
 import { BrowserRuntimePostMessageStream } from '@metamask/post-message-stream';
 import { ProxySnapExecutor } from '@metamask/snaps-execution-environments';
+import { OffscreenCommunicationTarget } from '../../shared/constants/offscreen-communication';
 import initLedger from './ledger';
 import initTrezor from './trezor';
 import initLattice from './lattice';
@@ -20,3 +21,8 @@ const parentStream = new BrowserRuntimePostMessageStream({
 });
 
 ProxySnapExecutor.initialize(parentStream, './snaps/index.html');
+
+chrome.runtime.sendMessage({
+  target: OffscreenCommunicationTarget.extensionMain,
+  isBooted: true,
+});

shared/constants/offscreen-communication.ts

@@ -7,6 +7,7 @@ export enum OffscreenCommunicationTarget {
   ledgerOffscreen = 'ledger-offscreen',
   latticeOffscreen = 'lattice-offscreen',
   extension = 'extension-offscreen',
+  extensionMain = 'extension',
 }
 
 /**