From b30cd8474d0a2338ff8683b80139cb61caf58056 Mon Sep 17 00:00:00 2001
From: Karen Rasmussen <rasmussenkaren@outlook.com>
Date: Tue, 27 Oct 2026 16:48:12 -0300
Subject: [PATCH] Serve ES modules with a wrapper the browser accepts

---
 tests/unit/viahtml/patch_test.py | 171 ++++++++++++++++++++++++++++++-
 viahtml/patch.py                 | 106 +++++++++++++++++++
 2 files changed, 275 insertions(+), 2 deletions(-)

diff --git a/tests/unit/viahtml/patch_test.py b/tests/unit/viahtml/patch_test.py
index a080124f..5000614d 100644
--- a/tests/unit/viahtml/patch_test.py
+++ b/tests/unit/viahtml/patch_test.py
@@ -1,12 +1,22 @@
 # pylint:disable=protected-access
 
-from unittest.mock import create_autospec, sentinel
+from unittest.mock import create_autospec, patch, sentinel
 
 import pytest
+from pywb.rewrite.content_rewriter import RewriteInfo
+from pywb.rewrite.default_rewriter import DefaultRewriter
+from pywb.rewrite.html_rewriter import HTMLRewriter
 from pywb.rewrite.url_rewriter import UrlRewriter
 
 from viahtml.hooks import Hooks
-from viahtml.patch import _PatchedHTMLRewriter
+from viahtml.patch import (
+    ESM_MOD,
+    ESM_TEXT_TYPE,
+    _ESMWombatProxyRewriter,
+    _patch_esm_rewriting,
+    _PatchedHTMLRewriter,
+    _resolve_text_type,
+)
 
 
 class TestRewriteLinkHref:
@@ -17,6 +27,13 @@ class TestRewriteLinkHref:
             # pywb's /id_/<URL> endpoint in order to disable any rewriting of
             # manifest files.
             ("https://example.com", [("rel", "manifest")], "id_/https://example.com"),
+            # Links that preload an ES module go through the module modifier, so
+            # that the module isn't served with pywb's classic script wrapper.
+            (
+                "https://example.com",
+                [("rel", "modulepreload")],
+                "esm_/https://example.com",
+            ),
             # Any other links get sent to the rw_mod endpoint which will rewrite
             # the file as usual.
             ("https://example.com", [], "rw_mod/https://example.com"),
@@ -87,6 +104,45 @@ def test_rewrite_tag_attrs_if_modify_tag_attrs_returns_stop_False(
         # It returns what pywb's _rewrite_tag_attrs returned.
         assert returned is True
 
+    def test_it_sends_module_scripts_to_the_module_modifier(
+        self, patched_html_rewriter
+    ):
+        self.returning_module_script(patched_html_rewriter)
+
+        patched_html_rewriter._rewrite_tag_attrs("script", sentinel.tag_attrs)
+
+        assert (
+            self.written(patched_html_rewriter)
+            == '<script type="module" src="esm_/https://example.com/app.js"'
+        )
+
+    def test_it_restores_the_script_modifier_afterwards(self, patched_html_rewriter):
+        self.returning_module_script(patched_html_rewriter)
+        original = patched_html_rewriter.rewrite_tags["script"]
+
+        patched_html_rewriter._rewrite_tag_attrs("script", sentinel.tag_attrs)
+
+        assert patched_html_rewriter.rewrite_tags["script"] == original
+
+    def test_it_restores_the_script_modifier_if_pywb_raises(
+        self, patched_html_rewriter
+    ):
+        self.returning_module_script(patched_html_rewriter)
+        original = patched_html_rewriter.rewrite_tags["script"]
+
+        with patch.object(HTMLRewriter, "_rewrite_tag_attrs", side_effect=RuntimeError):
+            with pytest.raises(RuntimeError):
+                patched_html_rewriter._rewrite_tag_attrs("script", sentinel.tag_attrs)
+
+        assert patched_html_rewriter.rewrite_tags["script"] == original
+
+    def returning_module_script(self, patched_html_rewriter):
+        """Make the hook return a `<script type="module">` with stop=False."""
+        patched_html_rewriter.hooks.modify_tag_attrs.return_value = (
+            [("type", "module"), ("src", "https://example.com/app.js")],
+            False,
+        )
+
     def written(self, patched_html_rewriter):
         """Return everything that was written to out.write() as one string."""
         return "".join(
@@ -94,6 +150,117 @@ def written(self, patched_html_rewriter):
         )
 
 
+class TestIsModuleScript:
+    @pytest.mark.parametrize(
+        "tag,tag_attrs,expected",
+        [
+            ("script", [("type", "module")], True),
+            # The type is matched case- and whitespace-insensitively.
+            ("script", [("type", " Module ")], True),
+            # Classic scripts, whether they declare a type or not.
+            ("script", [("type", "text/javascript")], False),
+            ("script", [("src", "app.js")], False),
+            ("script", [("type", None)], False),
+            # Only <script> tags can load a module this way.
+            ("link", [("type", "module")], False),
+        ],
+    )
+    def test_it(self, patched_html_rewriter, tag, tag_attrs, expected):
+        assert patched_html_rewriter._is_module_script(tag, tag_attrs) is expected
+
+
+class TestESMWombatProxyRewriter:
+    def test_it_declares_the_same_objects_as_pywb(self, esm_rewriter):
+        for obj in esm_rewriter.local_objs:
+            assert f'let {obj} = _____WB$wombat$assign$function_____("{obj}");' in (
+                esm_rewriter.first_buff
+            )
+
+    @pytest.mark.parametrize(
+        "forbidden,reason",
+        [
+            ("let arguments;", "strict mode, which modules always use, forbids it"),
+            ("{", "a block would enclose the module's export statement"),
+        ],
+    )
+    def test_it_omits_what_a_module_rejects(self, esm_rewriter, forbidden, reason):
+        declarations = esm_rewriter.first_buff.split("let window")[1]
+
+        assert forbidden not in declarations, reason
+
+    def test_it_appends_nothing(self, esm_rewriter):
+        # pywb appends '}}' to close the block it opened. There's no block here.
+        assert not esm_rewriter.last_buff
+
+    def test_it_reads_the_global_object_before_shadowing_it(self, esm_rewriter):
+        # The statements above the declarations share their scope, so reading
+        # `self` there would hit the temporal dead zone of `let self`.
+        preamble = esm_rewriter.first_buff.split("let window")[0]
+
+        assert "self." not in preamble
+        assert "globalThis." in preamble
+
+    @pytest.fixture
+    def esm_rewriter(self):
+        return _ESMWombatProxyRewriter(UrlRewriter("wburl"))
+
+
+class TestResolveTextType:
+    def test_it_reports_modules_as_such(self):
+        response = self.response_for(f"{ESM_MOD}/https://example.com/app.js")
+
+        assert _resolve_text_type(response, "js") == ESM_TEXT_TYPE
+
+    def test_it_leaves_everything_else_to_pywb(self):
+        response = self.response_for("js_/https://example.com/app.js")
+
+        with patch("viahtml.patch._PYWB_RESOLVE_TEXT_TYPE") as pywb_resolve_text_type:
+            returned = _resolve_text_type(response, sentinel.text_type)
+
+        pywb_resolve_text_type.assert_called_once_with(response, sentinel.text_type)
+        assert returned == pywb_resolve_text_type.return_value
+
+    def response_for(self, wburl):
+        response = create_autospec(RewriteInfo, instance=True)
+        response.url_rewriter = UrlRewriter(wburl)
+        return response
+
+
+@pytest.mark.usefixtures("unpatched_pywb")
+class TestPatchESMRewriting:
+    def test_it_registers_the_rewriter(self):
+        _patch_esm_rewriting()
+
+        assert (
+            DefaultRewriter.DEFAULT_REWRITERS[ESM_TEXT_TYPE] is _ESMWombatProxyRewriter
+        )
+
+    def test_it_replaces_pywbs_text_type_resolution(self):
+        _patch_esm_rewriting()
+
+        assert RewriteInfo._resolve_text_type is _resolve_text_type
+
+    def test_applying_it_twice_is_the_same_as_once(self):
+        _patch_esm_rewriting()
+        _patch_esm_rewriting()
+
+        # A patch which wrapped whatever it found would stack on each call.
+        assert RewriteInfo._resolve_text_type is _resolve_text_type
+
+    @pytest.fixture
+    def unpatched_pywb(self):
+        """Undo the patch for the duration of the test, and put it back after."""
+        resolve_text_type = RewriteInfo._resolve_text_type
+        rewriters = dict(DefaultRewriter.DEFAULT_REWRITERS)
+        DefaultRewriter.DEFAULT_REWRITERS.pop(ESM_TEXT_TYPE, None)
+
+        yield
+
+        RewriteInfo._resolve_text_type = resolve_text_type
+        DefaultRewriter.DEFAULT_REWRITERS.clear()
+        DefaultRewriter.DEFAULT_REWRITERS.update(rewriters)
+
+
 @pytest.fixture
 def patched_html_rewriter():
     patched_html_rewriter = _PatchedHTMLRewriter(UrlRewriter("wburl"))
diff --git a/viahtml/patch.py b/viahtml/patch.py
index 222adf47..d9d5a5ef 100644
--- a/viahtml/patch.py
+++ b/viahtml/patch.py
@@ -3,12 +3,23 @@
 from typing import Optional
 
 from pywb.apps.rewriterapp import RewriterApp
+from pywb.rewrite.content_rewriter import RewriteInfo
 from pywb.rewrite.default_rewriter import DefaultRewriter
 from pywb.rewrite.html_rewriter import HTMLRewriter
+from pywb.rewrite.regex_rewriters import JSWombatProxyRewriter
 from pywb.rewrite.url_rewriter import UrlRewriter
 
 from viahtml.hooks import Hooks
 
+# Our own URL modifier and text type for JS the page loads as an ES module.
+# `pywb` has neither: every script is a classic one as far as it's concerned.
+ESM_MOD = "esm_"
+ESM_TEXT_TYPE = "js-esm"
+
+# Stands in for `pywb`'s own `rewrite_tags["script"]`, which maps `src` to the
+# classic JS modifier whatever the script's type.
+ESM_SCRIPT_TAGS = {"src": ESM_MOD, "xlink:href": ESM_MOD}
+
 
 def apply_post_app_hooks(rewriter_app, hooks: Hooks):
     """Apply hooks after the app has been instantiated."""
@@ -19,6 +30,7 @@ def apply_pre_app_hooks(hooks: Hooks):
     """Apply hooks before the app has been instantiated."""
 
     _patch_url_rewriter(hooks)
+    _patch_esm_rewriting()
     _PatchedHTMLRewriter.patch(hooks)
 
 
@@ -29,6 +41,80 @@ def _patch_url_rewriter(hooks: Hooks):
     UrlRewriter.NO_REWRITE_URI_PREFIX = tuple(prefixes)
 
 
+class _ESMWombatProxyRewriter(JSWombatProxyRewriter):
+    """Serve an ES module with wombat's proxies set up.
+
+    So that wombat can proxy `window`, `location` and the rest, `pywb` wraps
+    every script it serves in a block which redeclares them. A module won't
+    parse with that wrapper and the browser discards the file without running
+    a line of it, so two of its parts go:
+
+    * `let arguments`. Strict mode forbids it, and modules are always strict.
+    * The block. A classic script's outermost scope is the global scope, so
+      `pywb` needs a block to declare into. A module's is already its own, and
+      a block would enclose the `export` which belongs at the top level.
+
+    The declarations themselves stay, and shadowing is why: it leaves telling
+    a variable from a property of the same name to the language. Substituting
+    the references instead, as `JSLinkAndLocationRewriter` does, rewrote
+    Kanopy's `document.body.style.top` along with the real thing.
+
+    `globalThis` in the statements above the declarations keeps them from
+    reading `let self` before it is initialised.
+    """
+
+    def __init__(self, rewriter, extra_rules=None):
+        super().__init__(rewriter, extra_rules=extra_rules)
+
+        self.first_buff = self._module_safe_first_buff()
+        self.last_buff = ""
+
+    def _module_safe_first_buff(self):
+        init_func = "_____WB$wombat$assign$function_____"
+
+        lines = [
+            f"var {init_func} = function(name) {{ return (globalThis._wb_wombat "
+            "&& globalThis._wb_wombat.local_init "
+            "&& globalThis._wb_wombat.local_init(name)) || globalThis[name]; };",
+            "if (!globalThis.__WB_pmw) { globalThis.__WB_pmw = function(obj) "
+            "{ this.__WB_source = obj; return this; } }",
+        ]
+        lines += [f'let {obj} = {init_func}("{obj}");' for obj in self.local_objs]
+
+        return "\n" + "\n".join(lines) + "\n"
+
+
+_PYWB_RESOLVE_TEXT_TYPE = (
+    RewriteInfo._resolve_text_type  # pylint:disable=protected-access
+)
+
+
+def _resolve_text_type(self, text_type):
+    """Report the text type of a response, recognising ES modules.
+
+    Replaces the `pywb` method of the same name. Everything which isn't a
+    module is left to the original.
+    """
+
+    if self.url_rewriter.wburl.mod == ESM_MOD:
+        return ESM_TEXT_TYPE
+
+    return _PYWB_RESOLVE_TEXT_TYPE(self, text_type)
+
+
+def _patch_esm_rewriting():
+    """Serve ES modules with a wrapper they can parse.
+
+    `_PatchedHTMLRewriter` tags module URLs with `ESM_MOD`; this teaches the
+    content rewriter what to do when the browser comes back asking for one.
+    """
+
+    DefaultRewriter.DEFAULT_REWRITERS[ESM_TEXT_TYPE] = _ESMWombatProxyRewriter
+    RewriteInfo._resolve_text_type = (  # pylint:disable=protected-access
+        _resolve_text_type
+    )
+
+
 class _PatchedHTMLRewriter(HTMLRewriter):
     hooks: Optional[Hooks] = None
 
@@ -49,6 +135,9 @@ def _rewrite_link_href(self, attr_value, tag_attrs, rw_mod):
             # 'id_' type hint appears to disable any rewriting
             return self._rewrite_url(attr_value, "id_")
 
+        if rel == "modulepreload":
+            return self._rewrite_url(attr_value, ESM_MOD)
+
         return super()._rewrite_link_href(attr_value, tag_attrs, rw_mod)
 
     def _rewrite_tag_attrs(self, tag, tag_attrs, set_parsing_context=True):
@@ -72,10 +161,27 @@ def _rewrite_tag_attrs(self, tag, tag_attrs, set_parsing_context=True):
 
             return True
 
+        if self._is_module_script(tag, new_attrs):
+            script_tags = self.rewrite_tags["script"]
+            self.rewrite_tags["script"] = ESM_SCRIPT_TAGS
+
+            try:
+                return super()._rewrite_tag_attrs(tag, new_attrs, set_parsing_context)
+            finally:
+                self.rewrite_tags["script"] = script_tags
+
         # If the hook returns stop=False we go on to all pywb's default
         # _rewrite_tag_attrs with the new attrs returned by the hook.
         return super()._rewrite_tag_attrs(tag, new_attrs, set_parsing_context)
 
+    def _is_module_script(self, tag, tag_attrs):
+        if tag != "script":
+            return False
+
+        type_value = self.get_attr(tag_attrs, "type")
+
+        return bool(type_value) and type_value.strip().lower() == "module"
+
 
 class _PatchedRewriterApp(RewriterApp):
     hooks: Optional[Hooks] = None