Avoid duplicate injection on MV3 worker restart (#73)

Author: fregante

package-lock.json

@@ -77,6 +77,7 @@
 		"content-scripts-register-polyfill": "^4.0.2",
 		"webext-content-scripts": "^2.6.1",
 		"webext-detect-page": "^5.0.1",
+		"webext-events": "^3.0.0",
 		"webext-patterns": "^1.4.0",
 		"webext-permissions": "^3.1.3",
 		"webext-polyfill-kinda": "^1.0.2"

package.json

@@ -44,7 +44,7 @@ navigator.importScripts('webext-dynamic-content-scripts.js');
 ```json
 // example manifest.json
 {
-	"permissions": ["scripting"],
+	"permissions": ["scripting", "storage"],
 	"optional_host_permissions": ["*://*/*"],
 	"background": {
 		"service_worker": "background.worker.js"

readme.md

@@ -1,3 +1,3 @@
 import {init} from './lib.js';
 
-void init();
+init();

source/index.ts

@@ -2,6 +2,8 @@ import {getTabsByUrl, injectContentScript} from 'webext-content-scripts';
 
 type ManifestContentScripts = NonNullable<chrome.runtime.Manifest['content_scripts']>;
 
+// May not be needed in the future in Firefox
+// https://bugzilla.mozilla.org/show_bug.cgi?id=1458947
 export async function injectToExistingTabs(
 	origins: string[],
 	scripts: ManifestContentScripts,

source/inject-to-existing-tabs.ts

@@ -3,11 +3,15 @@ import {
 	describe, it, vi, beforeEach, expect,
 } from 'vitest';
 import {queryAdditionalPermissions} from 'webext-permissions';
+import {onExtensionStart} from 'webext-events';
 import {init} from './lib.js';
 import {injectToExistingTabs} from './inject-to-existing-tabs.js';
 import {registerContentScript} from './register-content-script-shim.js';
 
+type AsyncFunction = () => void | Promise<void>;
+
 vi.mock('webext-permissions');
+vi.mock('webext-events');
 vi.mock('./register-content-script-shim.js');
 vi.mock('./inject-to-existing-tabs.js');
 
@@ -21,6 +25,7 @@ const baseManifest: chrome.runtime.Manifest = {
 			matches: ['https://content-script.example.com/*'],
 		},
 	],
+	permissions: ['storage'],
 	host_permissions: ['https://permission-only.example.com/*'],
 	optional_host_permissions: ['*://*/*'],
 };
@@ -34,6 +39,17 @@ const queryAdditionalPermissionsMock = vi.mocked(queryAdditionalPermissions);
 const injectToExistingTabsMock = vi.mocked(injectToExistingTabs);
 const registerContentScriptMock = vi.mocked(registerContentScript);
 
+const callbacks = new Set<AsyncFunction>();
+
+vi.mocked(onExtensionStart.addListener).mockImplementation((callback: AsyncFunction) => {
+	callbacks.add(callback);
+});
+
+async function simulateExtensionStart() {
+	await Promise.all(Array.from(callbacks).map(async callback => callback()));
+	callbacks.clear();
+}
+
 beforeEach(() => {
 	registerContentScriptMock.mockClear();
 	injectToExistingTabsMock.mockClear();
@@ -43,7 +59,9 @@ beforeEach(() => {
 
 describe('init', () => {
 	it('it should register the listeners and start checking permissions', async () => {
-		await init();
+		init();
+		await simulateExtensionStart();
+
 		expect(queryAdditionalPermissionsMock).toHaveBeenCalled();
 		expect(injectToExistingTabsMock).toHaveBeenCalledWith(
 			additionalPermissions.origins,
@@ -59,13 +77,16 @@ describe('init', () => {
 		const manifest = structuredClone(baseManifest);
 		delete manifest.content_scripts;
 		chrome.runtime.getManifest.mockReturnValue(manifest);
-		await expect(init()).rejects.toMatchInlineSnapshot('[Error: webext-dynamic-content-scripts tried to register scripts on the new host permissions, but no content scripts were found in the manifest.]');
+		init();
+		await expect(simulateExtensionStart).rejects
+			.toThrowErrorMatchingInlineSnapshot('[Error: webext-dynamic-content-scripts tried to register scripts on the new host permissions, but no content scripts were found in the manifest.]');
 	});
 });
 
 describe('init - registerContentScript', () => {
 	it('should register the manifest scripts on new permissions', async () => {
-		await init();
+		init();
+		await simulateExtensionStart();
 		expect(registerContentScriptMock).toMatchSnapshot();
 	});
 
@@ -78,7 +99,8 @@ describe('init - registerContentScript', () => {
 			permissions: [],
 		});
 
-		await init();
+		init();
+		await simulateExtensionStart();
 		expect(registerContentScriptMock).toMatchSnapshot();
 	});
 
@@ -90,7 +112,8 @@ describe('init - registerContentScript', () => {
 		});
 		chrome.runtime.getManifest.mockReturnValue(manifest);
 
-		await init();
+		init();
+		await simulateExtensionStart();
 		expect(registerContentScriptMock).toMatchSnapshot();
 	});
 });

source/lib.test.ts

@@ -1,4 +1,5 @@
 import {queryAdditionalPermissions} from 'webext-permissions';
+import {onExtensionStart} from 'webext-events';
 import {excludeDuplicateFiles} from './deduplicator.js';
 import {injectToExistingTabs} from './inject-to-existing-tabs.js';
 import {registerContentScript} from './register-content-script-shim.js';
@@ -13,25 +14,28 @@ function makePathRelative(file: string): string {
 	return new URL(file, location.origin).pathname;
 }
 
-// Automatically register the content scripts on the new origins
-async function registerOnOrigins({
-	origins: newOrigins,
-}: chrome.permissions.Permissions): Promise<void> {
-	if (!newOrigins?.length) {
-		return;
-	}
-
+function getContentScripts() {
 	const {content_scripts: rawManifest, manifest_version: manifestVersion} = chrome.runtime.getManifest();
 
 	if (!rawManifest) {
 		throw new Error('webext-dynamic-content-scripts tried to register scripts on the new host permissions, but no content scripts were found in the manifest.');
 	}
 
-	const cleanManifest = excludeDuplicateFiles(rawManifest, {warn: manifestVersion === 2});
+	return excludeDuplicateFiles(rawManifest, {warn: manifestVersion === 2});
+}
+
+// Automatically register the content scripts on the new origins
+async function registerOnOrigins(
+	origins: string[],
+	contentScripts: ReturnType<typeof getContentScripts>,
+): Promise<void> {
+	if (origins.length === 0) {
+		return;
+	}
 
 	// Register one at a time to allow removing one at a time as well
-	for (const origin of newOrigins) {
-		for (const config of cleanManifest) {
+	for (const origin of origins) {
+		for (const config of contentScripts) {
 			const registeredScript = registerContentScript({
 				// Always convert paths here because we don't know whether Firefox MV3 will accept full URLs
 				js: config.js?.map(file => makePathRelative(file)),
@@ -44,14 +48,10 @@ async function registerOnOrigins({
 			registeredScripts.set(origin, registeredScript);
 		}
 	}
-
-	// May not be needed in the future in Firefox
-	// https://bugzilla.mozilla.org/show_bug.cgi?id=1458947
-	void injectToExistingTabs(newOrigins, cleanManifest);
 }
 
-function handleNewPermissions(permissions: chrome.permissions.Permissions) {
-	void registerOnOrigins(permissions);
+async function handleNewPermissions({origins}: chrome.permissions.Permissions) {
+	await enableOnOrigins(origins);
 }
 
 async function handledDroppedPermissions({origins}: chrome.permissions.Permissions) {
@@ -68,12 +68,28 @@ async function handledDroppedPermissions({origins}: chrome.permissions.Permissio
 	}
 }
 
-export async function init() {
+async function enableOnOrigins(origins: string[] | undefined) {
+	if (!origins?.length) {
+		return;
+	}
+
+	const contentScripts = getContentScripts();
+	await Promise.all([
+		injectToExistingTabs(origins, contentScripts),
+		registerOnOrigins(origins, contentScripts),
+	]);
+}
+
+async function registerExistingOrigins() {
+	const {origins} = await queryAdditionalPermissions({
+		strictOrigins: false,
+	});
+
+	await enableOnOrigins(origins);
+}
+
+export function init() {
 	chrome.permissions.onRemoved.addListener(handledDroppedPermissions);
 	chrome.permissions.onAdded.addListener(handleNewPermissions);
-	await registerOnOrigins(
-		await queryAdditionalPermissions({
-			strictOrigins: false,
-		}),
-	);
+	onExtensionStart.addListener(registerExistingOrigins);
 }

source/lib.ts

@@ -2,7 +2,7 @@
 	"name": "webext-dynamic-content-scripts-mv3",
 	"version": "0.0.0",
 	"manifest_version": 3,
-	"permissions": ["webNavigation", "scripting", "contextMenus", "activeTab"],
+	"permissions": ["webNavigation", "scripting", "contextMenus", "activeTab", "storage"],
 	"host_permissions": [
 		"https://dynamic-ephiframe.vercel.app/*",
 		"https://accepted-ephiframe.vercel.app/*"