From d63ddd4d435a779996a14eeaaf93fcda73779a7e Mon Sep 17 00:00:00 2001 From: CloakHQ Date: Sun, 5 Apr 2026 22:41:33 +0200 Subject: [PATCH] feat(cloakserve): add connection tracking, configurable data dir, better status endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Move `import websockets` to top-level (guaranteed by [serve] extra) - Add --data-dir flag with smart default (Docker → /tmp/cloakserve, bare metal → ~/.cloakbrowser/cloakserve) - Store launch params (tz/locale/proxy) on ChromeProcess for conflict logging - Enhance GET / to return per-process detail (pid, port, seed, connections, config) - Add connection refcounting in WS handlers for status visibility - Add first-launch-wins note to README - Add tests for data-dir, Docker detection, and connection tracking --- README.md | 2 +- bin/cloakserve | 84 ++++++++++++++++++++++++++++++++++------ tests/test_cloakserve.py | 75 ++++++++++++++++++++++++++++++++++- 3 files changed, 147 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 969ba3f..7a7a319 100644 --- a/README.md +++ b/README.md @@ -775,7 +775,7 @@ b4 = pw.chromium.connect_over_cdp( ) ``` -Supported query params: `fingerprint`, `timezone`, `locale`, `platform`, `platform-version`, `brand`, `brand-version`, `gpu-vendor`, `gpu-renderer`, `hardware-concurrency`, `device-memory`, `screen-width`, `screen-height`, `proxy`, `geoip`. Same seed reuses the same process. No seed = shared default process (backward compatible). +Supported query params: `fingerprint`, `timezone`, `locale`, `platform`, `platform-version`, `brand`, `brand-version`, `gpu-vendor`, `gpu-renderer`, `hardware-concurrency`, `device-memory`, `screen-width`, `screen-height`, `proxy`, `geoip`. Same seed reuses the same process (first connection's params win). No seed = shared default process (backward compatible). Check active processes at `GET /` (returns JSON with PIDs, ports, and connection counts). **Persistent profiles** — mount a volume to keep cookies and sessions across container restarts: diff --git a/bin/cloakserve b/bin/cloakserve index 69f4380..dcb1ab0 100755 --- a/bin/cloakserve +++ b/bin/cloakserve @@ -30,7 +30,10 @@ import time from dataclasses import dataclass from urllib.parse import parse_qs +from pathlib import Path + import aiohttp +import websockets from aiohttp import web from cloakbrowser.browser import build_args, maybe_resolve_geoip @@ -69,6 +72,9 @@ class ChromeProcess: process: subprocess.Popen cdp_port: int user_data_dir: str + timezone: str | None = None + locale: str | None = None + proxy: str | None = None # --------------------------------------------------------------------------- @@ -81,14 +87,18 @@ class ChromePool: binary: str, global_args: list[str], headless: bool, + data_dir: str = "/tmp/cloakserve", ): self._binary = binary self._global_args = global_args self._headless = headless + self._data_dir = data_dir self._processes: dict[str, ChromeProcess] = {} self._default: ChromeProcess | None = None self._locks: dict[str, asyncio.Lock] = {} self._next_port = BASE_CDP_PORT + # Connection refcounting for status reporting + self._connections: dict[str, int] = {} def _get_lock(self, seed: str) -> asyncio.Lock: if seed not in self._locks: @@ -108,6 +118,18 @@ class ChromePool: continue raise RuntimeError("No free ports available for Chrome CDP") + def connect(self, seed_key: str) -> None: + """Increment connection refcount for a seed.""" + self._connections[seed_key] = self._connections.get(seed_key, 0) + 1 + + def disconnect(self, seed_key: str) -> None: + """Decrement connection refcount for a seed.""" + count = self._connections.get(seed_key, 0) - 1 + if count <= 0: + self._connections.pop(seed_key, None) + else: + self._connections[seed_key] = count + async def get_or_launch( self, seed: str | None, @@ -134,9 +156,10 @@ class ChromePool: if proc.process.poll() is None: if any([extra_args, timezone, locale, proxy, geoip]): logger.warning( - "Seed %s already running (port %d) — " + "Seed %s already running (port %d, tz=%s, locale=%s, proxy=%s) — " "ignoring new params (first-launch wins)", seed_key, proc.cdp_port, + proc.timezone, proc.locale, proc.proxy, ) return proc # Dead — clean up @@ -163,7 +186,7 @@ class ChromePool: # Allocate port and user data dir port = self._allocate_port() - user_data_dir = f"/tmp/cloakserve-{seed_key}" + user_data_dir = os.path.join(self._data_dir, seed_key) os.makedirs(user_data_dir, exist_ok=True) full_args = ( @@ -199,6 +222,9 @@ class ChromePool: process=process, cdp_port=port, user_data_dir=user_data_dir, + timezone=timezone, + locale=locale, + proxy=proxy, ) self._processes[seed_key] = cp @@ -224,6 +250,7 @@ class ChromePool: if self._default is proc: self._default = None self._locks.pop(key, None) + self._connections.pop(key, None) async def shutdown(self) -> None: """Terminate all Chrome processes.""" @@ -307,12 +334,24 @@ def _ws_scheme(request: web.Request) -> str: async def handle_root(request: web.Request) -> web.Response: - """Health check / info.""" + """Health check / process status.""" pool: ChromePool = request.app["pool"] - alive = sum(1 for p in pool._processes.values() if p.process.poll() is None) + processes = {} + for key, proc in pool._processes.items(): + if proc.process.poll() is None: + processes[key] = { + "pid": proc.process.pid, + "port": proc.cdp_port, + "seed": proc.seed, + "connections": pool._connections.get(key, 0), + "timezone": proc.timezone, + "locale": proc.locale, + "proxy": proc.proxy, + } return web.json_response({ "status": "ok", - "processes": alive, + "active": len(processes), + "processes": processes, }) @@ -410,8 +449,6 @@ async def proxy_cdp_websocket( label: str, ) -> None: """Bidirectional WebSocket proxy between client and Chrome CDP.""" - import websockets - try: async with websockets.connect( target_url, max_size=None, ping_interval=None, ping_timeout=None, @@ -463,8 +500,12 @@ async def handle_ws_default(request: web.Request) -> web.WebSocketResponse: ws = web.WebSocketResponse() await ws.prepare(request) - target_url = f"ws://127.0.0.1:{cp.cdp_port}/devtools/{path}" - await proxy_cdp_websocket(ws, target_url, f"CDP default [{path}]") + pool.connect("__default__") + try: + target_url = f"ws://127.0.0.1:{cp.cdp_port}/devtools/{path}" + await proxy_cdp_websocket(ws, target_url, f"CDP default [{path}]") + finally: + pool.disconnect("__default__") return ws @@ -479,8 +520,12 @@ async def handle_ws_seed(request: web.Request) -> web.WebSocketResponse: ws = web.WebSocketResponse() await ws.prepare(request) - target_url = f"ws://127.0.0.1:{cp.cdp_port}/devtools/{path}" - await proxy_cdp_websocket(ws, target_url, f"CDP seed={seed} [{path}]") + pool.connect(seed) + try: + target_url = f"ws://127.0.0.1:{cp.cdp_port}/devtools/{path}" + await proxy_cdp_websocket(ws, target_url, f"CDP seed={seed} [{path}]") + finally: + pool.disconnect(seed) return ws @@ -492,16 +537,25 @@ async def on_shutdown(app: web.Application) -> None: # CLI arg parsing # --------------------------------------------------------------------------- +def _default_data_dir() -> str: + """Smart default: Docker → /tmp/cloakserve, bare metal → ~/.cloakbrowser/cloakserve.""" + if os.path.exists("/.dockerenv"): + return "/tmp/cloakserve" + return str(Path.home() / ".cloakbrowser" / "cloakserve") + + def parse_cli_args(argv: list[str]) -> tuple[dict, list[str]]: """Parse cloakserve-specific args, return (config, passthrough_args).""" - config = { + config: dict = { "port": 9222, "headless": True, + "data_dir": None, } passthrough = [] # Flags consumed by cloakserve (not passed to Chrome) consumed_prefixes = ( "--port=", + "--data-dir=", "--remote-debugging-port=", "--remote-debugging-address=", ) @@ -509,6 +563,8 @@ def parse_cli_args(argv: list[str]) -> tuple[dict, list[str]]: for arg in argv: if arg.startswith("--port="): config["port"] = int(arg.split("=", 1)[1]) + elif arg.startswith("--data-dir="): + config["data_dir"] = arg.split("=", 1)[1] elif arg == "--headless=false" or arg == "--headless=False": config["headless"] = False passthrough.append(arg) @@ -517,6 +573,9 @@ def parse_cli_args(argv: list[str]) -> tuple[dict, list[str]]: else: passthrough.append(arg) + if config["data_dir"] is None: + config["data_dir"] = _default_data_dir() + return config, passthrough @@ -532,6 +591,7 @@ def main() -> None: binary=binary, global_args=global_args, headless=config["headless"], + data_dir=config["data_dir"], ) app = web.Application() diff --git a/tests/test_cloakserve.py b/tests/test_cloakserve.py index c857442..fa60711 100644 --- a/tests/test_cloakserve.py +++ b/tests/test_cloakserve.py @@ -1,9 +1,10 @@ -"""Unit tests for cloakserve — parse_connection_params, parse_cli_args, URL rewriting.""" +"""Unit tests for cloakserve — parse_connection_params, parse_cli_args, URL rewriting, connection tracking.""" import importlib.machinery import importlib.util import sys from pathlib import Path +from unittest.mock import patch import pytest @@ -19,6 +20,8 @@ _loader.exec_module(_mod) parse_connection_params = _mod.parse_connection_params parse_cli_args = _mod.parse_cli_args +ChromePool = _mod.ChromePool +_default_data_dir = _mod._default_data_dir # --------------------------------------------------------------------------- @@ -82,6 +85,7 @@ class TestParseCliArgs: config, passthrough = parse_cli_args([]) assert config["port"] == 9222 assert config["headless"] is True + assert config["data_dir"] is not None assert passthrough == [] def test_custom_port(self): @@ -109,6 +113,24 @@ class TestParseCliArgs: assert "--port=9222" not in passthrough assert "--no-sandbox" in passthrough + def test_custom_data_dir(self): + config, passthrough = parse_cli_args(["--data-dir=/custom/path", "--no-sandbox"]) + assert config["data_dir"] == "/custom/path" + assert "--data-dir=/custom/path" not in passthrough + + def test_data_dir_not_in_passthrough(self): + _, passthrough = parse_cli_args(["--data-dir=/tmp/test"]) + assert not any(a.startswith("--data-dir=") for a in passthrough) + + @patch("os.path.exists", return_value=True) + def test_default_data_dir_docker(self, _mock): + assert _default_data_dir() == "/tmp/cloakserve" + + @patch("os.path.exists", return_value=False) + def test_default_data_dir_bare_metal(self, _mock): + result = _default_data_dir() + assert result.endswith(".cloakbrowser/cloakserve") + # --------------------------------------------------------------------------- # URL rewriting logic (pure string manipulation, extracted from handlers) @@ -169,3 +191,54 @@ class TestURLRewriting: orig = "ws://127.0.0.1:5100/devtools/page/DEF-456" result = self._rewrite_list_entry(orig, "host:443", "seed1", scheme="wss") assert result == "wss://host:443/fingerprint/seed1/devtools/page/DEF-456" + + +# --------------------------------------------------------------------------- +# Connection refcounting +# --------------------------------------------------------------------------- + + +class TestConnectionTracking: + """Test ChromePool.connect() / disconnect() without real Chrome.""" + + def _make_pool(self): + return ChromePool( + binary="/fake/chrome", + global_args=[], + headless=True, + data_dir="/tmp/test-cloakserve", + ) + + def test_connect_increments(self): + pool = self._make_pool() + pool.connect("seed1") + assert pool._connections["seed1"] == 1 + pool.connect("seed1") + assert pool._connections["seed1"] == 2 + + def test_disconnect_decrements(self): + pool = self._make_pool() + pool.connect("seed1") + pool.connect("seed1") + pool.disconnect("seed1") + assert pool._connections["seed1"] == 1 + + def test_disconnect_to_zero_removes_key(self): + pool = self._make_pool() + pool.connect("seed1") + pool.disconnect("seed1") + assert "seed1" not in pool._connections + + def test_disconnect_below_zero_safe(self): + pool = self._make_pool() + pool.disconnect("nonexistent") + assert "nonexistent" not in pool._connections + + def test_multiple_seeds_independent(self): + pool = self._make_pool() + pool.connect("a") + pool.connect("b") + pool.connect("a") + pool.disconnect("a") + assert pool._connections["a"] == 1 + assert pool._connections["b"] == 1