make blocking rules queryable in pattern tree just as mappings are

3 files changed, 60 insertions, 26 deletions

diff --git a/test/unit/test_patterns_query_manager.py b/test/unit/test_patterns_query_manager.py
index ae1f490..35047f5 100644
--- a/test/unit/test_patterns_query_manager.py
+++ b/test/unit/test_patterns_query_manager.py
@@ -6,7 +6,7 @@ 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>
+# 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 CC0 1.0 Universal License as published by
@@ -64,8 +64,17 @@ sample_mappings = [simple_sample_mapping(pats, fruit) for pats, fruit in [
     ('https://***.gotmyowndoma.in/***', 'kiwi')
 ]]
 
+sample_blocking = [f'http{s}://{dw}gotmyown%sdoma.in{i}{pw}'
+                   for dw in ('', '***.', '**.', '*.')
+                   for i in ('/index.html', '')
+                   for pw in ('', '/', '/*')
+                   for s in ('', 's')]
+sample_blocking = [{'pattern': pattern % (i if i > 1 else ''),
+                    'allow': bool(i & 1)}
+                   for i, pattern in enumerate(sample_blocking)]
+
 # 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.
+# page, some tests can be as well performed running it from a normal page.
 @pytest.mark.get_page('https://gotmyowndoma.in')
 def test_pqm_tree_building(driver, execute_in_page):
     """
@@ -77,15 +86,19 @@ def test_pqm_tree_building(driver, execute_in_page):
     # Mock IndexedDB and build patterns tree.
     execute_in_page(
         '''
-        const initial_mappings = arguments[0]
-        let mappingchange;
-        function track_mock(cb)
-        {
+        const [initial_mappings, initial_blocking] = arguments.slice(0, 2);
+        let mappingchange, blockingchange;
+
+        haketilodb.track.mappings = function (cb) {
             mappingchange = cb;
 
             return [{}, initial_mappings];
         }
-        haketilodb.track.mappings = track_mock;
+        haketilodb.track.blocking = function (cb) {
+            blockingchange = cb;
+
+            return [{}, initial_blocking];
+        }
 
         let last_script;
         let unregister_called = 0;
@@ -99,7 +112,7 @@ def test_pqm_tree_building(driver, execute_in_page):
 
         returnval(start());
         ''',
-        sample_mappings[0:2])
+        sample_mappings[0:2], sample_blocking[0:2])
 
     found, tree, content_script, deregistrations = execute_in_page(
         '''
@@ -109,30 +122,44 @@ def test_pqm_tree_building(driver, execute_in_page):
         'https://gotmyowndoma.in/index.html')
     best_pattern = 'https://gotmyowndoma.in/index.html'
     assert found == \
-        dict([(f'inject-{fruit}', {'identifier': f'{fruit}-{best_pattern}'})
-              for fruit in ('banana', 'orange')])
+        dict([('~allow', 1),
+              *[(f'inject-{fruit}', {'identifier': f'{fruit}-{best_pattern}'})
+                for fruit in ('banana', 'orange')]])
     assert tree == extract_tree_data(content_script)
     assert deregistrations == 0
 
-    def condition_mappings_added(driver):
+    def condition_all_added(driver):
         last_script = execute_in_page('returnval(last_script);')
-        return all([m['identifier'] in last_script for m in sample_mappings])
+        nums = [i for i in range(len(sample_blocking)) if i > 1]
+        return (all([('gotmyown%sdoma' % i) in last_script for i in nums]) and
+                all([m['identifier'] in last_script for m in sample_mappings]))
 
     execute_in_page(
         '''
         for (const mapping of arguments[0])
             mappingchange({key: mapping.identifier, new_val: mapping});
+        for (const blocking of arguments[1])
+            blockingchange({key: blocking.pattern, new_val: blocking});
         ''',
-        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]
+        sample_mappings[2:], sample_blocking[2:])
+    WebDriverWait(driver, 10).until(condition_all_added)
+
+    odd_mappings = \
+        [m['identifier'] for i, m in enumerate(sample_mappings) if i & 1]
+    odd_blocking = \
+        [b['pattern'] for i, b in enumerate(sample_blocking) if i & 1]
+    even_mappings = \
+        [m['identifier'] for i, m in enumerate(sample_mappings) if 1 - i & 1]
+    even_blocking = \
+        [b['pattern'] for i, b in enumerate(sample_blocking) if 1 - i & 1]
 
     def condition_odd_removed(driver):
         last_script = execute_in_page('returnval(last_script);')
-        return (all([id not in last_script for id in odd]) and
-                all([id in last_script for id in even]))
+        nums = [i for i in range(len(sample_blocking)) if i > 1 and 1 - i & 1]
+        return (all([id not in last_script for id in odd_mappings]) and
+                all([id in last_script for id in even_mappings]) and
+                all([p not in last_script for p in odd_blocking[1:]]) and
+                all([('gotmyown%sdoma' % i) in last_script for i in nums]))
 
     def condition_all_removed(driver):
         content_script = execute_in_page('returnval(last_script);')
@@ -141,16 +168,18 @@ def test_pqm_tree_building(driver, execute_in_page):
     execute_in_page(
         '''
         arguments[0].forEach(identifier => mappingchange({key: identifier}));
+        arguments[1].forEach(pattern => blockingchange({key: pattern}));
         ''',
-        odd)
+        odd_mappings, odd_blocking)
 
     WebDriverWait(driver, 10).until(condition_odd_removed)
 
     execute_in_page(
         '''
         arguments[0].forEach(identifier => mappingchange({key: identifier}));
+        arguments[1].forEach(pattern => blockingchange({key: pattern}));
         ''',
-        even)
+        even_mappings, even_blocking)
 
     WebDriverWait(driver, 10).until(condition_all_removed)
 
diff --git a/test/unit/test_policy_deciding.py b/test/unit/test_policy_deciding.py
index a360537..88095af 100644
--- a/test/unit/test_policy_deciding.py
+++ b/test/unit/test_policy_deciding.py
@@ -61,11 +61,11 @@ def test_decide_policy(execute_in_page):
     policy = execute_in_page(
         '''{
         const tree = pqt.make();
-        pqt.register(tree, "http://kno.wn", "allowed", {allow: true});
+        pqt.register(tree, "http://kno.wn", "~allow", 1);
         returnval(decide_policy(tree, "http://kno.wn/", false, "abcd"));
         }''')
     assert policy['allow'] == True
-    assert policy['mapping'] == 'allowed'
+    assert policy['mapping'] == '~allow'
     for prop in ('payload', 'nonce', 'csp'):
         assert prop not in policy
 
@@ -87,11 +87,11 @@ def test_decide_policy(execute_in_page):
     policy = execute_in_page(
         '''{
         const tree = pqt.make();
-        pqt.register(tree, "http://kno.wn", "disallowed", {allow: false});
+        pqt.register(tree, "http://kno.wn", "~allow", 0);
         returnval(decide_policy(tree, "http://kno.wn/", true, "abcd"));
         }''')
     assert policy['allow'] == False
-    assert policy['mapping'] == 'disallowed'
+    assert policy['mapping'] == '~allow'
     for prop in ('payload', 'nonce'):
         assert prop not in policy
     assert parse_csp(policy['csp']) == {
diff --git a/test/unit/test_webrequest.py b/test/unit/test_webrequest.py
index 6af2758..ae617aa 100644
--- a/test/unit/test_webrequest.py
+++ b/test/unit/test_webrequest.py
@@ -29,8 +29,13 @@ def webrequest_js():
             ''';
             // Mock pattern tree.
             tree = pqt.make();
+
+            // Rule to block scripts.
             pqt.register(tree, "https://site.with.scripts.block.ed/***",
-                         "disallowed", {allow: false});
+                         "~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"});