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>

README.md

@@ -839,6 +839,26 @@ print(page.title())
 browser.close()
 ```
 
+If your framework needs a direct WebSocket endpoint, fetch Chrome's discovery document and use the rewritten `webSocketDebuggerUrl`. The URL points back through `cloakserve` so the CDP proxy can keep per-seed routing intact:
+
+```bash
+curl http://localhost:9222/json/version | jq -r .webSocketDebuggerUrl
+# ws://localhost:9222/devtools/browser/<browser-id>
+
+curl 'http://localhost:9222/json/version?fingerprint=11111' | jq -r .webSocketDebuggerUrl
+# ws://localhost:9222/fingerprint/11111/devtools/browser/<browser-id>
+```
+
+When `cloakserve` runs behind a reverse proxy or TLS terminator, forward the public host/protocol headers so generated WebSocket URLs use the address clients can actually reach:
+
+```nginx
+proxy_set_header Host $host;
+proxy_set_header X-Forwarded-Host $host;
+proxy_set_header X-Forwarded-Proto $scheme;
+```
+
+With those headers, `/json/version` returns public endpoints such as `wss://cdp.example.com/fingerprint/11111/devtools/browser/<browser-id>` instead of an internal container host.
+
 Pass extra flags to the browser:
 
 ```bash

bin/cloakserve

@@ -453,9 +453,21 @@ def parse_connection_params(query_string: str) -> dict:
 def _ws_scheme(request: web.Request) -> str:
     """Return 'wss' if client connected via HTTPS (e.g. TLS-terminating proxy), else 'ws'."""
     proto = request.headers.get("X-Forwarded-Proto", request.scheme)
+    proto = proto.split(",", 1)[0].strip().lower()
     return "wss" if proto == "https" else "ws"
 
 
+def _external_host(request: web.Request) -> str:
+    """Return the public host to use in rewritten CDP WebSocket URLs."""
+    fallback_host = request.headers.get("Host") or f"localhost:{request.app['port']}"
+    forwarded_host = request.headers.get("X-Forwarded-Host")
+    if forwarded_host:
+        public_host = forwarded_host.split(",", 1)[0].strip()
+        if public_host:
+            return public_host
+    return fallback_host
+
+
 async def handle_root(request: web.Request) -> web.Response:
     """Health check / process status."""
     pool: ChromePool = request.app["pool"]
@@ -504,7 +516,7 @@ async def handle_json_version(request: web.Request) -> web.Response:
         return web.json_response({"error": "CDP endpoint unreachable"}, status=502)
 
     # Rewrite webSocketDebuggerUrl to route through our multiplexer
-    host = request.headers.get("Host", f"localhost:{request.app['port']}")
+    host = _external_host(request)
     seed_key = params["seed"]
     if seed_key:
         ws_path = f"fingerprint/{seed_key}/devtools/browser"
@@ -545,7 +557,7 @@ async def handle_json_list(request: web.Request) -> web.Response:
         logger.error("Failed to reach Chrome CDP (port %d): %s", cp.cdp_port, exc)
         return web.json_response({"error": "CDP endpoint unreachable"}, status=502)
 
-    host = request.headers.get("Host", f"localhost:{request.app['port']}")
+    host = _external_host(request)
     scheme = _ws_scheme(request)
     seed_key = params["seed"]
 

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)
 # ---------------------------------------------------------------------------