Fixed some resource loading feedback glitches.

2 files changed, 44 insertions, 41 deletions

diff --git a/src/bg/RequestGuard.js b/src/bg/RequestGuard.js
index 688bcc0..5220658 100644
--- a/src/bg/RequestGuard.js
+++ b/src/bg/RequestGuard.js
@@ -89,7 +89,7 @@ var RequestGuard = (() => {
         records = this.initTab(tabId);
       }
 
-      if (what === "noscriptFrame") {
+      if (what === "noscriptFrame" && type !== "object") {
         let nsf = records.noscriptFrames;
         nsf[frameId] = optValue;
         what = optValue ? "blocked" : "allowed";
@@ -364,12 +364,11 @@ var RequestGuard = (() => {
       return ALLOW;
     },
 
-    async onHeadersReceived(request) {
+    onHeadersReceived(request) {
       // called for main_frame, sub_frame and object
       debug("onHeadersReceived", request);
       let {url, documentUrl, statusCode, tabId, responseHeaders} = request;
-      //if (statusCode >= 300 && statusCode < 400) return;
-
+      
       try {
         let header, blocker;
         let content = {}
@@ -386,7 +385,7 @@ var RequestGuard = (() => {
         if (ns.isEnforced(tabId)) {
           let policy = ns.policy;
           let perms = policy.get(url, documentUrl).perms;
-          if (policy.autoAllowTop && request.frameId === 0 && perms === policy.DEFAULT) {
+          if (policy.autoAllowTop && request.type === "main_frame" && perms === policy.DEFAULT) {
             policy.set(Sites.optimalKey(url), perms = policy.TRUSTED.tempTwin);
           }
 
@@ -420,19 +419,20 @@ var RequestGuard = (() => {
 
           if (canScript) {
             if (!capabilities.has("webgl")) {
-              await RequestUtil.executeOnStart(request, {
+               RequestUtil.executeOnStart(request, {
                 file: "/content/webglHook.js"
               });
             }
             if (!capabilities.has("media")) {
-              await RequestUtil.executeOnStart(request, {
+              RequestUtil.executeOnStart(request, {
                 code: "window.mediaBlocker = true;"
               });
             }
-            await RequestUtil.executeOnStart(request, {
+            
+            RequestUtil.executeOnStart(request, {
               file: "content/media.js"
             });
-          } else if (request.frameId === 0 && !TabStatus.map.has(tabId)) {
+          } else if (request.type === "main_frame" && !TabStatus.map.has(tabId)) {
             debug("No TabStatus data yet for noscriptFrame", tabId);
             TabStatus.record(request, "noscriptFrame", true);
           }
@@ -450,7 +450,7 @@ var RequestGuard = (() => {
 
         if (header) return {responseHeaders};
       } catch (e) {
-        error(e, "Error in onHeadersReceived", uneval(request));
+        error(e, "Error in onHeadersReceived", request);
       }
       return ALLOW;
     },
diff --git a/src/bg/RequestUtil.js b/src/bg/RequestUtil.js
index 8a6e621..2999cc1 100644
--- a/src/bg/RequestUtil.js
+++ b/src/bg/RequestUtil.js
@@ -3,14 +3,16 @@
   let xmlFeedOrImage = /^(?:(?:application|text)\/(?:(?:r(?:ss|df)|atom)\+)xml(;|$))|image\//i;
   let rawXml = /^(?:application|text)\/xml;/i;
   let brokenXMLOnLoad;
-
-  let pendingRequests = new Map();
-
+  (async () => brokenXMLOnLoad = parseInt((await browser.runtime.getBrowserInfo()).version) < 61)()
+  
+  let pendingScripts = new Map();
+  let NOP = () => {};
+  
   let reloadingTabs = new Map();
   let tabKey = (tabId, url) => `${tabId}:${url}`;
 
   let cleanup = r => {
-    pendingRequests.delete(r.requestId);
+    pendingScripts.delete(r.requestId);
     let key = tabKey(r.tabId, r.url);
     if (reloadingTabs.get(key) === false) {
       reloadingTabs.delete(key);
@@ -19,11 +21,11 @@
   
   let executeAll = async request => {
     let {url, tabId, frameId, requestId} = request;
-    let scripts = pendingRequests.get(requestId);
+    let scripts = pendingScripts.get(requestId);
     if (!scripts) return -1;
-    pendingRequests.delete(requestId);
+    pendingScripts.delete(requestId);
     let count = 0;
-    for (let details of scripts.values()) {
+    await Promise.all([...scripts.values()].map(async details => {
       details = Object.assign({
         runAt: "document_start",
         matchAboutBlank: true,
@@ -43,7 +45,7 @@
       } catch (e) {
         error(e, "Execute on start failed", url, details);
       }
-    }
+    }));
     return count;
   };
   
@@ -53,11 +55,15 @@
       types:  ["main_frame", "sub_frame", "object"]
     };
     let wr = browser.webRequest;
-    for (let event of ["onCompleted", "onErrorOccurred"])
+    for (let event of ["onCompleted", "onErrorOccurred"]) {
       wr[event].addListener(cleanup, filter);
-  
+    }
+    
     filter.types = ["main_frame"];
-    wr.onResponseStarted.addListener(executeAll, filter);
+    wr.onResponseStarted.addListener(r => {
+      let scripts = pendingScripts.get(r.requestId);
+      if (scripts) scripts.runAndFlush();
+    }, filter);
   }
   
   var RequestUtil = {
@@ -66,7 +72,7 @@
       return request.response || (request.response = new ResponseMetaData(request));
     },
 
-    async executeOnStart(request, details) {
+    executeOnStart(request, details) {
       let {requestId, url, tabId, frameId, statusCode} = request;
 
       if (statusCode >= 300 && statusCode < 400) return;
@@ -89,10 +95,10 @@
       
       debug("Injecting script on start in %s (%o).", url, response);
 
-      let scripts = pendingRequests.get(requestId);
+      let scripts = pendingScripts.get(requestId);
       let scriptKey = JSON.stringify(details);
       if (!scripts) {
-        pendingRequests.set(requestId, scripts = new Map());
+        pendingScripts.set(requestId, scripts = new Map());
         scripts.set(scriptKey, details);
       } else {
         scripts.set(scriptKey, details);
@@ -105,18 +111,17 @@
         return;
       }
       
-      if (typeof brokenXMLOnLoad === "undefined") {
-        brokenXMLOnLoad = await (async () => parseInt((await browser.runtime.getBrowserInfo()).version) < 61)();
-      }
-      
       let mustCheckFeed = brokenXMLOnLoad && frameId === 0 && rawXml.test(contentType);
       debug("mustCheckFeed = %s, brokenXMLOnLoad = %s", mustCheckFeed, brokenXMLOnLoad);
       let filter = browser.webRequest.filterResponseData(requestId);
       let buffer = [];
-      let first = true;
-      let done = false;
+      let responseCompleted = false;
       let mustReload = false;
-      let runAndFlush = async () => {
+      scripts.runAndFlush = async () => {
+        scripts.runAndFlush = NOP;
+        if (responseCompleted && buffer && !buffer.length) {
+          filter.disconnect();
+        }
         let scriptsRan = await executeAll(request);
         if (mustCheckFeed && !scriptsRan) {
           mustReload = true;
@@ -124,23 +129,20 @@
           reloadingTabs.set(tabKey(tabId, url), true);
         }
         if (buffer && buffer.length) {
-          debug("Flushing %s buffer chunks", buffer.length);
+          debug("Flushing %s buffer chunks on %s", buffer.length, url);
           for (let chunk of buffer) {
             filter.write(chunk);
           }
-          filter.disconnect();
           buffer = null;
         }
-        if (done) {
+        filter.disconnect();
+        if (responseCompleted) {
           filter.onstop(null);
         }
       };
-
+      
       filter.ondata = event => {
-        if (first) {
-          runAndFlush();
-          first = false;
-        }
+        scripts.runAndFlush();
         if (buffer) {
           debug("buffering", url);
           buffer.push(event.data);
@@ -152,8 +154,9 @@
         filter.disconnect();
       };
 
-      filter.onstop = event => {
-        done = true;
+      filter.onstop = async event => {
+        responseCompleted = true;
+        await scripts.runAndFlush();
         if (mustReload && !buffer) {
           mustReload = false;
           browser.tabs.update(tabId, {url});