make blocking rules queryable in pattern tree just as mappings are

6 files changed, 107 insertions, 46 deletions

diff --git a/background/patterns_query_manager.js b/background/patterns_query_manager.js
index e657448..8b563ef 100644
--- a/background/patterns_query_manager.js
+++ b/background/patterns_query_manager.js
@@ -4,7 +4,7 @@
  * Function: Instantiate the Pattern Tree data structure, filled with mappings
  *           from IndexedDB.
  *
- * Copyright (C) 2021 Wojtek Kosior
+ * Copyright (C) 2021,2022 Wojtek Kosior
  *
  * This program is free software: you can redistribute it and/or modify
  * it under the terms of the GNU General Public License as published by
@@ -54,7 +54,7 @@ let secret;
 const tree = pqt.make();
 #EXPORT tree
 
-const current_mappings = new Map();
+const currently_registered = new Map();
 
 #IF MOZILLA || MV3
 let registered_script = null;
@@ -92,28 +92,46 @@ if (this.haketilo_content_script_main)
 
     script_update_occuring = false;
 }
+#ENDIF
 
-function register_mapping(mapping)
+function register(kind, object)
 {
-    for (const [pattern, resource] of Object.entries(mapping.payloads))
-	pqt.register(tree, pattern, mapping.identifier, resource);
-    current_mappings.set(mapping.identifier, mapping);
-}
+    if (kind === "mappings") {
+	for (const [pattern, resource] of Object.entries(object.payloads))
+	    pqt.register(tree, pattern, object.identifier, resource);
+    } else /* if (kind === "blocking") */ {
+	/*
+	 * All simple block/allow rules use "~allow" in place of mapping id.
+	 * This way it won't collide with any real mapping id and will always
+	 * be sorted as higher value than mapping ids.
+	 */
+	pqt.register(tree, object.pattern, "~allow", object.allow + 0);
+    }
+
+#IF MOZILLA || MV3
+    const id = kind === "mappings" ? object.identifier : object.pattern;
+    currently_registered.set(id, object);
 #ENDIF
+}
 
-function mapping_changed(change)
+function changed(kind, change)
 {
-    console.log('mapping changes!', arguments);
-    const old_version = current_mappings.get(change.key);
+    const old_version = currently_registered.get(change.key);
     if (old_version !== undefined) {
-	for (const pattern in old_version.payloads)
-	    pqt.deregister(tree, pattern, change.key);
+	if (kind === "mappings") {
+	    for (const pattern in old_version.payloads)
+		pqt.deregister(tree, pattern, change.key);
+	} else /* if (kind === "blocking") */ {
+	    pqt.deregister(tree, change.key, "~allow");
+	}
 
-	current_mappings.delete(change.key);
+#IF MOZILLA || MV3
+	currently_registered.delete(change.key);
+#ENDIF
     }
 
     if (change.new_val !== undefined)
-	register_mapping(change.new_val);
+	register(kind, change.new_val);
 
 #IF MOZILLA || MV3
     script_update_needed = true;
@@ -125,10 +143,14 @@ async function start(secret_)
 {
     secret = secret_;
 
-    const [tracking, initial_mappings] =
-	  await haketilodb.track.mappings(mapping_changed);
+    const [mapping_tracking, initial_mappings] =
+	  await haketilodb.track.mappings(ch => changed("mappings", ch));
+    const [blocking_tracking, initial_blocking] =
+	  await haketilodb.track.blocking(ch => changed("blocking", ch));
+
+    initial_mappings.forEach(m => register("mappings", m));
+    initial_blocking.forEach(b => register("blocking", b));
 
-    initial_mappings.forEach(register_mapping);
 #IF MOZILLA || MV3
     script_update_needed = true;
     await update_content_script();
diff --git a/common/policy.js b/common/policy.js
index 0ac71d6..7ab9b5d 100644
--- a/common/policy.js
+++ b/common/policy.js
@@ -70,10 +70,15 @@ function decide_policy(patterns_tree, url, default_allow, secret)
     }
 
     if (payloads !== undefined) {
+	/*
+	 * mapping will be either the actual mapping identifier or "~allow" if
+	 * we matched a simple script block/allow rule.
+	 */
 	policy.mapping = Object.keys(payloads).sort()[0];
 	const payload = payloads[policy.mapping];
-	if (payload.allow !== undefined) {
-	    policy.allow = payload.allow;
+	if (policy.mapping === "~allow") {
+	    /* Convert 0/1 back to false/true. */
+	    policy.allow = !!payload;
 	} else /* if (payload.identifier) */ {
 	    policy.allow = false;
 	    policy.payload = payload;
diff --git a/html/settings.html b/html/settings.html
index 0bba5e3..df5e751 100644
--- a/html/settings.html
+++ b/html/settings.html
@@ -104,7 +104,7 @@
 	  /* Leave space for default policy dialog and headings. */
 	  --content-height: calc(var(--tab-content-height) - 3em);
       }
-      
+
       /* Pass height information to html in all tabs. */
       .tab {
 	  --content-height: var(--tab-content-height);
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"});