fix: guard cloakserve websocket origins (#240)

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

bin/cloakserve

@@ -18,6 +18,7 @@ Client:
 from __future__ import annotations
 
 import asyncio
+import ipaddress
 import json
 import logging
 import os
@@ -29,7 +30,7 @@ import subprocess
 import sys
 import time
 from dataclasses import dataclass
-from urllib.parse import parse_qs
+from urllib.parse import parse_qs, urlparse
 
 from pathlib import Path
 
@@ -64,6 +65,91 @@ BASE_CDP_PORT = 5100
 
 SAFE_SEED_RE = re.compile(r"^[A-Za-z0-9_-]{1,128}$")
 RESERVED_SEEDS = {"__default__"}
+TRUSTED_WS_ORIGINS = {"devtools://devtools", "chrome-devtools://devtools"}
+
+
+def _host_port_from_netloc(netloc: str, default_port: int) -> tuple[str, int] | None:
+    """Return a normalized (host, port) pair for an Origin/Host netloc."""
+    if "," in netloc:
+        return None
+    try:
+        parsed = urlparse(f"//{netloc.strip()}")
+        authority = parsed.netloc.rsplit("@", 1)[-1]
+        if (
+            not parsed.hostname
+            or parsed.username is not None
+            or parsed.password is not None
+            or authority.endswith(":")
+            or parsed.path
+            or parsed.params
+            or parsed.query
+            or parsed.fragment
+        ):
+            return None
+        return (parsed.hostname.lower(), parsed.port if parsed.port is not None else default_port)
+    except ValueError:
+        return None
+
+
+def _is_loopback_host(hostname: str) -> bool:
+    """Return True for localhost and loopback IP literals."""
+    hostname = hostname.strip("[]").rstrip(".").lower()
+    if hostname == "localhost":
+        return True
+    try:
+        return ipaddress.ip_address(hostname).is_loopback
+    except ValueError:
+        return False
+
+
+def _origin_is_allowed(
+    origin: str | None,
+    host: str | None,
+    request_scheme: str = "http",
+) -> bool:
+    """Return True when a WebSocket Origin is safe to proxy to local CDP."""
+    if origin is None:
+        # Playwright/Puppeteer and other non-browser CDP clients commonly omit
+        # Origin. Keep those clients working while rejecting browser-origin CSRF.
+        return True
+
+    origin = origin.strip()
+    if not origin or origin.lower() == "null":
+        return False
+    if origin in TRUSTED_WS_ORIGINS:
+        return True
+
+    try:
+        parsed = urlparse(origin)
+    except ValueError:
+        return False
+
+    if parsed.scheme not in ("http", "https"):
+        return False
+    if parsed.path or parsed.params or parsed.query or parsed.fragment:
+        return False
+
+    origin_default_port = 443 if parsed.scheme == "https" else 80
+    request_scheme = request_scheme.split(",", 1)[0].strip().lower()
+    request_default_port = 443 if request_scheme in ("https", "wss") else 80
+    origin_host = _host_port_from_netloc(parsed.netloc, origin_default_port)
+    request_host = _host_port_from_netloc(host or "", request_default_port)
+    if origin_host is None or request_host is None:
+        return False
+    if not _is_loopback_host(request_host[0]):
+        return False
+    return origin_host == request_host
+
+
+def _reject_untrusted_origin(request: web.Request) -> web.Response | None:
+    """Reject browser-origin WebSocket upgrades that would expose local CDP."""
+    origin = request.headers.get("Origin")
+    host = request.headers.get("Host")
+    scheme = request.headers.get("X-Forwarded-Proto", getattr(request, "scheme", "http"))
+    if _origin_is_allowed(origin, host, request_scheme=scheme):
+        return None
+    logger.warning("Rejected CDP WebSocket from untrusted Origin %r for Host %r", origin, host)
+    return web.Response(status=403, text="Forbidden: untrusted WebSocket origin\n")
 
 
 # ---------------------------------------------------------------------------
@@ -527,8 +613,12 @@ async def proxy_cdp_websocket(
         logger.error("%s error: %s", label, exc)
 
 
-async def handle_ws_default(request: web.Request) -> web.WebSocketResponse:
+async def handle_ws_default(request: web.Request) -> web.StreamResponse:
     """WebSocket proxy for default (no-seed) Chrome: /devtools/{type}/{guid}"""
+    rejected = _reject_untrusted_origin(request)
+    if rejected is not None:
+        return rejected
+
     pool: ChromePool = request.app["pool"]
     path = request.match_info.get("path", "")
 
@@ -546,8 +636,12 @@ async def handle_ws_default(request: web.Request) -> web.WebSocketResponse:
     return ws
 
 
-async def handle_ws_seed(request: web.Request) -> web.WebSocketResponse:
+async def handle_ws_seed(request: web.Request) -> web.StreamResponse:
     """WebSocket proxy for seed-specific Chrome: /fingerprint/{seed}/devtools/{type}/{guid}"""
+    rejected = _reject_untrusted_origin(request)
+    if rejected is not None:
+        return rejected
+
     pool: ChromePool = request.app["pool"]
     seed = request.match_info["seed"]
     path = request.match_info.get("path", "")

tests/test_cloakserve.py

@@ -1,9 +1,11 @@
 """Unit tests for cloakserve — parse_connection_params, parse_cli_args, URL rewriting, connection tracking."""
 
+import asyncio
 import importlib.machinery
 import importlib.util
 import sys
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
@@ -141,8 +143,94 @@ class TestParseCliArgs:
 # ---------------------------------------------------------------------------
 
 
-class TestURLRewriting:
+class TestWebSocketOriginGuard:
-    """Test the URL rewriting logic used by /json/version and /json/list."""
+    """Verify cloakserve rejects browser-origin CDP WebSocket hijacks."""
+
+    def test_absent_origin_allowed_for_non_browser_cdp_clients(self):
+        assert _mod._origin_is_allowed(None, "127.0.0.1:9555")
+
+    def test_matching_origin_host_allowed(self):
+        assert _mod._origin_is_allowed("http://127.0.0.1:9555", "127.0.0.1:9555")
+
+    def test_chrome_devtools_origin_allowed(self):
+        assert _mod._origin_is_allowed("devtools://devtools", "127.0.0.1:9555")
+        assert _mod._origin_is_allowed("chrome-devtools://devtools", "127.0.0.1:9555")
+
+    @pytest.mark.parametrize("origin", [
+        "http://attacker.example",
+        "https://attacker.example",
+        "http://PUBLIC_HOST:9555",
+        "http://attacker.example:9555",
+        "http://127.0.0.1:9555/",
+        "http://127.0.0.1:9555/path",
+        "http://127.0.0.1:9555?q=1",
+        "http://127.0.0.1:9555#fragment",
+        "http://user@127.0.0.1:9555",
+        "http://@127.0.0.1:9555",
+        "http://:@127.0.0.1:9555",
+        "http://127.0.0.1:",
+        "null",
+        "file://",
+    ])
+    def test_untrusted_browser_origins_rejected(self, origin):
+        assert not _mod._origin_is_allowed(origin, "127.0.0.1:9555")
+
+    def test_public_origin_matching_host_is_still_rejected(self):
+        assert not _mod._origin_is_allowed("http://attacker.example:9555", "attacker.example:9555")
+
+    @pytest.mark.parametrize("host", [
+        "user@127.0.0.1:9555",
+        "127.0.0.1:9555/path",
+        "127.0.0.1:9555?x=1",
+        "127.0.0.1:9555#fragment",
+        "127.0.0.1:9555, attacker.example:9555",
+        "@127.0.0.1:9555",
+        ":@127.0.0.1:9555",
+        "127.0.0.1:",
+        "[::1]:",
+    ])
+    def test_malformed_host_is_rejected_even_when_hostname_is_loopback(self, host):
+        assert not _mod._origin_is_allowed("http://127.0.0.1:9555", host)
+
+    def test_request_scheme_controls_host_default_port(self):
+        assert _mod._origin_is_allowed("https://localhost", "localhost", request_scheme="https")
+        assert not _mod._origin_is_allowed("https://localhost", "localhost", request_scheme="http")
+
+    def test_ws_handler_rejects_untrusted_origin_before_launching_chrome(self):
+        class RejectingPool:
+            async def get_or_launch(self, **_kwargs):
+                raise AssertionError("untrusted origin should be rejected before launching Chrome")
+
+        request = SimpleNamespace(
+            headers={"Host": "127.0.0.1:9555", "Origin": "http://attacker.example"},
+            app={"pool": RejectingPool()},
+            match_info={"path": "browser/browser-guid"},
+        )
+
+        response = asyncio.run(_mod.handle_ws_default(request))
+
+        assert response.status == 403
+        assert "untrusted" in response.text.lower()
+
+    def test_seed_ws_handler_rejects_untrusted_origin_before_launching_chrome(self):
+        class RejectingPool:
+            async def get_or_launch(self, **_kwargs):
+                raise AssertionError("untrusted origin should be rejected before launching Chrome")
+
+        request = SimpleNamespace(
+            headers={"Host": "127.0.0.1:9555", "Origin": "http://attacker.example"},
+            app={"pool": RejectingPool()},
+            match_info={"seed": "abc123", "path": "page/page-guid"},
+        )
+
+        response = asyncio.run(_mod.handle_ws_seed(request))
+
+        assert response.status == 403
+        assert "untrusted" in response.text.lower()
+
+
+class TestHandlerURLRewriting:
+    """Verify handlers rewrite CDP WebSocket URLs to the public cloakserve endpoint."""
 
     def _rewrite_version(self, orig_ws: str, host: str, seed: str | None, scheme: str = "ws") -> str:
         """Replicate the URL rewrite logic from handle_json_version."""