fix: rewrite cloakserve CDP WebSocket URLs (#234)

* fix: rewrite cloakserve CDP WebSocket URLs

* fix: guard against blank forwarded host

---------

Co-authored-by: honor2030 <19909783+honor2030@users.noreply.github.com>

tests/test_cloakserve.py

@@ -3,6 +3,7 @@
 import asyncio
 import importlib.machinery
 import importlib.util
+import json
 import sys
 from pathlib import Path
 from types import SimpleNamespace
@@ -24,6 +25,8 @@ parse_connection_params = _mod.parse_connection_params
 parse_cli_args = _mod.parse_cli_args
 ChromePool = _mod.ChromePool
 _default_data_dir = _mod._default_data_dir
+_external_host = _mod._external_host
+_ws_scheme = _mod._ws_scheme
 SAFE_SEED_RE = _mod.SAFE_SEED_RE
 RESERVED_SEEDS = _mod.RESERVED_SEEDS
 
@@ -138,6 +141,141 @@ class TestParseCliArgs:
         assert result.endswith(".cloakbrowser/cloakserve")
 
 
+# ---------------------------------------------------------------------------
+# External host detection
+# ---------------------------------------------------------------------------
+
+
+class TestExternalHost:
+    """Test public host selection for rewritten CDP WebSocket URLs."""
+
+    class _Request:
+        def __init__(self, headers, port=9222, scheme="http", query_string=""):
+            self.headers = headers
+            self.app = {"port": port}
+            self.scheme = scheme
+            self.query_string = query_string
+
+    def test_forwarded_host_overrides_internal_host(self):
+        request = self._Request({
+            "Host": "localhost:8080",
+            "X-Forwarded-Host": "cdp.example.com:443",
+        })
+        assert _external_host(request) == "cdp.example.com:443"
+
+    def test_forwarded_host_uses_first_value(self):
+        request = self._Request({
+            "Host": "internal:9222",
+            "X-Forwarded-Host": "public.example.com, internal:9222",
+        })
+        assert _external_host(request) == "public.example.com"
+
+    def test_blank_forwarded_host_falls_back_to_host_header(self):
+        request = self._Request({
+            "Host": "internal:9222",
+            "X-Forwarded-Host": "   ",
+        })
+        assert _external_host(request) == "internal:9222"
+
+    def test_falls_back_to_host_header(self):
+        request = self._Request({"Host": "localhost:9222"})
+        assert _external_host(request) == "localhost:9222"
+
+    def test_falls_back_to_app_port_without_host_header(self):
+        request = self._Request({}, port=9333)
+        assert _external_host(request) == "localhost:9333"
+
+    def test_forwarded_proto_selects_wss(self):
+        request = self._Request({"X-Forwarded-Proto": "https"}, scheme="http")
+        assert _ws_scheme(request) == "wss"
+
+    def test_forwarded_proto_uses_first_value(self):
+        request = self._Request({"X-Forwarded-Proto": "https, http"}, scheme="http")
+        assert _ws_scheme(request) == "wss"
+
+
+class TestHandlerURLRewriting:
+    """Verify handlers rewrite CDP WebSocket URLs to the public cloakserve endpoint."""
+
+    class _Request:
+        def __init__(self, headers, query_string="fingerprint=seed1", port=9222, scheme="http"):
+            self.headers = headers
+            self.query_string = query_string
+            self.scheme = scheme
+            self.app = {"port": port, "pool": self._Pool()}
+
+        class _Pool:
+            async def get_or_launch(self, **_kwargs):
+                return SimpleNamespace(cdp_port=5100)
+
+    class _FakeResponse:
+        def __init__(self, data):
+            self._data = data
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *_exc):
+            return None
+
+        async def json(self):
+            return self._data
+
+    class _FakeSession:
+        def __init__(self, data):
+            self._data = data
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *_exc):
+            return None
+
+        def get(self, *_args, **_kwargs):
+            return TestHandlerURLRewriting._FakeResponse(self._data)
+
+    def _patch_session(self, monkeypatch, data):
+        monkeypatch.setattr(
+            _mod.aiohttp,
+            "ClientSession",
+            lambda *_args, **_kwargs: self._FakeSession(data),
+        )
+
+    def test_json_version_uses_forwarded_host_and_proto(self, monkeypatch):
+        self._patch_session(monkeypatch, {
+            "webSocketDebuggerUrl": "ws://127.0.0.1:5100/devtools/browser/browser-guid",
+        })
+        request = self._Request({
+            "Host": "internal:9222",
+            "X-Forwarded-Host": "cdp.example.com",
+            "X-Forwarded-Proto": "https",
+        })
+
+        response = asyncio.run(_mod.handle_json_version(request))
+        payload = json.loads(response.text)
+
+        assert payload["webSocketDebuggerUrl"] == (
+            "wss://cdp.example.com/fingerprint/seed1/devtools/browser/browser-guid"
+        )
+
+    def test_json_list_uses_forwarded_host_and_proto(self, monkeypatch):
+        self._patch_session(monkeypatch, [{
+            "webSocketDebuggerUrl": "ws://127.0.0.1:5100/devtools/page/page-guid",
+        }])
+        request = self._Request({
+            "Host": "internal:9222",
+            "X-Forwarded-Host": "cdp.example.com",
+            "X-Forwarded-Proto": "https",
+        })
+
+        response = asyncio.run(_mod.handle_json_list(request))
+        payload = json.loads(response.text)
+
+        assert payload[0]["webSocketDebuggerUrl"] == (
+            "wss://cdp.example.com/fingerprint/seed1/devtools/page/page-guid"
+        )
+
+
 # ---------------------------------------------------------------------------
 # URL rewriting logic (pure string manipulation, extracted from handlers)
 # ---------------------------------------------------------------------------