add actual payload injection functionality to new content script

2 files changed, 117 insertions, 17 deletions

diff --git a/content/content.js b/content/content.js
index 804a473..feef5db 100644
--- a/content/content.js
+++ b/content/content.js
@@ -48,16 +48,19 @@
 #FROM common/policy.js            IMPORT decide_policy
 #FROM content/policy_enforcing.js IMPORT enforce_blocking
 
-let already_run = false, page_info;
+let already_run = false, resolve_page_info,
+    page_info_prom = new Promise(cb => resolve_page_info = cb);
 
 function on_page_info_request([type], sender, respond_cb) {
     if (type !== "page_info")
 	return;
 
-    respond_cb(page_info);
+    page_info_prom.then(respond_cb);
+
+    return true;
 }
 
-globalThis.haketilo_content_script_main = function() {
+globalThis.haketilo_content_script_main = async function() {
     if (already_run)
         return;
 
@@ -73,10 +76,36 @@ globalThis.haketilo_content_script_main = function() {
 				 document.URL,
 				 globalThis.haketilo_defualt_allow,
 				 globalThis.haketilo_secret);
-    page_info = Object.assign({url: document.URL}, policy);
+    const page_info = Object.assign({url: document.URL}, policy);
     ["csp", "nonce"].forEach(prop => delete page_info[prop]);
 
-    enforce_blocking(policy);
+    if ("payload" in policy) {
+	const msg = ["indexeddb_files", policy.payload.identifier];
+	var scripts_prom = browser.runtime.sendMessage(msg);
+    }
+
+    await enforce_blocking(policy);
+
+    if ("payload" in policy) {
+	const script_response = await scripts_prom;
+
+	if ("error" in script_response) {
+	    resolve_page_info(Object.assign(page_info, script_response));
+	    return;
+	} else {
+	    for (const script_contents of script_response) {
+		const html_ns = "http://www.w3.org/1999/xhtml";
+		const script = document.createElementNS(html_ns, "script");
+
+		script.innerText = script_contents;
+		script.setAttribute("nonce", policy.nonce);
+		document.documentElement.append(script);
+		script.remove();
+	    }
+	}
+    }
+
+    resolve_page_info(page_info);
 }
 
 function main() {
diff --git a/test/unit/test_content.py b/test/unit/test_content.py
index c8e0987..35ab027 100644
--- a/test/unit/test_content.py
+++ b/test/unit/test_content.py
@@ -42,7 +42,7 @@ dynamic_script = \
 content_script = \
     '''
     /* Mock dynamic content script - case 'before'. */
-    if (/#dynamic_before$/.test(document.URL)) {
+    if (/dynamic_before/.test(document.URL)) {
         %s;
     }
 
@@ -50,6 +50,37 @@ content_script = \
     %s;
 
     /* Rest of mocks */
+
+    function mock_decide_policy() {
+        nonce = "12345";
+        return {
+            allow: false,
+            mapping: "what-is-programmers-favorite-drinking-place",
+            payload: {identifier: "foo-bar"},
+            nonce,
+            csp: "prefetch-src 'none'; script-src-attr 'none'; script-src 'nonce-12345'; script-src-elem 'nonce-12345';"
+        };
+    }
+
+    async function mock_payload_error([type, res_id]) {
+        if (type === "indexeddb_files")
+            return {error: {haketilo_error_type: "missing", id: res_id}};
+    }
+
+    async function mock_payload_ok([type, res_id]) {
+        if (type === "indexeddb_files")
+            return [1, 2].map(n => `window.haketilo_injected_${n} = ${n}${n};`);
+    }
+
+    if (/payload_error/.test(document.URL)) {
+        browser.runtime.sendMessage = mock_payload_error;
+        decide_policy = mock_decide_policy;
+    } else if (/payload_ok/.test(document.URL)) {
+        browser.runtime.sendMessage = mock_payload_ok;
+        decide_policy = mock_decide_policy;
+    }
+    /* Otherwise, script blocking policy without payload to inject is used. */
+
     const data_to_verify = {};
     function data_set(prop, val) {
         data_to_verify[prop] = val;
@@ -61,22 +92,24 @@ content_script = \
     enforce_blocking = policy => data_set("enforcing", policy);
 
     browser.runtime.onMessage.addListener = async function (listener_cb) {
-        await new Promise(cb => setTimeout(cb, 0));
+        await new Promise(cb => setTimeout(cb, 10));
 
         /* Mock a good request. */
         const set_good = val => data_set("good_request_result", val);
-        listener_cb(["page_info"], {}, val => set_good(val));
+        data_set("good_request_returned",
+                 !!listener_cb(["page_info"], {}, val => set_good(val)));
 
         /* Mock a bad request. */
         const set_bad = val => data_set("bad_request_result", val);
-        listener_cb(["???"], {}, val => set_bad(val));
+        data_set("bad_request_returned",
+                 !!listener_cb(["???"], {}, val => set_bad(val)));
     }
 
     /* main() call - normally present in content.js, inside '#IF !UNIT_TEST'. */
     main();
 
     /* Mock dynamic content script - case 'after'. */
-    if (/#dynamic_after$/.test(document.URL)) {
+    if (/#dynamic_after/.test(document.URL)) {
         %s;
     }
 
@@ -85,26 +118,64 @@ content_script = \
 
 @pytest.mark.ext_data({'content_script': content_script})
 @pytest.mark.usefixtures('webextension')
-@pytest.mark.parametrize('target', ['dynamic_before', 'dynamic_after'])
-def test_content_unprivileged_page(driver, execute_in_page, target):
+@pytest.mark.parametrize('target1', ['dynamic_before'])#, 'dynamic_after'])
+@pytest.mark.parametrize('target2', [
+    'scripts_blocked',
+    'payload_error',
+    'payload_ok'
+])
+def test_content_unprivileged_page(driver, execute_in_page, target1, target2):
     """
     Test functioning of content.js on an page using unprivileged schema (e.g.
     'https://' and not 'about:').
     """
-    driver.get(f'https://gotmyowndoma.in/index.html#{target}')
-    data = json.loads(driver.execute_script('return window.data_to_verify;'))
+    driver.get(f'https://gotmyowndoma.in/index.html#{target1}-{target2}')
+
+    def get_data(driver):
+        data = driver.execute_script('return window.data_to_verify;')
+        return data if 'good_request_result' in data else False
+
+    data = json.loads(WebDriverWait(driver, 10).until(get_data))
 
     assert 'gotmyowndoma.in' in data['good_request_result']['url']
     assert 'bad_request_result' not in data
 
+    assert data['good_request_returned'] == True
+    assert data['bad_request_returned'] == False
+
     assert data['cacher_started'] == True
 
-    assert data['enforcing']['allow']  == False
-    assert 'mapping' not in data['enforcing']
-    assert 'error'   not in data['enforcing']
+    for obj in (data['good_request_result'], data['enforcing']):
+        assert obj['allow'] == False
+
+    assert 'error' not in data['enforcing']
+
+    if target2.startswith('payload'):
+        for obj in (data['good_request_result'], data['enforcing']):
+            assert obj['payload']['identifier'] == 'foo-bar'
+            assert 'mapping' in obj
+    else:
+        assert 'payload' not in data['enforcing']
+        assert 'mapping' not in data['enforcing']
 
     assert data['script_run_without_errors'] == True
 
+    def vars_made_by_payload(driver):
+        vars_values = driver.execute_script(
+            'return [1, 2].map(n => window[`haketilo_injected_${n}`]);'
+        )
+        if vars_values != [None, None]:
+            return vars_values
+
+    if target2 == 'payload_error':
+        assert data['good_request_result']['error'] == {
+            'haketilo_error_type': 'missing',
+            'id': 'foo-bar'
+        }
+    elif target2 == 'payload_ok':
+        vars_values = WebDriverWait(driver, 10).until(vars_made_by_payload)
+        assert vars_values == [11, 22]
+
 @pytest.mark.ext_data({'content_script': content_script})
 @pytest.mark.usefixtures('webextension')
 @pytest.mark.parametrize('target', ['dynamic_before', 'dynamic_after'])