From e3f3ed20e690a53bde98be9bcf371a591ca102e1 Mon Sep 17 00:00:00 2001
From: Ajay <dev@ajay.app>
Date: Fri, 11 Aug 2023 11:39:06 -0400
Subject: [PATCH] Enable non persistent background page on Firefox

---
 manifest/firefox-manifest-extra.json |  8 +++-
 manifest/safari-manifest-extra.json  |  5 ++-
 src/background.ts                    | 59 +++++++++++++++++++++-------
 src/types.ts                         |  4 +-
 src/utils.ts                         | 24 +++++------
 5 files changed, 68 insertions(+), 32 deletions(-)

diff --git a/manifest/firefox-manifest-extra.json b/manifest/firefox-manifest-extra.json
index f481cb1c..714fef58 100644
--- a/manifest/firefox-manifest-extra.json
+++ b/manifest/firefox-manifest-extra.json
@@ -3,5 +3,11 @@
     "gecko": {
       "id": "sponsorBlocker@ajay.app"
     }
-  }
+  },
+  "background": {
+    "persistent": false
+  },
+  "permissions": [
+    "scripting"
+  ]
 }
diff --git a/manifest/safari-manifest-extra.json b/manifest/safari-manifest-extra.json
index 30f46118..bf18bbe5 100644
--- a/manifest/safari-manifest-extra.json
+++ b/manifest/safari-manifest-extra.json
@@ -1,5 +1,8 @@
 {
   "background": {
     "persistent": false
-  }
+  },
+  "permissions": [
+    "scripting"
+  ]
 }
diff --git a/src/background.ts b/src/background.ts
index 7ed4d2b2..5b0bceb8 100644
--- a/src/background.ts
+++ b/src/background.ts
@@ -148,27 +148,58 @@ chrome.runtime.onInstalled.addListener(function () {
  *
  * @param {JSON} options
  */
-function registerFirefoxContentScript(options: Registration) {
-    const oldRegistration = contentScriptRegistrations[options.id];
-    if (oldRegistration) oldRegistration.unregister();
+async function registerFirefoxContentScript(options: Registration) {
+    await unregisterFirefoxContentScript(options.id);
+
+    if ("scripting" in chrome && "getRegisteredContentScripts" in chrome.scripting) {
+        await chrome.scripting.registerContentScripts([{
+            id: options.id,
+            runAt: "document_start",
+            matches: options.matches,
+            allFrames: options.allFrames,
+            js: options.js,
+            css: options.css,
+            persistAcrossSessions: true,
+        }]);
+    } else {
+        chrome.contentScripts.register({
+            allFrames: options.allFrames,
+            js: options.js?.map?.(file => ({file})),
+            css: options.css?.map?.(file => ({file})),
+            matches: options.matches
+        }).then((registration) => void (contentScriptRegistrations[options.id] = registration));
+    }
 
-    chrome.contentScripts.register({
-        allFrames: options.allFrames,
-        js: options.js,
-        css: options.css,
-        matches: options.matches
-    }).then((registration) => void (contentScriptRegistrations[options.id] = registration));
 }
 
 /**
  * Only works on Firefox.
  * Firefox requires that this is handled by the background script
- *
  */
-function unregisterFirefoxContentScript(id: string) {
-    if (contentScriptRegistrations[id]) {
-        contentScriptRegistrations[id].unregister();
-        delete contentScriptRegistrations[id];
+async function unregisterFirefoxContentScript(id: string) {
+    if ("scripting" in chrome && "getRegisteredContentScripts" in chrome.scripting) {
+        // Bug in Firefox where you need to use browser namespace for this call
+        const getContentScripts = async (filter: browser.scripting.ContentScriptFilter) => {
+            if (isFirefoxOrSafari()) {
+                return await browser.scripting.getRegisteredContentScripts(filter);
+            } else {
+                return await chrome.scripting.getRegisteredContentScripts(filter);
+            }
+        };
+
+        const existingRegistrations = await getContentScripts({
+            ids: [id]
+        });
+        if (existingRegistrations?.length > 0) {
+            await chrome.scripting.unregisterContentScripts({
+                ids: existingRegistrations.map((script) => script.id),
+            });
+        }
+    } else {
+        if (contentScriptRegistrations[id]) {
+            contentScriptRegistrations[id].unregister();
+            delete contentScriptRegistrations[id];
+        }
     }
 }
 
diff --git a/src/types.ts b/src/types.ts
index 9ea0841c..460a308d 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -99,8 +99,8 @@ export interface Registration {
     message: string;
     id: string;
     allFrames: boolean;
-    js: browser.extensionTypes.ExtensionFileOrCode[];
-    css: browser.extensionTypes.ExtensionFileOrCode[];
+    js: string[];
+    css: string[];
     matches: string[];
 }
 
diff --git a/src/utils.ts b/src/utils.ts
index bf6199d5..bc55f467 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -47,9 +47,14 @@ export default class Utils {
      * @param {CallableFunction} callback
      */
     setupExtraSitePermissions(callback: (granted: boolean) => void): void {
-        let permissions = ["webNavigation"];
-        if (!isSafari()) permissions.push("declarativeContent");
-        if (isFirefoxOrSafari() && !isSafari()) permissions = [];
+        const permissions = [];
+        if (!isFirefoxOrSafari()) {
+            permissions.push("declarativeContent");
+        }
+        if (!isFirefoxOrSafari() || isSafari()) {
+            permissions.push("webNavigation");
+        }
+        console.log(permissions)
 
         chrome.permissions.request({
             origins: this.getPermissionRegex(),
@@ -73,21 +78,12 @@ export default class Utils {
      * For now, it is just SB.config.invidiousInstances.
      */
     setupExtraSiteContentScripts(): void {
-        const firefoxJS = [];
-        for (const file of this.js) {
-            firefoxJS.push({file});
-        }
-        const firefoxCSS = [];
-        for (const file of this.css) {
-            firefoxCSS.push({file});
-        }
-
         const registration: Registration = {
             message: "registerContentScript",
             id: "invidious",
             allFrames: true,
-            js: firefoxJS,
-            css: firefoxCSS,
+            js: this.js,
+            css: this.css,
             matches: this.getPermissionRegex()
         };