code review: make onHeadersReceive() able to cancel responses

platform/firefox/vapi-background.js

@@ -1749,9 +1749,6 @@ var httpObserver = {
     ACCEPT: Components.results.NS_SUCCEEDED,
     // Request types:
     // https://developer.mozilla.org/en-US/docs/Mozilla/Tech/XPCOM/Reference/Interface/nsIContentPolicy#Constants
-    MAIN_FRAME: Ci.nsIContentPolicy.TYPE_DOCUMENT,
-    VALID_CSP_TARGETS: 1 << Ci.nsIContentPolicy.TYPE_DOCUMENT |
-                       1 << Ci.nsIContentPolicy.TYPE_SUBDOCUMENT,
     typeMap: {
         1: 'other',
         2: 'script',
@@ -1769,6 +1766,10 @@ var httpObserver = {
         19: 'beacon',
         21: 'image'
     },
+    onBeforeRequest: function(){},
+    onBeforeRequestTypes: null,
+    onHeadersReceived: function(){},
+    onHeadersReceivedTypes: null,
 
     get componentRegistrar() {
         return Components.manager.QueryInterface(Ci.nsIComponentRegistrar);
@@ -1935,14 +1936,12 @@ var httpObserver = {
     },
 
     handleRequest: function(channel, URI, details) {
-        var onBeforeRequest = vAPI.net.onBeforeRequest;
         var type = this.typeMap[details.rawtype] || 'other';
-
+        if ( this.onBeforeRequestTypes && this.onBeforeRequestTypes.has(type) === false ) {
-        if ( onBeforeRequest.types && onBeforeRequest.types.has(type) === false ) {
             return false;
         }
 
-        var result = onBeforeRequest.callback({
+        var result = this.onBeforeRequest({
             frameId: details.frameId,
             hostname: URI.asciiHost,
             parentFrameId: details.parentFrameId,
@@ -1963,65 +1962,85 @@ var httpObserver = {
         return false;
     },
 
+    getResponseHeader: function(channel, name) {
+        var value;
+        try {
+            value = channel.getResponseHeader(name);
+        } catch (ex) {
+        }
+        return value;
+    },
+
+    handleResponseHeaders: function(channel, URI, channelData) {
+        var type = this.typeMap[channelData[4]] || 'other';
+        if ( this.onHeadersReceivedTypes && this.onHeadersReceivedTypes.has(type) === false ) {
+            return;
+        }
+
+        // 'Content-Security-Policy' MUST come last in the array. Need to
+        // revised this eventually.
+        var responseHeaders = [];
+        var value = this.getResponseHeader(channel, 'Content-Security-Policy');
+        if ( value !== undefined ) {
+            responseHeaders.push({ name: 'Content-Security-Policy', value: value });
+        }
+
+        var result = this.onHeadersReceived({
+            hostname: URI.asciiHost,
+            parentFrameId: channelData[1],
+            responseHeaders: responseHeaders,
+            tabId: channelData[3],
+            type: this.typeMap[channelData[4]] || 'other',
+            url: URI.asciiSpec
+        });
+
+        if ( !result ) {
+            return;
+        }
+
+        if ( result.cancel ) {
+            channel.cancel(this.ABORT);
+            return;
+        }
+
+        if ( result.responseHeaders ) {
+            channel.setResponseHeader(
+                'Content-Security-Policy',
+                result.responseHeaders.pop().value,
+                true
+            );
+            return;
+        }
+    },
+
     observe: function(channel, topic) {
         if ( channel instanceof Ci.nsIHttpChannel === false ) {
             return;
         }
 
         var URI = channel.URI;
-        var channelData, result;
 
         if ( topic === 'http-on-examine-response' ) {
-            if ( !(channel instanceof Ci.nsIWritablePropertyBag) ) {
+            if ( channel instanceof Ci.nsIWritablePropertyBag === false ) {
                 return;
             }
 
+            var channelData;
             try {
                 channelData = channel.getProperty(this.REQDATAKEY);
             } catch (ex) {
-                return;
             }
-
             if ( !channelData ) {
                 return;
             }
 
-            if ( (1 << channelData[4] & this.VALID_CSP_TARGETS) === 0 ) {
+            this.handleResponseHeaders(channel, URI, channelData);
-                return;
-            }
-
-            topic = 'Content-Security-Policy';
-
-            try {
-                result = channel.getResponseHeader(topic);
-            } catch (ex) {
-                result = null;
-            }
-
-            result = vAPI.net.onHeadersReceived.callback({
-                hostname: URI.asciiHost,
-                parentFrameId: channelData[1],
-                responseHeaders: result ? [{name: topic, value: result}] : [],
-                tabId: channelData[3],
-                type: this.typeMap[channelData[4]] || 'other',
-                url: URI.asciiSpec
-            });
-
-            if ( result ) {
-                channel.setResponseHeader(
-                    topic,
-                    result.responseHeaders.pop().value,
-                    true
-                );
-            }
 
             return;
         }
 
         // http-on-opening-request
 
-        //console.log('http-on-opening-request:', URI.spec);
-
         var pendingRequest = this.lookupPendingRequest(URI.spec);
         var rawtype = channel.loadInfo && channel.loadInfo.contentPolicyType || 1;
 
@@ -2108,6 +2127,7 @@ var httpObserver = {
     }
 };
 
+/******************************************************************************/
 /******************************************************************************/
 
 vAPI.net = {};
@@ -2118,9 +2138,19 @@ vAPI.net.registerListeners = function() {
     // Since it's not used
     this.onBeforeSendHeaders = null;
 
-    this.onBeforeRequest.types = this.onBeforeRequest.types ?
+    if ( typeof this.onBeforeRequest.callback === 'function' ) {
-        new Set(this.onBeforeRequest.types) :
+        httpObserver.onBeforeRequest = this.onBeforeRequest.callback;
-        null;
+        httpObserver.onBeforeRequestTypes = this.onBeforeRequest.types ?
+            new Set(this.onBeforeRequest.types) :
+            null;
+    }
+
+    if ( typeof this.onHeadersReceived.callback === 'function' ) {
+        httpObserver.onHeadersReceived = this.onHeadersReceived.callback;
+        httpObserver.onHeadersReceivedTypes = this.onHeadersReceived.types ?
+            new Set(this.onHeadersReceived.types) :
+            null;
+    }
 
     var shouldBlockPopup = function(details) {
         var sourceTabId = null;

src/js/traffic.js

@@ -331,28 +331,33 @@ var onHeadersReceived = function(details) {
         return;
     }
 
-    // Special handling for root document.
+    var requestType = details.type;
-    if ( details.type === 'main_frame' ) {
+
+    if ( requestType === 'main_frame' ) {
         return onRootFrameHeadersReceived(details);
     }
 
-    // Just in case...
+    if ( requestType === 'sub_frame' ) {
-    if ( details.type !== 'sub_frame' ) {
+        return onFrameHeadersReceived(details);
-        return;
     }
+};
 
-    // If we reach this point, we are dealing with a sub_frame
+/******************************************************************************/
+
+var onRootFrameHeadersReceived = function(details) {
+    var µb = µBlock;
+    var tabId = details.tabId;
+
+    µb.tabContextManager.push(tabId, details.url);
 
     // Lookup the page store associated with this tab id.
-    var µb = µBlock;
     var pageStore = µb.pageStoreFromTabId(tabId);
     if ( !pageStore ) {
-        return;
+        pageStore = µb.bindTabToPageStats(tabId, 'beforeRequest');
     }
+    // I can't think of how pageStore could be null at this point.
 
-    // Frame id of frame request is their own id, while the request is made
+    var context = pageStore.createContextFromPage();
-    // in the context of the parent.
-    var context = pageStore.createContextFromFrameId(details.parentFrameId);
     context.requestURL = details.url;
     context.requestHostname = details.hostname;
     context.requestType = 'inline-script';
@@ -385,20 +390,19 @@ var onHeadersReceived = function(details) {
 
 /******************************************************************************/
 
-var onRootFrameHeadersReceived = function(details) {
+var onFrameHeadersReceived = function(details) {
-    var tabId = details.tabId;
     var µb = µBlock;
-
+    var tabId = details.tabId;
-    µb.tabContextManager.push(tabId, details.url);
 
     // Lookup the page store associated with this tab id.
     var pageStore = µb.pageStoreFromTabId(tabId);
     if ( !pageStore ) {
-        pageStore = µb.bindTabToPageStats(tabId, 'beforeRequest');
+        return;
     }
-    // I can't think of how pageStore could be null at this point.
 
-    var context = pageStore.createContextFromPage();
+    // Frame id of frame request is their own id, while the request is made
+    // in the context of the parent.
+    var context = pageStore.createContextFromFrameId(details.parentFrameId);
     context.requestURL = details.url;
     context.requestHostname = details.hostname;
     context.requestType = 'inline-script';