2 files changed, 65 insertions, 6 deletions

diff --git a/test/haketilo_test/unit/test_policy_enforcing.py b/test/haketilo_test/unit/test_policy_enforcing.py
index c5dd20e..98b5044 100644
--- a/test/haketilo_test/unit/test_policy_enforcing.py
+++ b/test/haketilo_test/unit/test_policy_enforcing.py
@@ -73,12 +73,15 @@ def get(driver, page, what_to_do):
 @pytest.mark.parametrize('csp_off_setting', [{}, {'csp_off': True}])
 def test_policy_enforcing_html(driver, execute_in_page, csp_off_setting):
     """
-    A test case of sanitizing <script>s and intrinsic javascript in pages.
+    A test case of sanitizing <script>s and intrinsic JavaScript in HTML pages.
     """
-    def assert_properly_blocked():
+    def click_all():
         for i in range(1, 3):
             driver.find_element_by_id(f'clickme{i}').click()
 
+    def assert_properly_blocked():
+        click_all()
+
         assert set(driver.execute_script('return window.__run || [];')) == set()
         assert bool(csp_off_setting) == are_scripts_allowed(driver)
 
@@ -98,8 +101,7 @@ def test_policy_enforcing_html(driver, execute_in_page, csp_off_setting):
         **csp_off_setting
     })
 
-    for i in range(1, 3):
-        driver.find_element_by_id(f'clickme{i}').click()
+    click_all()
 
     assert set(driver.execute_script('return window.__run || [];')) == \
         {'inline', 'on', 'href', 'src', 'data'}
@@ -121,3 +123,59 @@ def test_policy_enforcing_html(driver, execute_in_page, csp_off_setting):
 
     assert_properly_blocked()
     assert are_scripts_allowed(driver, nonce)
+
+# Test function analogous to that for HTML page.
+@pytest.mark.ext_data({'content_script': content_script})
+@pytest.mark.usefixtures('webextension')
+@pytest.mark.parametrize('csp_off_setting', [{}, {'csp_off': True}])
+def test_policy_enforcing_xml(driver, execute_in_page, csp_off_setting):
+    """
+    A test case of sanitizing <script>s and intrinsic JavaScript in XML
+    documents.
+    """
+    def click_all():
+        for name in ('idaret', 'nowamak', 'mango', 'annoying'):
+            elem = driver.find_element_by_id(f'{name}_circle')
+            try:
+                elem.click()
+            except:
+                pass
+
+    def assert_properly_blocked():
+        click_all()
+
+        try:
+            assert set(driver.execute_script('return window.__run || [];')) == set()
+        except:
+            from time import sleep
+            sleep(100000)
+        assert bool(csp_off_setting) == are_scripts_allowed(driver)
+
+    # First, see if scripts run when not blocked.
+    get(driver, 'https://gotmyowndoma.in/scripts_to_block_2.xml', {
+        'policy': allow_policy,
+        **csp_off_setting
+    })
+
+    click_all()
+
+    assert set(driver.execute_script('return window.__run || [];')) == \
+        {'grape', 'raspberry', 'idaret', 'melon'}
+    assert are_scripts_allowed(driver)
+
+    # Now, verify scripts don't run when blocked.
+    get(driver, 'https://gotmyowndoma.in/scripts_to_block_2.xml', {
+        'policy': block_policy,
+        **csp_off_setting
+    })
+
+    assert_properly_blocked()
+
+    # Now, verify only scripts with nonce can run when payload is injected.
+    get(driver, 'https://gotmyowndoma.in/scripts_to_block_2.xml', {
+        'policy': payload_policy,
+        **csp_off_setting
+    })
+
+    assert_properly_blocked()
+    assert are_scripts_allowed(driver, nonce)
diff --git a/test/haketilo_test/unit/utils.py b/test/haketilo_test/unit/utils.py
index b27a209..7ddf92a 100644
--- a/test/haketilo_test/unit/utils.py
+++ b/test/haketilo_test/unit/utils.py
@@ -228,11 +228,12 @@ def are_scripts_allowed(driver, nonce=None):
         return driver.execute_script(
             '''
             document.haketilo_scripts_allowed = false;
-            const script = document.createElement("script");
+            const html_ns = "http://www.w3.org/1999/xhtml";
+            const script = document.createElementNS(html_ns, "script");
             script.innerHTML = "document.haketilo_scripts_allowed = true;";
             if (arguments[0])
                 script.setAttribute("nonce", arguments[0]);
-            document.head.append(script);
+            (document.head || document.documentElement).append(script);
             return document.haketilo_scripts_allowed;
             ''',
             nonce)