fix(security): Use import.meta.hot for communication with Vite 6.0.9+ (and 5.4.14+) dev server

packages/wxt/package.json

@@ -67,6 +67,10 @@
     "./modules": {
       "types": "./dist/modules.d.ts",
       "default": "./dist/modules.mjs"
+    },
+    "./background-client": {
+      "types": "./dist/background-client.d.ts",
+      "default": "./dist/background-client.mjs"
     }
   },
   "scripts": {
@@ -123,7 +127,7 @@
     "publish-browser-extension": "^2.2.2",
     "scule": "^1.3.0",
     "unimport": "^3.13.1",
-    "vite": "^5.0.0 || ^6.0.0",
+    "vite": "^5.0.0 || <=6.0.8",
     "vite-node": "^2.1.4",
     "web-ext-run": "^0.2.1",
     "webextension-polyfill": "^0.12.0"

packages/wxt/src/@types/globals.d.ts

@@ -1,8 +1,15 @@
+/**
+ * Replaced by an import to `wxt/background-client`, which changes based on the
+ * manifest version being targeted.
+ */
+declare const __WXT_BACKGROUND_CLIENT_IMPORT__: void;
+
+// Dev server globals, replaced during development
 declare const __DEV_SERVER_PROTOCOL__: string;
 declare const __DEV_SERVER_HOSTNAME__: string;
 declare const __DEV_SERVER_PORT__: string;
 
-// Globals defined by the vite-plugins/devServerGlobals.ts and utils/globals.ts
+// Global env vars provided by WXT
 interface ImportMetaEnv {
   readonly COMMAND: WxtCommand;
   readonly MANIFEST_VERSION: 2 | 3;

packages/wxt/src/virtual/utils/reload-content-scripts.ts → packages/wxt/src/background-client.ts

@@ -1,9 +1,56 @@
+/**
+ * This file contains all the code required for the background to communicate
+ * with the dev server during development:
+ *
+ * - Reload extension or parts of extension when files are saved.
+ * - Keep MV3 service worker alive indefinitely.
+ *
+ * It is not imported directly (`import "wxt/background-client"`), but from
+ * the dev server (`import "http://localhost:3000/@id/wxt/background-client"`)
+ * to ensure `import.meta.hot` is defined by Vite.
+ */
+
 import { browser } from 'wxt/browser';
-import { logger } from '../../sandbox/utils/logger';
+import { logger } from './sandbox/utils/logger';
 import { MatchPattern } from 'wxt/sandbox';
-import type { ReloadContentScriptPayload } from '../../sandbox/dev-server-websocket';
 
-export function reloadContentScript(payload: ReloadContentScriptPayload) {
+if (import.meta.hot) {
+  import.meta.hot.on('wxt:reload-extension', () => browser.runtime.reload());
+  import.meta.hot.on('wxt:reload-content-script', (event) =>
+    reloadContentScript(event),
+  );
+
+  if (import.meta.env.MANIFEST_VERSION === 3) {
+    let backgroundInitialized = false;
+    // Tell the server the background script is loaded and ready to receive events
+    import.meta.hot.on('vite:ws:connect', () => {
+      if (backgroundInitialized) return;
+
+      import.meta.hot?.send('wxt:background-initialized');
+      backgroundInitialized = true;
+    });
+
+    // Web Socket will disconnect if the service worker is killed. Supposedly,
+    // just having a web socket connection active should keep the service worker
+    // alive, but when this was originally implemented on older versions of
+    // Chrome, that was not true. So this code has stayed around.
+    // See: https://developer.chrome.com/blog/longer-esw-lifetimes/
+    setInterval(async () => {
+      // Calling an async browser API resets the service worker's timeout
+      await browser.runtime.getPlatformInfo();
+    }, 5e3);
+  }
+
+  browser.commands.onCommand.addListener((command) => {
+    if (command === 'wxt:reload-extension') {
+      browser.runtime.reload();
+    }
+  });
+} else {
+  console.error('[wxt] HMR context, import.meta.hot, not found');
+}
+
+function reloadContentScript(payload: ReloadContentScriptPayload) {
   const manifest = browser.runtime.getManifest();
   if (manifest.manifest_version == 2) {
     void reloadContentScriptMv2(payload);
@@ -12,7 +59,7 @@ export function reloadContentScript(payload: ReloadContentScriptPayload) {
   }
 }
 
-export async function reloadContentScriptMv3({
+async function reloadContentScriptMv3({
   registration,
   contentScript,
 }: ReloadContentScriptPayload) {
@@ -25,9 +72,7 @@ export async function reloadContentScriptMv3({
 
 type ContentScript = ReloadContentScriptPayload['contentScript'];
 
-export async function reloadManifestContentScriptMv3(
-  contentScript: ContentScript,
-) {
+async function reloadManifestContentScriptMv3(contentScript: ContentScript) {
   const id = `wxt:${contentScript.js![0]}`;
   logger.log('Reloading content script:', contentScript);
   const registered = await browser.scripting.getRegisteredContentScripts();
@@ -46,9 +91,7 @@ export async function reloadManifestContentScriptMv3(
   await reloadTabsForContentScript(contentScript);
 }
 
-export async function reloadRuntimeContentScriptMv3(
-  contentScript: ContentScript,
-) {
+async function reloadRuntimeContentScriptMv3(contentScript: ContentScript) {
   logger.log('Reloading content script:', contentScript);
   const registered = await browser.scripting.getRegisteredContentScripts();
   logger.debug('Existing scripts:', registered);
@@ -92,8 +135,15 @@ async function reloadTabsForContentScript(contentScript: ContentScript) {
   );
 }
 
-export async function reloadContentScriptMv2(
-  _payload: ReloadContentScriptPayload,
-) {
+async function reloadContentScriptMv2(_payload: ReloadContentScriptPayload) {
   throw Error('TODO: reloadContentScriptMv2');
 }
+
+interface ReloadContentScriptPayload {
+  registration?: 'manifest' | 'runtime';
+  contentScript: {
+    matches: string[];
+    js?: string[];
+    css?: string[];
+  };
+}

packages/wxt/src/core/builders/vite/plugins/devServerGlobals.ts

@@ -11,10 +11,23 @@ export function devServerGlobals(
   return {
     name: 'wxt:dev-server-globals',
     config() {
-      if (server == null || config.command == 'build') return;
+      if (server == null || config.command == 'build')
+        return {
+          define: {
+            __WXT_BACKGROUND_CLIENT_IMPORT__: '',
+          },
+        };
+
+      const wxtBackgroundClientUrl = `http://${server.hostname}:${server.port}/@id/wxt/background-client`;
+      console.log({ wxtBackgroundClientUrl });
 
       return {
         define: {
+          __WXT_BACKGROUND_CLIENT_IMPORT__: JSON.stringify(
+            config.manifestVersion === 2
+              ? `import(/* @vite-ignore */ "${wxtBackgroundClientUrl}")`
+              : `/* @vite-ignore */\nimport "${wxtBackgroundClientUrl}"`,
+          ),
           __DEV_SERVER_PROTOCOL__: JSON.stringify('ws:'),
           __DEV_SERVER_HOSTNAME__: JSON.stringify(server.hostname),
           __DEV_SERVER_PORT__: JSON.stringify(server.port),

packages/wxt/src/core/utils/manifest.ts

@@ -116,8 +116,11 @@ export async function generateManifest(
 
   addEntrypoints(manifest, entrypoints, buildOutput);
 
-  if (wxt.config.command === 'serve') addDevModeCsp(manifest);
-  if (wxt.config.command === 'serve') addDevModePermissions(manifest);
+  if (wxt.config.command === 'serve') {
+    addDevModeCsp(manifest);
+    addDevModePermissions(manifest);
+    setDevModeBackgroundToEsm(manifest);
+  }
 
   // TODO: Remove in v1
   wxt.config.transformManifest?.(manifest);
@@ -509,6 +512,15 @@ function addDevModePermissions(manifest: Manifest.WebExtensionManifest) {
   if (wxt.config.manifestVersion === 3) addPermission(manifest, 'scripting');
 }
 
+/** MV3 service worker must by `type: "module"` to connect to dev server. */
+function setDevModeBackgroundToEsm(manifest: Manifest.WebExtensionManifest) {
+  if (wxt.config.manifestVersion === 3) {
+    // Background must be ESM for it to be able to connect to the dev server
+    (manifest as any).background ??= {};
+    (manifest as any).background.type = 'module';
+  }
+}
+
 /**
  * Returns the bundle paths to CSS files associated with a list of content scripts, or undefined if
  * there is no associated CSS.

packages/wxt/src/virtual/background-entrypoint.ts

@@ -1,40 +1,8 @@
 import definition from 'virtual:user-background-entrypoint';
 import { initPlugins } from 'virtual:wxt-plugins';
-import { getDevServerWebSocket } from '../sandbox/dev-server-websocket';
 import { logger } from '../sandbox/utils/logger';
-import { browser } from 'wxt/browser';
-import { keepServiceWorkerAlive } from './utils/keep-service-worker-alive';
-import { reloadContentScript } from './utils/reload-content-scripts';
-
-if (import.meta.env.COMMAND === 'serve') {
-  try {
-    const ws = getDevServerWebSocket();
-    ws.addWxtEventListener('wxt:reload-extension', () => {
-      browser.runtime.reload();
-    });
-    ws.addWxtEventListener('wxt:reload-content-script', (event) => {
-      reloadContentScript(event.detail);
-    });
-
-    if (import.meta.env.MANIFEST_VERSION === 3) {
-      // Tell the server the background script is loaded and ready to go
-      ws.addEventListener('open', () =>
-        ws.sendCustom('wxt:background-initialized'),
-      );
-
-      // Web Socket will disconnect if the service worker is killed
-      keepServiceWorkerAlive();
-    }
-  } catch (err) {
-    logger.error('Failed to setup web socket connection with dev server', err);
-  }
-
-  browser.commands.onCommand.addListener((command) => {
-    if (command === 'wxt:reload-extension') {
-      browser.runtime.reload();
-    }
-  });
-}
+console.log('__WXT_BACKGROUND_CLIENT_IMPORT__');
+__WXT_BACKGROUND_CLIENT_IMPORT__;
 
 let result;
 

packages/wxt/src/virtual/reload-html.ts

@@ -1,14 +1,6 @@
-import { logger } from '../sandbox/utils/logger';
-import { getDevServerWebSocket } from '../sandbox/dev-server-websocket';
-
-if (import.meta.env.COMMAND === 'serve') {
-  try {
-    const ws = getDevServerWebSocket();
-    ws.addWxtEventListener('wxt:reload-page', (event) => {
-      // "popup.html" === "/popup.html".substring(1)
-      if (event.detail === location.pathname.substring(1)) location.reload();
-    });
-  } catch (err) {
-    logger.error('Failed to setup web socket connection with dev server', err);
-  }
+if (import.meta.hot) {
+  import.meta.hot.on('wxt:reload-page', (event) => {
+    // "popup.html" === "/popup.html".substring(1)
+    if (event.detail === location.pathname.substring(1)) location.reload();
+  });
 }