Simplified and apparently more reliable+flexible+efficient dynamic script injection method.

2 files changed, 47 insertions, 116 deletions

diff --git a/src/bg/RequestGuard.js b/src/bg/RequestGuard.js
index 0f1558c..3304d02 100644
--- a/src/bg/RequestGuard.js
+++ b/src/bg/RequestGuard.js
@@ -422,7 +422,7 @@ var RequestGuard = (() => {
           }
         }
 
-        debug(`CSP blocker:`, blocker);
+        debug(`CSP blocker on %s:`, request.url, blocker);
         if (blocker) {
           if (header) {
             header.value = CSP.inject(header.value, blocker);
diff --git a/src/bg/RequestUtil.js b/src/bg/RequestUtil.js
index 1ebbbaa..2020286 100644
--- a/src/bg/RequestUtil.js
+++ b/src/bg/RequestUtil.js
@@ -1,130 +1,61 @@
 'use strict';
 {
-  let runningScripts = new Map();
-
+  let pendingRequests = new Map();
+  let cleanup = r => {
+    pendingRequests.delete(r.requestId);
+  };
+  let filter = {
+    urls: ["<all_urls>"],
+    types:  ["main_frame", "sub_frame", "object"]
+  };
+  browser.webRequest.onCompleted.addListener(cleanup, filter);
+  browser.webRequest.onErrorOccurred.addListener(cleanup, filter);
   var RequestUtil = {
     async executeOnStart(request, details) {
       let {requestId, tabId, frameId} = request;
-      details = Object.assign({
-        runAt: "document_start",
-        frameId,
-      }, details);
-      browser.tabs.executeScript(tabId, details);
-      return;
-      let filter = browser.webRequest.filterResponseData(requestId);
-      filter.onstart = event => {
-        browser.tabs.executeScript(tabId, details);
-        debug("Execute on start", details);
-        filter.write(new Uint8Array());
-      };
-      filter.ondata = event => {
-        filter.write(event.data);
-        filter.disconnect();
-
-      }
-    },
-    async executeOnStartCS(request, details) {
-      let {url, requestId, tabId, frameId} = request;
-
-      let urlObj = new URL(url);
-      if (urlObj.hash || urlObj.port || urlObj.username) {
-        urlObj.hash = urlObj.port = urlObj.username = "";
-        url = urlObj.toString();
+      let scripts = pendingRequests.get(requestId);
+      let scriptKey = JSON.stringify(details);
+      if (!scripts) {
+        pendingRequests.set(requestId, scripts = new Map());
+        scripts.set(scriptKey, details);
+      } else {
+        scripts.set(scriptKey, details);
+        return;
       }
-      let wr = browser.webRequest;
-      let filter = {
-        urls: [`${urlObj.origin}/*`],
-        types:  ["main_frame", "sub_frame", "object"]
-      };
-      let finalize;
-      let cleanup = r => {
-        if (cleanup && r.requestId === requestId) {
-          wr.onCompleted.removeListener(cleanup);
-          wr.onErrorOccurred.removeListener(cleanup);
-          cleanup = null;
-          if (finalize) {
-            finalize();
-          }
-        }
-      };
-      wr.onCompleted.addListener(cleanup, filter);
-      wr.onErrorOccurred.addListener(cleanup, filter);
-
-      details = Object.assign({
-        runAt: "document_start",
-        frameId,
-      }, details);
-
-      if (browser.contentScripts) {
-        let js = [{}];
-        if (details.file) js[0].file = details.file;
-        else if (details.code) js[0].code = details.code;
-        let settings = {
-          "runAt": details.runAt,
-          js,
-          matches: [url],
-          allFrames: frameId !== 0,
-        }
-        // let's try to avoid duplicates
-        let key = JSON.stringify(settings);
-        if (runningScripts.has(key)) {
-          let scriptRef = runningScripts.get(key);
-          scriptRef.count++;
-          return;
-        }
-        if (settings.allFrames) {
-          // let's check whether the same script is registered for top frames:
-          // if it is, let's unregister it first to avoid duplicates
-          settings.allFrames = false;
-          let topKey = JSON.stringify(settings);
-          settings.allFrames = true;
-          if (runningScripts.has(topKey)) {
-            let topScript = runningScripts.get(topKey);
-            try {
-              topScript.unregister();
-            } catch (e) {
-              error(e);
-            } finally {
-              runningScripts.delete(topKey);
-            }
-          }
-        }
 
-        let script = await browser.contentScripts.register(settings);
-        debug("Content script %o registered.", settings);
-        finalize = () => {
-          debug("Finalizing content script %o...", settings);
+      let filter = browser.webRequest.filterResponseData(requestId);
+      let buffer = [];
+      filter.onstart = async event => {
+        filter.write(new Uint8Array());
+        for (let details of scripts.values()) {
+          details = Object.assign({
+            runAt: "document_start",
+            frameId,
+          }, details);
           try {
-            script.unregister();
-            runningScripts.delete(key);
-            debug("Content script %o unregistered!", settings);
-          } finally {
-            finalize = null;
+            await browser.tabs.executeScript(tabId, details);
+            debug("Execute on start OK", request.url, details);
+          } catch (e) {
+            error(e, "Execute on start failed", request.url, details);
           }
         }
-        runningScripts.set(key, script);
-        if (!cleanup) { // the request has already been interrupted
-          finalize();
+        if (buffer.length) {
+          debug("Flushing %s buffer chunks", buffer.length);
+          for (let chunk of buffer) {
+            filter.write(chunk);
+          }
+          filter.disconnect();
+          buffer = null;
         }
-
-        return;
-      }
-
-      function listener(r) {
-        if (r.requestId === requestId) {
-          browser.tabs.executeScript(tabId, details);
-          finalize();
-          finalize = null;
+      };
+      filter.ondata = event => {
+        if (buffer) {
+          buffer.push(event.data);
+          return;
         }
+        filter.write(event.data);
+        filter.disconnect();
       }
-      finalize = () => {
-        wr.onResponseStarted.removeListener(listener);
-      }
-      wr.onResponseStarted.addListener(listener, filter);
-      debug("Executing %o", details);
-
-    },
-
-
+    }
   }
 }