From 4cac9d185bf522c69b0a262c8778990e3cf2b14d Mon Sep 17 00:00:00 2001
From: Raymond Hill <rhill@raymondhill.net>
Date: Mon, 2 Oct 2023 08:42:03 -0400
Subject: [PATCH] Reduce race conditions in scriptlet injection on Firefox

This is done by taking advantage through Firefox-specific
contentScripts.register() API:

https://developer.mozilla.org/en-US/docs/Mozilla/Add-ons/WebExtensions/API/contentScripts
---
 src/js/scriptlet-filtering.js | 82 ++++++++++++++++++++++++++++-------
 1 file changed, 66 insertions(+), 16 deletions(-)

diff --git a/src/js/scriptlet-filtering.js b/src/js/scriptlet-filtering.js
index 19c8d1308..5b50ae675 100644
--- a/src/js/scriptlet-filtering.js
+++ b/src/js/scriptlet-filtering.js
@@ -19,6 +19,8 @@
     Home: https://github.com/gorhill/uBlock
 */
 
+/* globals browser */
+
 'use strict';
 
 /******************************************************************************/
@@ -61,6 +63,55 @@ const scriptletFilteringEngine = {
     },
 };
 
+const contentScriptRegisterer = new (class {
+    constructor() {
+        this.hostnameToDetails = new Map();
+    }
+    register(hostname, code) {
+        if ( browser.contentScripts === undefined ) { return false; }
+        const details = this.hostnameToDetails.get(hostname);
+        if ( details !== undefined ) {
+            if ( code === details.code ) {
+                return details.handle instanceof Promise === false;
+            }
+            details.handle.unregister();
+            this.hostnameToDetails.delete(hostname);
+        }
+        const promise = browser.contentScripts.register({
+            js: [ { code } ],
+            allFrames: true,
+            matches: [ `*://*.${hostname}/*` ],
+            matchAboutBlank: true,
+            runAt: 'document_start',
+        }).then(handle => {
+            this.hostnameToDetails.set(hostname, { handle, code });
+        });
+        this.hostnameToDetails.set(hostname, { handle: promise, code });
+        return false;
+    }
+    unregister(hostname) {
+        if ( this.hostnameToDetails.size === 0 ) { return; }
+        const details = this.hostnameToDetails.get(hostname);
+        if ( details === undefined ) { return; }
+        this.hostnameToDetails.delete(hostname);
+        this.unregisterHandle(details.handle);
+    }
+    reset() {
+        if ( this.hostnameToDetails.size === 0 ) { return; }
+        for ( const details of this.hostnameToDetails.values() ) {
+            this.unregisterHandle(details.handle);
+        }
+        this.hostnameToDetails.clear();
+    }
+    unregisterHandle(handle) {
+        if ( handle instanceof Promise ) {
+            handle.then(handle => { handle.unregister(); });
+        } else {
+            handle.unregister();
+        }
+    }
+})();
+
 // Purpose of `contentscriptCode` below is too programmatically inject
 // content script code which only purpose is to inject scriptlets. This
 // essentially does the same as what uBO's declarative content script does,
@@ -141,7 +192,6 @@ const isolatedWorldInjector = (( ) => {
     return {
         parts,
         jsonSlot: parts.indexOf('json-slot'),
-        scriptletSlot: parts.indexOf('scriptlet-slot'),
         assemble: function(hostname, scriptlets) {
             this.parts[this.jsonSlot] = JSON.stringify({ hostname });
             const code = this.parts.join('');
@@ -239,6 +289,7 @@ scriptletFilteringEngine.logFilters = function(tabId, url, filters) {
 scriptletFilteringEngine.reset = function() {
     scriptletDB.clear();
     duplicates.clear();
+    contentScriptRegisterer.reset();
     scriptletCache.reset();
     acceptedCount = 0;
     discardedCount = 0;
@@ -441,25 +492,24 @@ scriptletFilteringEngine.injectNow = function(details) {
     request.domain = domainFromHostname(request.hostname);
     request.entity = entityFromDomain(request.domain);
     const scriptletDetails = this.retrieve(request);
-    if ( scriptletDetails === undefined ) { return; }
+    if ( scriptletDetails === undefined ) {
+        contentScriptRegisterer.unregister(request.hostname);
+        return;
+    }
+    const contentScript = [];
+    if ( µb.hiddenSettings.debugScriptletInjector ) {
+        contentScript.push('debugger');
+    }
     const { mainWorld = '', isolatedWorld = '', filters } = scriptletDetails;
     if ( mainWorld !== '' ) {
-        let code = mainWorldInjector.assemble(request.hostname, mainWorld, filters);
-        if ( µb.hiddenSettings.debugScriptletInjector ) {
-            code = 'debugger;\n' + code;
-        }
-        vAPI.tabs.executeScript(details.tabId, {
-            code,
-            frameId: details.frameId,
-            matchAboutBlank: true,
-            runAt: 'document_start',
-        });
+        contentScript.push(mainWorldInjector.assemble(request.hostname, mainWorld, filters));
     }
     if ( isolatedWorld !== '' ) {
-        let code = isolatedWorldInjector.assemble(request.hostname, isolatedWorld);
-        if ( µb.hiddenSettings.debugScriptletInjector ) {
-            code = 'debugger;\n' + code;
-        }
+        contentScript.push(isolatedWorldInjector.assemble(request.hostname, isolatedWorld));
+    }
+    const code = contentScript.join('\n\n');
+    const isAlreadyInjected = contentScriptRegisterer.register(request.hostname, code);
+    if ( isAlreadyInjected !== true ) {
         vAPI.tabs.executeScript(details.tabId, {
             code,
             frameId: details.frameId,