1 files changed, 67 insertions, 26 deletions

diff --git a/background/policy_injector.js b/background/policy_injector.js
index 9f79425..ee97333 100644
--- a/background/policy_injector.js
+++ b/background/policy_injector.js
@@ -2,6 +2,7 @@
  * Myext injecting policy to page using webRequest
  *
  * Copyright (C) 2021 Wojtek Kosior
+ * Copyright (C) 2021 jahoti
  * Redistribution terms are gathered in the `copyright' file.
  */
 
@@ -11,8 +12,13 @@
  * IMPORT get_storage
  * IMPORT browser
  * IMPORT is_chrome
- * IMPORT retrieve_nonce
+ * IMPORT is_mozilla
+ * IMPORT gen_unique
+ * IMPORT gen_nonce
+ * IMPORT is_privileged_url
  * IMPORT url_item
+ * IMPORT url_extract_target
+ * IMPORT sign_policy
  * IMPORT get_query_best
  * IMPORT csp_rule
  * IMPORTS_END
@@ -34,34 +40,60 @@ function is_csp_header(header)
     return !!csp_header_names[header.name.toLowerCase()];
 }
 
-function is_our_header(header, rule)
+function url_inject(details)
 {
-    return header.value === rule
+    if (is_privileged_url(details.url))
+	return;
+
+    const targets = url_extract_target(details.url);
+    if (targets.current)
+	return;
+
+    /* Redirect; update policy */
+    if (targets.policy)
+	targets.target = "";
+
+    let [pattern, settings] = query_best(targets.base_url);
+    /* Defaults */
+    if (!pattern)
+	settings = {};
+
+    const policy = encodeURIComponent(
+	JSON.stringify({
+	    allow: settings.allow,
+	    nonce: gen_nonce(),
+	    base_url: targets.base_url
+	})
+    );
+
+    return {
+	redirectUrl: [
+	    targets.base_url,
+	    '#', sign_policy(policy, new Date()), policy,
+	    targets.target,
+	    targets.target2
+	].join("")
+    };
 }
 
-function inject(details)
+function headers_inject(details)
 {
-    const url = url_item(details.url);
-
-    const [pattern, settings] = query_best(url);
-
-    const nonce = retrieve_nonce(details.tabId, details.frameId, true);
-    const rule = csp_rule(nonce);
-
-    var headers;
-
-    if (settings !== undefined && settings.allow) {
-	/*
-	 * Chrome doesn't have the buggy behavior of repeatedly injecting a
-	 * header we injected once. Firefox does and we have to remove it there.
-	 */
-	if (is_chrome)
-	    return {cancel: false};
-
-	headers = details.responseHeaders.filter(h => !is_our_header(h, rule));
-    } else {
-	headers = details.responseHeaders.filter(h => !is_csp_header(h));
-
+    const targets = url_extract_target(details.url);
+    /* Block mis-/unsigned requests */
+    if (!targets.current)
+	return {cancel: true};
+
+    const rule = csp_rule(targets.policy.nonce);
+    var headers = details.responseHeaders;
+
+    /*
+     * 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 (!targets.policy.allow) {
 	headers.push({
 	    name : header_name,
 	    value : rule
@@ -80,8 +112,17 @@ async function start_policy_injector()
     if (is_chrome)
 	extra_opts.push("extraHeaders");
 
+    browser.webRequest.onBeforeRequest.addListener(
+	url_inject,
+	{
+	    urls: ["<all_urls>"],
+	    types: ["main_frame", "sub_frame"]
+	},
+	["blocking"]
+    );
+
     browser.webRequest.onHeadersReceived.addListener(
-	inject,
+	headers_inject,
 	{
 	    urls: ["<all_urls>"],
 	    types: ["main_frame", "sub_frame"]