1 files changed, 114 insertions, 22 deletions

diff --git a/background/webrequest.js b/background/webrequest.js
index a523772..5ec7b7f 100644
--- a/background/webrequest.js
+++ b/background/webrequest.js
@@ -3,7 +3,7 @@
  *
  * Function: Modify HTTP traffic usng webRequest API.
  *
- * Copyright (C) 2021 Wojtek Kosior <koszko@koszko.org>
+ * Copyright (C) 2021, 2022 Wojtek Kosior <koszko@koszko.org>
  *
  * This program is free software: you can redistribute it and/or modify
  * it under the terms of the GNU General Public License as published by
@@ -41,53 +41,145 @@
  * proprietary program, I am not going to enforce this in court.
  */
 
-#IMPORT common/indexeddb.js        AS haketilodb
+#IMPORT common/indexeddb.js AS haketilodb
+
 #IF MOZILLA
 #IMPORT background/stream_filter.js
 #ENDIF
 
 #FROM common/browser.js IMPORT browser
-#FROM common/misc.js    IMPORT is_privileged_url, csp_header_regex
+#FROM common/misc.js    IMPORT is_privileged_url, csp_header_regex, \
+                               sha256_async AS sha256
 #FROM common/policy.js  IMPORT decide_policy
 
 #FROM background/patterns_query_manager.js IMPORT tree, default_allow
 
 let secret;
 
-function on_headers_received(details)
-{
+#IF MOZILLA
+/*
+ * Under Mozilla-based browsers, responses are cached together with headers as
+ * they appear *after* modifications by Haketilo. This means Haketilo's CSP
+ * script-blocking headers might be present in responses loaded from cache. In
+ * the meantime the user might have changes Haketilo settings to instead allow
+ * the scripts on the page in question. This causes a problem and creates the
+ * need to somehow restore the response headers to the state in which they
+ * arrived from the server.
+ * To cope with this, Haketilo will inject some additional headers with private
+ * data. Those will include a hard-to-guess value derived from extension's
+ * internal ID. It is assumed the internal ID has a longer lifetime than cached
+ * responses.
+ */
+
+const settings_page_url = browser.runtime.getURL("html/settings.html");
+const header_prefix_prom = sha256(settings_page_url)
+      .then(hash => `X-Haketilo-${hash}`);
+
+/*
+ * Mozilla, unlike Chrome, allows webRequest callbacks to return promises. Here
+ * we leverage that to be able to use asynchronous sha256 computation.
+ */
+async function on_headers_received(details) {
+#IF NEVER
+} /* Help auto-indent in editors. */
+#ENDIF
+#ELSE
+function on_headers_received(details) {
+#ENDIF
     const url = details.url;
     if (is_privileged_url(details.url))
 	return;
 
     let headers = details.responseHeaders;
 
+#IF MOZILLA
+    const prefix = await header_prefix_prom;
+
+    /*
+     * We assume that the original CSP headers of a response are always
+     * preserved under names of the form:
+     *     X-Haketilo-<some_secret>-<original_name>
+     * In some cases the original response may contain no CSP headers. To still
+     * be able to tell whether the headers we were provided were modified by
+     * Haketilo in the past, all modifications are accompanied by addition of an
+     * extra header with name:
+     *     X-Haketilo-<some_secret>
+     */
+
+    const restore_old_headers = details.fromCache &&
+	  !!headers.filter(h => h.name === prefix).length;
+
+    if (restore_old_headers) {
+	const restored_headers = [];
+
+	for (const h of headers) {
+	    if (csp_header_regex.test(h.name) || h.name === prefix)
+		continue;
+
+	    if (h.name.startsWith(prefix)) {
+		restored_headers.push({
+		    name:  h.name.substring(prefix.length + 1),
+		    value: h.value
+		});
+	    } else {
+		restored_headers.push(h);
+	    }
+	}
+
+	headers = restored_headers;
+    }
+#ENDIF
+
     const policy =
 	  decide_policy(tree, details.url, !!default_allow.value, secret);
-    if (policy.allow)
-	return;
 
-    if (policy.payload)
-	headers = headers.filter(h => !csp_header_regex.test(h.name));
+    if (!policy.allow) {
+#IF MOZILLA
+	const to_append = [{name: prefix, value: ":)"}];
+
+	for (const h of headers.filter(h => csp_header_regex.test(h.name))) {
+	    if (!policy.payload)
+		to_append.push(Object.assign({}, h));
+
+	    h.name = `${prefix}-${h.name}`;
+	}
 
-    headers.push({name: "Content-Security-Policy", value: policy.csp});
+	headers.push(...to_append);
+#ELSE
+	if (policy.payload)
+	    headers = headers.filter(h => !csp_header_regex.test(h.name));
+#ENDIF
+
+	headers.push({name: "Content-Security-Policy", value: policy.csp});
+    }
 
 #IF MOZILLA
-    let skip = false;
-    for (const header of headers) {
-	if (header.name.toLowerCase().trim() !== "content-disposition")
-	    continue;
-
-	if (/^\s*attachment\s*(;.*)$/i.test(header.value)) {
-		skip = true;
-	} else {
-	    skip = false;
-	    break;
+    /*
+     * When page is meant to be viewed in the browser, use streamFilter to
+     * inject a dummy <script> at the very beginning of it. This <script>
+     * will cause extension's content scripts to run before page's first <meta>
+     * tag is rendered so that they can prevent CSP rules inside <meta> tags
+     * from blocking the payload we want to inject.
+     */
+
+    let use_stream_filter = !!policy.payload;
+    if (use_stream_filter) {
+	for (const header of headers) {
+	    if (header.name.toLowerCase().trim() !== "content-disposition")
+		continue;
+
+	    if (/^\s*attachment\s*(;.*)$/i.test(header.value)) {
+		use_stream_filter = false;
+	    } else {
+		use_stream_filter = true;
+		break;
+	    }
 	}
     }
-    skip = skip || (details.statusCode >= 300 && details.statusCode < 400);
+    use_stream_filter = use_stream_filter &&
+	(details.statusCode < 300 || details.statusCode >= 400);
 
-    if (!skip)
+    if (use_stream_filter)
 	headers = stream_filter.apply(details, headers, policy);
 #ENDIF