From ac86ec33724d4ef0a5d2d6f28bbf39bbcfbd55f8 Mon Sep 17 00:00:00 2001
From: Krzysztof Modras <chrmod@chrmod.net>
Date: Tue, 19 Nov 2024 14:03:13 +0100
Subject: [PATCH] POC: on Firefox inject scriplets with browser.contentScripts

---
 src/background/adblocker.js | 93 ++++++++++++++++++++++++++++++++++---
 1 file changed, 87 insertions(+), 6 deletions(-)

diff --git a/src/background/adblocker.js b/src/background/adblocker.js
index 6d8c42f8a..8344fc8a1 100644
--- a/src/background/adblocker.js
+++ b/src/background/adblocker.js
@@ -86,6 +86,9 @@ export async function reloadMainEngine() {
     await engines.create(engines.MAIN_ENGINE);
     console.info('[adblocker] Main engine reloaded with no filters');
   }
+  if (__PLATFORM__ === 'firefox') {
+    contentScripts.unregisterAll();
+  }
 }
 
 let updating = false;
@@ -163,7 +166,50 @@ export const setup = asyncSetup('adblocker', [
   ),
 ]);
 
-function injectScriptlets(filters, tabId, frameId) {
+const contentScripts = (() => {
+  const map = new Map();
+  return {
+    async register(hostname, code) {
+      this.unregister(hostname);
+      try {
+        const contentScript = await browser.contentScripts.register({
+          js: [
+            {
+              code,
+            },
+          ],
+          allFrames: true,
+          matches: [`https://*.${hostname}/*`, `http://*.${hostname}/*`],
+          matchAboutBlank: true,
+          matchOriginAsFallback: true,
+          runAt: 'document_start',
+        });
+        map.set(hostname, contentScript);
+      } catch (e) {
+        console.warn(e);
+        this.unregister(hostname);
+      }
+    },
+    isRegistered(hostname) {
+      return map.has(hostname);
+    },
+    unregister(hostname) {
+      const contentScript = map.get(hostname);
+      if (contentScript) {
+        contentScript.unregister();
+        map.delete(hostname);
+      }
+    },
+    unregisterAll() {
+      for (const hostname of map.keys()) {
+        this.unregister(hostname);
+      }
+    },
+  };
+})();
+
+function injectScriptlets(filters, tabId, frameId, hostname) {
+  let contentScript = '';
   for (const filter of filters) {
     const parsed = filter.parseScript();
 
@@ -177,12 +223,19 @@ function injectScriptlets(filters, tabId, frameId) {
 
     const scriptletName = `${parsed.name}${parsed.name.endsWith('.js') ? '' : '.js'}`;
     const scriptlet = scriptlets[scriptletName];
+    const func = scriptlet.func;
+    const args = parsed.args.map((arg) => decodeURIComponent(arg));
 
     if (!scriptlet) {
       console.warn('[adblocker] unknown scriptlet with name:', scriptletName);
       continue;
     }
 
+    if (__PLATFORM__ === 'firefox') {
+      contentScript += `(function () { ${func.toString()} })(...${JSON.stringify(args)})`
+      continue;
+    }
+
     chrome.scripting.executeScript(
       {
         injectImmediately: true,
@@ -193,8 +246,8 @@ function injectScriptlets(filters, tabId, frameId) {
           tabId,
           frameIds: [frameId],
         },
-        func: scriptlet.func,
-        args: parsed.args.map((arg) => decodeURIComponent(arg)),
+        func,
+        args,
       },
       () => {
         if (chrome.runtime.lastError) {
@@ -203,6 +256,19 @@ function injectScriptlets(filters, tabId, frameId) {
       },
     );
   }
+
+  if (__PLATFORM__ === 'firefox') {
+    if (filters.length === 0) {
+      contentScripts.unregister(hostname);
+    } else if (!contentScripts.isRegistered(hostname)) {
+      contentScripts.register(
+        hostname,
+        contentScript,
+      );
+    } else {
+      // do nothing if already registered
+    }
+  }
 }
 
 function injectStyles(styles, tabId, frameId) {
@@ -224,6 +290,9 @@ function injectStyles(styles, tabId, frameId) {
 }
 
 async function injectCosmetics(details, config) {
+  const isBootstrap = config.bootstrap;
+  const scriptletsOnly = Boolean(config.scriptletsOnly);
+
   try {
     setup.pending && (await setup.pending);
   } catch (e) {
@@ -237,6 +306,8 @@ async function injectCosmetics(details, config) {
   const domain = parsed.domain || '';
   const hostname = parsed.hostname || '';
 
+  if (scriptletsOnly && contentScripts.isRegistered(hostname)) return;
+
   if (isPaused(options, hostname)) return;
 
   const tabHostname = tabStats.get(tabId)?.hostname;
@@ -245,7 +316,6 @@ async function injectCosmetics(details, config) {
   }
 
   const engine = engines.get(engines.MAIN_ENGINE);
-  const isBootstrap = config.bootstrap;
 
   {
     const { matches } = engine.matchCosmeticFilters({
@@ -282,8 +352,12 @@ async function injectCosmetics(details, config) {
       }
     }
 
-    if (isBootstrap && scriptFilters.length > 0) {
-      injectScriptlets(scriptFilters, tabId, frameId);
+    if (isBootstrap) {
+      injectScriptlets(scriptFilters, tabId, frameId, hostname);
+    }
+
+    if (scriptletsOnly) {
+      return;
     }
 
     const { styles } = engine.injectCosmeticFilters(styleFilters, {
@@ -377,6 +451,13 @@ function isTrusted(request, type) {
 }
 
 if (__PLATFORM__ === 'firefox') {
+  chrome.webNavigation.onBeforeNavigate.addListener(
+    (details) => {
+      injectCosmetics(details, { bootstrap: true, scriptletsOnly: true });
+    },
+    { url: [{ urlPrefix: 'http://' }, { urlPrefix: 'https://' }] },
+  );
+
   function isExtensionRequest(details) {
     return (
       (details.tabId === -1 && details.url.startsWith('moz-extension://')) ||