1 files changed, 79 insertions, 22 deletions

diff --git a/background/policy_injector.js b/background/policy_injector.js
index b3d85e8..9725e99 100644
--- a/background/policy_injector.js
+++ b/background/policy_injector.js
@@ -20,24 +20,28 @@
  * IMPORT url_extract_target
  * IMPORT sign_policy
  * IMPORT query_best
- * IMPORT csp_rule
+ * IMPORT sanitize_csp_header
  * IMPORTS_END
  */
 
 var storage;
 
-const csp_header_names = {
-    "content-security-policy" : true,
-    "x-webkit-csp" : true,
-    "x-content-security-policy" : true
-};
+const csp_header_names = new Set([
+    "content-security-policy",
+    "x-webkit-csp",
+    "x-content-security-policy"
+]);
 
-const header_name = "content-security-policy";
+/* TODO: variable no longer in use; remove if not needed */
+const unwanted_csp_directives = new Set([
+    "report-to",
+    "report-uri",
+    "script-src",
+    "script-src-elem",
+    "prefetch-src"
+]);
 
-function is_csp_header(header)
-{
-    return !!csp_header_names[header.name.toLowerCase()];
-}
+const report_only = "content-security-policy-report-only";
 
 function url_inject(details)
 {
@@ -82,20 +86,73 @@ function headers_inject(details)
     if (!targets.current)
 	return {cancel: true};
 
-    const rule = csp_rule(targets.policy.nonce);
-    var headers = details.responseHeaders;
+    let orig_csp_headers = is_chrome ? null : [];
+    let headers = [];
+    let csp_headers = is_chrome ? headers : [];
+
+    const rule = `'nonce-${targets.policy.nonce}'`;
+    const block = !targets.policy.allow;
+
+    for (const header of details.responseHeaders) {
+	if (!csp_header_names.has(header)) {
+	    /* Remove headers that only snitch on us */
+	    if (header.name.toLowerCase() === report_only && block)
+		continue;
+	    headers.push(header);
+
+	    /* If these are the original CSP headers, use them instead */
+	    /* Test based on url_extract_target() in misc.js */
+	    if (is_mozilla && header.name === "x-orig-csp") {
+		let index = header.value.indexOf('%5B');
+		if (index === -1)
+		    continue;
+
+		let sig = header.value.substring(0, index);
+		let data = header.value.substring(index);
+		if (sig !== sign_policy(data, 0))
+		    continue;
+
+		/* Confirmed- it's the originals, smuggled in! */
+		try {
+		    data = JSON.parse(decodeURIComponent(data));
+		} catch (e) {
+		    /* This should not be reached -
+			it's our self-produced valid JSON. */
+		    console.log("Unexpected internal error - invalid JSON smuggled!", e);
+		}
+
+		orig_csp_headers = csp_headers = null;
+		for (const header of data)
+		    headers.push(sanitize_csp_header(header, rule, block));
+	    }
+	} else if (is_chrome || !orig_csp_headers) {
+	    csp_headers.push(sanitize_csp_header(header, rule, block));
+	    if (is_mozilla)
+		orig_csp_headers.push(header);
+	}
+    }
 
-    /*
-     * Chrome doesn't have the buggy behavior of caching headers
-     * we injected. Firefox does and we have to remove it there.
-     */
-    if (!targets.policy.allow || is_mozilla)
-	headers = headers.filter(h => !is_csp_header(h));
+    if (orig_csp_headers) {
+	/** Smuggle in the original CSP headers for future use.
+	  * These are signed with a time of 0, as it's not clear there
+	  * is a limit on how long Firefox might retain these headers in
+	  * the cache.
+	  */
+	orig_csp_headers = encodeURIComponent(JSON.stringify(orig_csp_headers));
+	headers.push({
+	    name: "x-orig-csp",
+	    value: sign_policy(orig_csp_headers, 0) + orig_csp_headers
+	});
+
+	headers = headers.concat(csp_headers);
+    }
 
-    if (!targets.policy.allow) {
+    /* To ensure there is a CSP header if required */
+    if (block) {
 	headers.push({
-	    name : header_name,
-	    value : rule
+	    name: "content-security-policy",
+	    value: `script-src ${rule}; script-src-elem ${rule}; ` +
+		"script-src-attr 'none'; prefetch-src 'none';"
 	});
     }