facilitate egistering dynamic content scripts with mappings data

2 files changed, 236 insertions, 2 deletions

diff --git a/test/unit/test_indexeddb.py b/test/unit/test_indexeddb.py
index af60e1c..476690c 100644
--- a/test/unit/test_indexeddb.py
+++ b/test/unit/test_indexeddb.py
@@ -85,8 +85,6 @@ def test_haketilodb_save_remove(execute_in_page):
     # Mock some unwanted imports.
     execute_in_page(
         '''{
-        initial_data = {};
-
         const broadcast_mock = {};
         const nop = () => {};
         for (const key in broadcast)
diff --git a/test/unit/test_patterns_query_manager.py b/test/unit/test_patterns_query_manager.py
new file mode 100644
index 0000000..8ae7c28
--- /dev/null
+++ b/test/unit/test_patterns_query_manager.py
@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: CC0-1.0
+
+"""
+Haketilo unit tests - building pattern tree and putting it in a content script
+"""
+
+# This file is part of Haketilo
+#
+# Copyright (C) 2021, Wojtek Kosior <koszko@koszko.org>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the CC0 1.0 Universal License as published by
+# the Creative Commons Corporation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# CC0 1.0 Universal License for more details.
+
+import pytest
+import re
+import json
+from selenium.webdriver.support.ui import WebDriverWait
+
+from ..script_loader import load_script
+
+def simple_sample_mapping(patterns, fruit):
+    if type(patterns) is list:
+        payloads = dict([(p, {'identifier': fruit}) for p in patterns])
+    else:
+        payloads = {patterns: {'identifier': fruit}}
+    return {
+        'source_copyright': [],
+        'type': 'mapping',
+        'identifier': f'inject-{fruit}',
+        'payloads': payloads
+    }
+
+content_script_re = re.compile(r'this.haketilo_pattern_tree = (.*);')
+def extract_tree_data(content_script_text):
+    return json.loads(content_script_re.search(content_script_text)[1])
+
+# Fields that are not relevant for testing are omitted from these mapping
+# definitions.
+sample_mappings = [simple_sample_mapping(pats, fruit) for pats, fruit in [
+    (['https://gotmyowndoma.in/index.html',
+      'http://gotmyowndoma.in/index.html'], 'banana'),
+    (['https://***.gotmyowndoma.in/index.html',
+      'https://**.gotmyowndoma.in/index.html',
+      'https://*.gotmyowndoma.in/index.html',
+      'https://gotmyowndoma.in/index.html'], 'orange'),
+    ('https://gotmyowndoma.in/index.html/***', 'grape'),
+    ('http://gotmyowndoma.in/index.html/***', 'melon'),
+    ('https://gotmyowndoma.in/index.html', 'peach'),
+    ('https://gotmyowndoma.in/*', 'pear'),
+    ('https://gotmyowndoma.in/**', 'raspberry'),
+    ('https://gotmyowndoma.in/***', 'strawberry'),
+    ('https://***.gotmyowndoma.in/index.html', 'apple'),
+    ('https://***.gotmyowndoma.in/*', 'avocado'),
+    ('https://***.gotmyowndoma.in/**', 'papaya'),
+    ('https://***.gotmyowndoma.in/***', 'kiwi')
+]]
+
+# Even though patterns_query_manager.js is normally meant to run from background
+# page, tests can be as well performed running it from extension's bundled page.
+@pytest.mark.get_page('https://gotmyowndoma.in')
+def test_pqm_tree_building(driver, execute_in_page):
+    """
+    patterns_query_manager.js tracks Haketilo's internal database and builds a
+    constantly-updated pattern tree based on its contents. Mock the database and
+    verify tree building works properly.
+    """
+    execute_in_page(load_script('background/patterns_query_manager.js'))
+    # Mock IndexedDB and build patterns tree.
+    execute_in_page(
+        '''
+        const initial_mappings = arguments[0]
+        let mappingchange;
+        function track_mock(cb)
+        {
+            mappingchange = cb;
+
+            return [{}, initial_mappings];
+        }
+        haketilodb.track_mappings = track_mock;
+
+        let last_script;
+        let unregister_called = 0;
+        async function register_mock(injection)
+        {
+            await new Promise(resolve => setTimeout(resolve, 1));
+            last_script = injection.js[0].code;
+            return {unregister: () => unregister_called++};
+        }
+        browser = {contentScripts: {register: register_mock}};
+
+        returnval(start());
+        ''',
+        sample_mappings[0:2])
+
+    found, tree, content_script, deregistrations = execute_in_page(
+        '''
+        returnval([pqt.search(tree, arguments[0]).next().value,
+                   tree, last_script, unregister_called]);
+        ''',
+        'https://gotmyowndoma.in/index.html')
+    assert found == dict([(m['identifier'], m) for m in sample_mappings[0:2]])
+    assert tree == extract_tree_data(content_script)
+    assert deregistrations == 0
+
+    def condition_mappings_added(driver):
+        last_script = execute_in_page('returnval(last_script);')
+        return all([m['identifier'] in last_script for m in sample_mappings])
+
+    execute_in_page(
+        '''
+        for (const mapping of arguments[0]) {
+            mappingchange({
+                identifier: mapping.identifier,
+                new_val: mapping
+            });
+        }
+        ''',
+        sample_mappings[2:])
+    WebDriverWait(driver, 10).until(condition_mappings_added)
+
+    odd = [m['identifier'] for i, m in enumerate(sample_mappings) if i % 2]
+    even = [m['identifier'] for i, m in enumerate(sample_mappings) if 1 - i % 2]
+
+    def condition_odd_removed(driver):
+        last_script = execute_in_page('returnval(last_script);')
+        return all([id not in last_script for id in odd])
+
+    def condition_all_removed(driver):
+        content_script = execute_in_page('returnval(last_script);')
+        return extract_tree_data(content_script) == {}
+
+    execute_in_page(
+        '''
+        arguments[0].forEach(identifier => mappingchange({identifier}));
+        ''',
+        odd)
+
+    WebDriverWait(driver, 10).until(condition_odd_removed)
+
+    execute_in_page(
+        '''
+        arguments[0].forEach(identifier => mappingchange({identifier}));
+        ''',
+        even)
+
+    WebDriverWait(driver, 10).until(condition_all_removed)
+
+content_js = '''
+let already_run = false;
+this.haketilo_content_script_main = function() {
+    if (already_run)
+        return;
+    already_run = true;
+    document.documentElement.innerHTML = "<body><div id='tree-json'>";
+    document.getElementById("tree-json").innerText =
+        JSON.stringify(this.haketilo_pattern_tree);
+}
+if (this.haketilo_pattern_tree !== undefined)
+    this.haketilo_content_script_main();
+'''
+
+def background_js():
+    pqm_js = load_script('background/patterns_query_manager.js',
+                         "#IMPORT background/broadcast_broker.js")
+    return pqm_js + '; broadcast_broker.start(); start();'
+
+@pytest.mark.ext_data({
+    'content_script':    content_js,
+    'background_script': background_js
+})
+@pytest.mark.usefixtures('webextension')
+def test_pqm_script_injection(driver, execute_in_page):
+    # Let's open a normal page in a second window. Window 0 will be used to make
+    # changed to IndexedDB and window 1 to test the working of content scripts.
+    driver.execute_script('window.open("about:blank", "_blank");')
+    windows = [*driver.window_handles]
+    assert len(windows) == 2
+
+    def run_content_script():
+        driver.switch_to.window(windows[1])
+        driver.get('https://gotmyowndoma.in/index.html')
+        windows[1] = driver.window_handles[1]
+        return driver.execute_script(
+            '''
+            return (document.getElementById("tree-json") || {}).innerText;
+            ''')
+
+    for attempt in range(10):
+        json_txt = run_content_script()
+        if json.loads(json_txt) == {}:
+            break;
+        assert attempt != 9
+
+    driver.switch_to.window(windows[0])
+    execute_in_page(load_script('common/indexeddb.js'))
+
+    sample_data = {
+        'mappings': dict([(sm['identifier'], {'1.0': sm})
+                          for sm in sample_mappings]),
+        'resources': {},
+        'files': {}
+    }
+    execute_in_page('returnval(save_items(arguments[0]));', sample_data)
+
+    for attempt in range(10):
+        tree_json = run_content_script()
+        json.loads(tree_json)
+        if all([m['identifier'] in tree_json for m in sample_mappings]):
+            break
+        assert attempt != 9
+
+    driver.switch_to.window(windows[0])
+    execute_in_page(
+        '''{
+        const identifiers = arguments[0];
+        async function remove_items()
+        {
+            const ctx = await start_items_transaction(["mappings"], {});
+            for (const id of identifiers)
+                await remove_mapping(id, ctx);
+            await finalize_items_transaction(ctx);
+        }
+        returnval(remove_items());
+        }''',
+        [sm['identifier'] for sm in sample_mappings])
+
+    for attempt in range(10):
+        if json.loads(run_content_script()) == {}:
+            break
+        assert attempt != 9