1 files changed, 16 insertions, 14 deletions
diff --git a/background/policy_injector.js b/background/policy_injector.js
index 72318d4..e5af055 100644
--- a/background/policy_injector.js
+++ b/background/policy_injector.js
@@ -10,9 +10,8 @@
  * IMPORTS_START
  * IMPORT sign_data
  * IMPORT extract_signed
- * IMPORT sanitize_csp_header
- * IMPORT csp_rule
- * IMPORT is_csp_header_name
+ * IMPORT make_csp_rule
+ * IMPORT csp_header_regex
  * IMPORTS_END
  */
 
@@ -43,22 +42,25 @@ function inject_csp_headers(headers, policy)
 	break;
     }
 
+    if (policy.has_payload) {
+	csp_headers = [];
+	const non_csp_headers = [];
+	const header_list =
+	      h => csp_header_regex.test(h) ? csp_headers : non_csp_headers;
+	headers.forEach(h => header_list(h.name).push(h));
+	headers = non_csp_headers;
+    } else {
+	headers.push(...csp_headers || []);
+    }
+
     if (!hachette_header) {
 	hachette_header = {name: "x-hachette"};
 	headers.push(hachette_header);
     }
 
-    csp_headers = csp_headers ||
-	headers.filter(h => is_csp_header_name(h.name));
-
-    /* When blocking remove report-only CSP headers that snitch on us. */
-    headers = headers.filter(h => !is_csp_header_name(h.name, !policy.allow));
-
     if (old_signature)
 	headers = headers.filter(h => h.value.search(old_signature) === -1);
 
-    headers.push(...csp_headers.map(h => sanitize_csp_header(h, policy)));
-
     const policy_str = encodeURIComponent(JSON.stringify(policy));
     const signed_policy = sign_data(policy_str, new Date().getTime());
     const later_30sec = new Date(new Date().getTime() + 30000).toGMTString();
@@ -76,12 +78,12 @@ function inject_csp_headers(headers, policy)
     hachette_data = encodeURIComponent(JSON.stringify(hachette_data));
     hachette_header.value = sign_data(hachette_data, 0).join("_");
 
-    /* To ensure there is a CSP header if required */
-    if (!policy.allow)
+    if (!policy.allow) {
 	headers.push({
 	    name: "content-security-policy",
-	    value: csp_rule(policy.nonce)
+	    value: make_csp_rule(policy)
 	});
+    }
 
     return headers;
 }