Best effort to run webRequest.onHeaderReceived listener last (issue #6)

2 files changed, 74 insertions, 6 deletions

diff --git a/src/bg/RequestGuard.js b/src/bg/RequestGuard.js
index 1833c24..fcd6fa7 100644
--- a/src/bg/RequestGuard.js
+++ b/src/bg/RequestGuard.js
@@ -369,6 +369,14 @@ var RequestGuard = (() => {
 
     onHeadersReceived(request) {
       // called for main_frame, sub_frame and object
+      
+      // check for duplicate calls
+      let pending = pendingRequests.get(request.requestId);
+      if (pending && pending.headersProcessed) {
+        debug("[WARNING] already processed ", request);
+      }
+      pending.headersProcessed = true;
+      
       debug("onHeadersReceived", request);
       let {url, documentUrl, statusCode, tabId, responseHeaders} = request;
       
@@ -478,7 +486,12 @@ var RequestGuard = (() => {
       debug("%s scriptBlocked=%s setting noscriptFrame on ", request.url, scriptBlocked, request.tabId, request.frameId);
       TabStatus.record(request, "noscriptFrame", scriptBlocked);
       let pending = pendingRequests.get(request.requestId);
-      if (pending) pending.scriptBlocked = scriptBlocked;
+      if (pending) {
+        pending.scriptBlocked = scriptBlocked;
+        if (!pending.headersProcessed) {
+          debug("[WARNING] onHeadersReceived could not process", request);
+        }
+      }
     },
 
     onCompleted(request) {
@@ -533,12 +546,13 @@ var RequestGuard = (() => {
     }
     return ABORT;
   }
-
+  
   const RequestGuard = {
     async start() {
       let wr = browser.webRequest;
       let listen = (what, ...args) => wr[what].addListener(listeners[what], ...args);
-
+      let listenLast = (what, ...args) => new LastListener(wr[what], listeners[what], ...args).install();
+      
       let allUrls = ["<all_urls>"];
       let docTypes = ["main_frame", "sub_frame", "object"];
 
@@ -546,7 +560,7 @@ var RequestGuard = (() => {
         {urls: allUrls, types: allTypes},
         ["blocking"]
       );
-      listen("onHeadersReceived",
+      listenLast("onHeadersReceived",
         {urls: allUrls, types: docTypes},
         ["blocking", "responseHeaders"]
       );
@@ -570,8 +584,12 @@ var RequestGuard = (() => {
 
     stop() {
       let wr = browser.webRequest;
-      for (let [name, listener] of Object.entries(this.listeners)) {
-        wr[name].removeListener(listener);
+      for (let [name, listener] of Object.entries(listeners)) {
+        if (typeof listener === "function") {
+          wr[name].removeListener(listener);
+        } else if (listener instanceof LastListener) {
+          listener.uninstall();
+        }
       }
       wr.onBeforeRequest.removeListener(onViolationReport);
     }
diff --git a/src/lib/LastListener.js b/src/lib/LastListener.js
new file mode 100644
index 0000000..51600eb
--- /dev/null
+++ b/src/lib/LastListener.js
@@ -0,0 +1,50 @@
+/**
+* Wrapper around listeners on various WebExtensions
+* APIs (e.g. webRequest.on*), as a best effort to 
+* let them run last by removing and re-adding them 
+* on each call (swapping 2 copies because 
+* addListener() calls are asynchronous).
+* Note: we rely on implementation details Like
+* listeners being called in addition order; also,
+* clients should ensure they're not called twice for
+* the same event, if that's important.
+}
+*/
+
+class LastListener {
+  constructor(observed, listener, ...extras) {
+    this.observed = observed;
+    this.listener = listener;
+    this.extras = extras;
+    let ww = this._wrapped = [listener, listener].map(l => {
+      let w = (...args) => {
+        if (this.observed.hasListener(w._other)) {
+          this.observed.removeListener(w._other);
+          if (this.last === w) return this.defaultResult;
+        } else if (this.installed) { 
+          this.observed.addListener(w._other, ...this.extras);
+          this.last = w._other;
+        }
+        return this.installed ? this.listener(...args)
+          : this.defaultResult;
+      }
+      return w;
+    });
+    
+    ww[0]._other = ww[1];
+    ww[1]._other = ww[0];
+    this.installed = false;
+    this.defaultResult = null;
+  }
+
+  install() {
+    if (this.installed) return;
+    this.observed.addListener(this._wrapped[0], ...this.extras);
+    this.installed = true;
+  }
+  
+  uninstall() {
+    this.installed = false;
+    for (let l of this._wrapped) this.observed.removeListener(l);
+  }
+}