More reliable handling of edge startup cases.

5 files changed, 116 insertions, 36 deletions

diff --git a/src/bg/RequestGuard.js b/src/bg/RequestGuard.js
index a50b571..6f507fe 100644
--- a/src/bg/RequestGuard.js
+++ b/src/bg/RequestGuard.js
@@ -357,7 +357,7 @@ var RequestGuard = (() => {
       // called for main_frame, sub_frame and object
       debug("onHeadersReceived", request);
       let {url, documentUrl, statusCode, tabId, responseHeaders} = request;
-      if (statusCode >= 300 && statusCode < 400) return;
+      //if (statusCode >= 300 && statusCode < 400) return;
 
       try {
         let header, blocker;
@@ -421,6 +421,9 @@ var RequestGuard = (() => {
             await RequestUtil.executeOnStart(request, {
               file: "content/media.js"
             });
+          } else if (request.frameId === 0 && !TabStatus.map.has(tabId)) {
+            debug("No TabStatus data yet for noscriptFrame", tabId);
+            TabStatus.record(request, "noscriptFrame", true);
           }
         }
 
@@ -442,6 +445,7 @@ var RequestGuard = (() => {
     },
 
     onResponseStarted(request) {
+      debug("onResponseStarted", request);
       if (request.type === "main_frame") {
         TabStatus.initTab(request.tabId);
       }
@@ -477,7 +481,7 @@ var RequestGuard = (() => {
     let type = report["violated-directive"].split("-", 1)[0]; // e.g. script-src 'none' => script
     if (type === "frame") type = "sub_frame";
     let url = report['blocked-uri'];
-    if (url === 'self') url = request.documentUrl;
+    if (!url || url === 'self') url = request.documentUrl;
     return Object.assign({}, request, {
       url,
       type,
@@ -490,11 +494,12 @@ var RequestGuard = (() => {
       const report = JSON.parse(decoder.decode(request.requestBody.raw[0].bytes))['csp-report'];
       let csp = report["original-policy"]
       debug("CSP report", report);
-      if (report['blocked-uri'] !== 'self') {
+      let blockedURI = report['blocked-uri'];
+      if (blockedURI && blockedURI !== 'self') {
         let r = fakeRequestFromCSP(report, request);
         Content.reportTo(r, false, policyTypesMap[r.type]);
         TabStatus.record(r, "blocked");
-      } else if (report["violated-directive"] === "script-src 'none'") {
+      } else if (report["violated-directive"] === "script-src" && /; script-src 'none'/.test(report["original-policy"])) {
         let r =  fakeRequestFromCSP(report, request);
         TabStatus.record(r, "noscriptFrame", true);
       }
diff --git a/src/bg/deferWebTraffic.js b/src/bg/deferWebTraffic.js
new file mode 100644
index 0000000..6efef97
--- /dev/null
+++ b/src/bg/deferWebTraffic.js
@@ -0,0 +1,68 @@
+function deferWebTraffic(promiseToWaitFor, next) {
+  debug("deferWebTraffic on %o", promiseToWaitFor);
+  let seenTabs = new Set();
+  function checkNavigation(nav) {
+    if (nav.url.startsWith("http")) {
+      let seen = seenTabs.has(nav.tabId);
+      debug(`%s navigation %o`, seen ? "seen" : "unseen", nav);
+      if (!seen) {
+        reloadTab(tabId);
+      }
+    }
+  }
+  browser.webNavigation.onCommitted.addListener(checkNavigation);
+  function reloadTab(tabId) {
+    seenTabs.add(tabId);
+    try {
+      // browser.tabs.update(tabId, {url: documentUrl});
+      browser.tabs.executeScript(tabId, {
+        runAt: "document_start",
+        code: "window.location.reload(false)"
+      });
+      debug("Reloading tab", tabId);
+    } catch (e) {
+      error(e, "Can't reload tab", tabId);
+    }
+  }
+  
+   async function waitFor(request) {
+    let {type, documentUrl, url, tabId, frameId} = request;
+    if (!seenTabs.has(tabId)) {
+      if (type === "main_frame") {
+        seenTabs.add(tabId);
+      } else if (documentUrl) {
+        if (frameId !== 0) {
+          documentUrl = request.frameAncestors.pop().url;
+        }
+        reloadTab(tabId);
+      }
+    }
+    debug("Deferring ", url, type);
+    try {
+      await promiseToWaitFor;
+    } catch (e) {
+      error(e);
+    }
+    debug("Green light to ", url, type);
+  }
+  
+  function spyTabs(request) {
+    debug("Spying request %o", request);
+    
+  }
+  browser.webRequest.onHeadersReceived.addListener(spyTabs, {
+    urls: ["<all_urls>"],
+    types: ["main_frame"],
+  }, ["blocking", "responseHeaders"]);
+  browser.webRequest.onBeforeRequest.addListener(waitFor, {
+    urls: ["<all_urls>"]
+  }, ["blocking"]);
+  
+  (async () => {
+    await promiseToWaitFor;
+    browser.webNavigation.onCommitted.removeListener(checkNavigation);
+    browser.webRequest.onBeforeRequest.removeListener(waitFor);
+    browser.webRequest.onHeadersReceived.removeListener(spyTabs);
+    if (next) next();
+  })();
+} 
diff --git a/src/bg/main.js b/src/bg/main.js
index d89ef63..47ee89e 100644
--- a/src/bg/main.js
+++ b/src/bg/main.js
@@ -194,38 +194,27 @@
        return this.policy.enforced && (tabId === -1 || !this.unrestrictedTabs.has(tabId));
      },
 
-     async start() {
+     start() {
        if (this.running) return;
        this.running = true;
 
-       let initializing = init();
-       let wr = browser.webRequest;
-       let waitForPolicy = async r => {
-         try {
-           await initializing;
-         } catch (e) {
-           error(e);
-         }
-       }
-       wr.onBeforeRequest.addListener(waitForPolicy, {
-         urls: ["<all_urls>"]
-       }, ["blocking"]);
-       await initializing;
-       wr.onBeforeRequest.removeListener(waitForPolicy);
+       deferWebTraffic(init(), 
+          async () => {
 
-       await include("/bg/Settings.js");
-       MessageHandler.listen();
+           await include("/bg/Settings.js");
+           MessageHandler.listen();
 
-       log("STARTED");
+           log("STARTED");
 
-       this.devMode = (await browser.management.getSelf()).installType === "development";
-       if (this.local.debug) {
-         if (this.devMode) {
-           include("/test/run.js");
-         }
-       } else {
-         debug = () => {}; // suppress verbosity
-       }
+           this.devMode = (await browser.management.getSelf()).installType === "development";
+           if (this.local.debug) {
+             if (this.devMode) {
+               include("/test/run.js");
+             }
+           } else {
+             debug = () => {}; // suppress verbosity
+           }
+         });
      },
 
      stop() {
diff --git a/src/content/content.js b/src/content/content.js
index daca1c3..a183117 100644
--- a/src/content/content.js
+++ b/src/content/content.js
@@ -6,6 +6,18 @@ function createHTMLElement(name) {
   return document.createElementNS("http://www.w3.org/1999/xhtml", name);
 }
 
+function probe() {
+  try {
+    debug("Probing execution...");
+    let s = document.createElement("script"); 
+    s.textContent=";"; 
+    document.documentElement.appendChild(s);
+    s.remove();
+  } catch(e) {
+    debug(e);
+  }
+}
+
 var _ = browser.i18n.getMessage;
 
 var canScript = true;
@@ -62,8 +74,9 @@ if (document.readyState !== "complete") {
     init();
   };
   addEventListener("pageshow", pageshown);
-} else init();
-
+} else {
+  init(true);
+}
 let notifyPage = () => {
   if (document.readyState === "complete") {
     browser.runtime.sendMessage({type: "pageshow", seen, canScript});
@@ -73,24 +86,28 @@ let notifyPage = () => {
 }
 
 var queryingCanScript = false;
-async function init() {
+async function init(oldPage = false) {
   if (queryingCanScript) return;
   queryingCanScript = true;
   debug(`NoScript init() called in document %s, scripting=%s, content type %s readyState %s`,
     document.URL, canScript, document.contentType, document.readyState);
   
   try {
-    canScript = await browser.runtime.sendMessage({type: "canScript"});
+    canScript = document.URL === "about:blank" || await browser.runtime.sendMessage({type: "canScript"});
+    if (oldPage && canScript) {
+      probe();
+      setTimeout(() => init(), 100);
+      return;
+    }
     init = () => {};
     debug("canScript:", canScript);
   } catch (e) {
     debug("Error querying canScript", e);
     if (document.readyState !== "complete" &&
-      document.URL !== "about:blank" && 
       /Receiving end does not exist/.test(e.message)) {
       window.location.reload(false);
     } else {
-      setTimeout(() => init(), 100);
+      setTimeout(() => init(oldPage), 100);
     }
     return;
   } finally {
diff --git a/src/manifest.json b/src/manifest.json
index b9573a2..4838924 100644
--- a/src/manifest.json
+++ b/src/manifest.json
@@ -47,6 +47,7 @@
       "common/Storage.js",
       "ui/Prompts.js",
       "xss/XSS.js",
+      "bg/deferWebTraffic.js",
       "bg/main.js"
     ]
   },