1 files changed, 63 insertions, 13 deletions

diff --git a/background/policy_injector.js b/background/policy_injector.js
index 318190b..881595b 100644
--- a/background/policy_injector.js
+++ b/background/policy_injector.js
@@ -1,5 +1,7 @@
 /**
- * Hachette injecting policy to page using webRequest
+ * This file is part of Haketilo.
+ *
+ * Function: Injecting policy to page by modifying HTTP headers.
  *
  * Copyright (C) 2021 Wojtek Kosior
  * Copyright (C) 2021 jahoti
@@ -10,33 +12,81 @@
  * IMPORTS_START
  * IMPORT sign_data
  * IMPORT extract_signed
- * IMPORT sanitize_csp_header
  * IMPORT make_csp_rule
- * IMPORT is_csp_header_name
+ * IMPORT csp_header_regex
  * IMPORTS_END
  */
 
 function inject_csp_headers(headers, policy)
 {
-    if (!policy.allow || policy.has_payload) {
-	/* Remove report-only CSP headers that snitch on us. */
-	headers = headers.filter(h => !is_csp_header_name(h.name, true));
+    let csp_headers;
+    let old_signature;
+    let haketilo_header;
 
-	/* Add our own CSP header */
-	headers.push({
-	    name: "content-security-policy",
-	    value: make_csp_rule(policy)
-	});
+    for (const header of headers.filter(h => h.name === "x-haketilo")) {
+	/* x-haketilo header has format: <signature>_0_<data> */
+	const match = /^([^_]+)_(0_.*)$/.exec(header.value);
+	if (!match)
+	    continue;
+
+	const result = extract_signed(...match.slice(1, 3));
+	if (result.fail)
+	    continue;
+
+	/* This should succeed - it's our self-produced valid JSON. */
+	const old_data = JSON.parse(decodeURIComponent(result.data));
+
+	/* Confirmed- it's the originals, smuggled in! */
+	csp_headers = old_data.csp_headers;
+	old_signature = old_data.policy_sig;
+
+	haketilo_header = header;
+	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 (!haketilo_header) {
+	haketilo_header = {name: "x-haketilo"};
+	headers.push(haketilo_header);
+    }
+
+    if (old_signature)
+	headers = headers.filter(h => h.value.search(old_signature) === -1);
+
     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();
     headers.push({
 	name: "Set-Cookie",
-	value: `hachette-${signed_policy.join("=")}; Expires=${later_30sec};`
+	value: `haketilo-${signed_policy.join("=")}; Expires=${later_30sec};`
     });
 
+    /*
+     * Smuggle in the signature and 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 headers in the cache.
+     */
+    let haketilo_data = {csp_headers, policy_sig: signed_policy[0]};
+    haketilo_data = encodeURIComponent(JSON.stringify(haketilo_data));
+    haketilo_header.value = sign_data(haketilo_data, 0).join("_");
+
+    if (!policy.allow) {
+	headers.push({
+	    name: "content-security-policy",
+	    value: make_csp_rule(policy)
+	});
+    }
+
     return headers;
 }