1 files changed, 109 insertions, 11 deletions
diff --git a/test/haketilo_test/unit/test_webrequest.py b/test/haketilo_test/unit/test_webrequest.py
index fb24b3d..1244117 100644
--- a/test/haketilo_test/unit/test_webrequest.py
+++ b/test/haketilo_test/unit/test_webrequest.py
@@ -24,6 +24,10 @@ import pytest
 from ..script_loader import load_script
 from .utils import are_scripts_allowed
 
+allowed_url = 'https://site.with.scripts.allow.ed/'
+blocked_url = 'https://site.with.scripts.block.ed/'
+payload_url = 'https://site.with.paylo.ad/'
+
 def webrequest_js():
     return (load_script('background/webrequest.js',
                         '#IMPORT common/patterns_query_tree.js AS pqt') +
@@ -34,24 +38,118 @@ def webrequest_js():
             default_allow = {name: "default_allow", value: true};
 
             // Rule to block scripts.
-            pqt.register(tree, "https://site.with.scripts.block.ed/***",
+            pqt.register(tree, "%(blocked)s***",
                          "~allow", 0);
 
             // Rule to allow scripts, but overridden by payload assignment.
-            pqt.register(tree, "https://site.with.paylo.ad/***", "~allow", 1);
-            pqt.register(tree, "https://site.with.paylo.ad/***",
-                         "somemapping", {identifier: "someresource"});
+            pqt.register(tree, "%(payload)s***", "~allow", 1);
+            pqt.register(tree, "%(payload)s***", "somemapping",
+                         {identifier: "someresource"});
 
             // Mock stream_filter.
             stream_filter.apply = (details, headers, policy) => headers;
+            ''' % {'blocked': blocked_url, 'payload': payload_url})
+
+def webrequest_js_start_called():
+    return webrequest_js() + ';\nstart("somesecret");'
+
+ext_url = 'moz-extension://49de6ce9-49fc-49e1-8102-7ef35286389c/html/settings.html'
+prefix = 'X-Haketilo-' + sha256(ext_url.encode()).digest().hex()
+
+# Prepare a list of headers as could be sent by a website.
+sample_csp_header = {
+    'name': 'Content-Security-Policy',
+    'value': "script-src 'self';"
+}
+sample_csp_header_idx = 7
+
+sample_headers = [
+    {'name': 'Content-Type',     'value': 'text/html;charset=utf-8'},
+    {'name': 'Content-Length',   'value': '61954'},
+    {'name': 'Content-Language', 'value': 'en'},
+    {'name': 'Expires',          'value': 'Mon, 12 Mar 2012 11:04...'},
+    {'name': 'Last-Modified',    'value': 'Fri, 26 Jul 2013 22:50...'},
+    {'name': 'Cache-Control',    'value': 'max-age=0, s-maxage=86...'},
+    {'name': 'Age',              'value': '224'},
+    {'name': 'Server',           'value': 'nginx/1.1.19'},
+    {'name': 'Date',             'value': 'Thu, 10 Mar 2022 12:09...'}
+]
+
+sample_headers.insert(sample_csp_header_idx, sample_csp_header)
+
+# Prepare a list of headers as would be crafted by Haketilo when there is a
+# payload to inject.
+nonce_source = f'somemapping:someresource:{payload_url}:somesecret'.encode()
+nonce = f'nonce-{sha256(nonce_source).digest().hex()}'
+
+payload_csp_header = {
+    'name': f'Content-Security-Policy',
+    'value': ("prefetch-src 'none'; script-src-attr 'none'; "
+              f"script-src '{nonce}'; script-src-elem '{nonce}';")
+}
+
+sample_payload_headers = [
+    *sample_headers,
+    {'name': prefix, 'value': ':)'},
+    payload_csp_header
+]
+
+sample_payload_headers[sample_csp_header_idx] = {
+    **sample_csp_header,
+    'name': f'{prefix}-{sample_csp_header["name"]}',
+}
+
+# Prepare a list of headers as would be crafted by Haketilo when scripts are
+# blocked.
+sample_blocked_headers = [*sample_payload_headers]
+sample_blocked_headers.pop()
+sample_blocked_headers.append(sample_csp_header)
+sample_blocked_headers.append({
+    'name': f'Content-Security-Policy',
+    'value': ("prefetch-src 'none'; script-src-attr 'none'; "
+              f"script-src 'none'; script-src-elem 'none';")
+})
+
+@pytest.mark.get_page('https://gotmyowndoma.in')
+@pytest.mark.parametrize('params', [
+    (sample_headers,         allowed_url),
+    (sample_blocked_headers, blocked_url),
+    (sample_payload_headers, payload_url),
+])
+def test_webrequest_on_headers_received(driver, execute_in_page, params):
+    """Unit-test the on_headers_received() function."""
+    headers_out, url = params
+
+    execute_in_page(
+        '''{
+        // Mock browser object.
+        const url = arguments[0];
+        this.browser = {runtime: {getURL: () => url}};
+        }''',
+        ext_url)
+
+    execute_in_page(webrequest_js())
+
+    execute_in_page('secret = "somesecret";')
+
+    for headers_in in [
+            sample_headers,
+            sample_blocked_headers,
+            sample_payload_headers
+    ]:
+        details = {'url': url, 'responseHeaders': headers_in, 'fromCache': True}
+        res = execute_in_page('returnval(on_headers_received(arguments[0]));',
+                              details)
 
-            // Mock secret and start webrequest operations.
-            start("somesecret");
-            ''')
+        assert res == {'responseHeaders': headers_out}
 
-@pytest.mark.ext_data({'background_script': webrequest_js})
+@pytest.mark.ext_data({'background_script': webrequest_js_start_called})
 @pytest.mark.usefixtures('webextension')
-def test_on_headers_received(driver, execute_in_page):
+def test_webrequest_real_pages(driver, execute_in_page):
+    """
+    Test webRequest-based header modifications by loading actual pages and
+    attempting to run scripts within them.
+    """
     for attempt in range(10):
         driver.get('https://site.with.scripts.block.ed/')
 
@@ -59,10 +157,10 @@ def test_on_headers_received(driver, execute_in_page):
             break
         assert attempt != 9
 
-    driver.get('https://site.with.scripts.allow.ed/')
+    driver.get(allowed_url)
     assert are_scripts_allowed(driver)
 
-    driver.get('https://site.with.paylo.ad/')
+    driver.get(payload_url)
     assert not are_scripts_allowed(driver)
     source = 'somemapping:someresource:https://site.with.paylo.ad/index.html:somesecret'
     assert are_scripts_allowed(driver, sha256(source.encode()).digest().hex())