From bc212e727478c263a0bca34906d32a3b7d270c4c Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 17 Apr 2026 22:17:42 +0000 Subject: [PATCH] feat(computer-server): gate service on CONTAINER_NAME via middleware Add UNAVAILABLE_WITHOUT_CONTAINER_NAME env var that, when truthy, causes the server to reject every HTTP and WebSocket request with a configurable status code (UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE, default 503) whenever CONTAINER_NAME is not set -- replacing the silent pass-through to local development mode. Implemented as a single ASGI middleware so the gate applies uniformly across /ws, /cmd, /pty/*, /responses, /playwright_exec, and /status, instead of duplicating an inline check at each auth site. Backwards-compatible: with neither env var set, behavior is unchanged. Integration tests cover the helper, HTTP endpoints, PTY endpoints, the WebSocket handshake, custom status codes, invalid status codes, multiple truthy value variants, and CONTAINER_NAME taking precedence over the flag. https://claude.ai/code/session_01VUPZ5JsZUpYKVKxgaSwYXt --- .../computer-server/computer_server/main.py | 103 ++++++++++ .../tests/test_auth_availability.py | 189 ++++++++++++++++++ 2 files changed, 292 insertions(+) create mode 100644 libs/python/computer-server/tests/test_auth_availability.py diff --git a/libs/python/computer-server/computer_server/main.py b/libs/python/computer-server/computer_server/main.py index 85bdffdc9c..12fc47bdd2 100644 --- a/libs/python/computer-server/computer_server/main.py +++ b/libs/python/computer-server/computer_server/main.py @@ -40,6 +40,38 @@ # Authentication session TTL (in seconds). Override via env var CUA_AUTH_TTL_SECONDS. Default: 60s AUTH_SESSION_TTL_SECONDS: int = int(os.environ.get("CUA_AUTH_TTL_SECONDS", "60")) +# Status code returned when UNAVAILABLE_WITHOUT_CONTAINER_NAME is set and CONTAINER_NAME is missing. +DEFAULT_UNAVAILABLE_STATUS_CODE: int = 503 + + +def _parse_bool_env(name: str) -> bool: + return os.environ.get(name, "").lower().strip() in ("1", "true", "yes", "y", "on") + + +def _unavailable_status_code() -> Optional[int]: + """Return the HTTP status code to use when CONTAINER_NAME is required but unset. + + When ``UNAVAILABLE_WITHOUT_CONTAINER_NAME`` is truthy and ``CONTAINER_NAME`` is not + set, the server should reject requests with the configured status code + (``UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE``, default 503) rather + than passing through to local dev mode. Returns ``None`` when the server should + proceed normally (either because ``CONTAINER_NAME`` is set, or because the + unavailable-without-container flag is not enabled). + """ + if os.environ.get("CONTAINER_NAME"): + return None + if not _parse_bool_env("UNAVAILABLE_WITHOUT_CONTAINER_NAME"): + return None + try: + return int( + os.environ.get( + "UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", + str(DEFAULT_UNAVAILABLE_STATUS_CODE), + ) + ) + except ValueError: + return DEFAULT_UNAVAILABLE_STATUS_CODE + try: from cua_agent import ComputerAgent @@ -76,6 +108,77 @@ redirect_slashes=False, ) +class UnavailableWithoutContainerMiddleware: + """ASGI middleware that rejects all requests when CONTAINER_NAME is required but unset. + + Controlled by env vars (read per-request so tests and dynamic config work): + - ``UNAVAILABLE_WITHOUT_CONTAINER_NAME``: if truthy and ``CONTAINER_NAME`` is unset, + every HTTP and WebSocket request is rejected. + - ``UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE``: HTTP status code for + rejections (default 503). + + When disabled (either env var not set), requests pass through unchanged, preserving + the original "local development mode" behavior for backwards compatibility. + """ + + _DETAIL = "Service unavailable: CONTAINER_NAME is required but not configured" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + scope_type = scope.get("type") + if scope_type in ("http", "websocket"): + status_code = _unavailable_status_code() + if status_code is not None: + if scope_type == "http": + await self._reject_http(send, status_code) + else: + await self._reject_websocket(receive, send, status_code) + return + await self.app(scope, receive, send) + + @classmethod + async def _reject_http(cls, send, status_code): + body = json.dumps({"detail": cls._DETAIL}).encode() + await send( + { + "type": "http.response.start", + "status": status_code, + "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(body)).encode()), + ], + } + ) + await send({"type": "http.response.body", "body": body}) + + @classmethod + async def _reject_websocket(cls, receive, send, status_code): + # Accept first so we can send a structured JSON error before closing — this + # preserves the existing error shape that clients already handle. + event = await receive() + if event.get("type") != "websocket.connect": + return + await send({"type": "websocket.accept"}) + await send( + { + "type": "websocket.send", + "text": json.dumps( + { + "success": False, + "error": cls._DETAIL, + "status_code": status_code, + } + ), + } + ) + # 1008 = Policy Violation + await send({"type": "websocket.close", "code": 1008}) + + +app.add_middleware(UnavailableWithoutContainerMiddleware) + # CORS configuration origins = ["*"] app.add_middleware( diff --git a/libs/python/computer-server/tests/test_auth_availability.py b/libs/python/computer-server/tests/test_auth_availability.py new file mode 100644 index 0000000000..e9c81326dc --- /dev/null +++ b/libs/python/computer-server/tests/test_auth_availability.py @@ -0,0 +1,189 @@ +"""Integration tests for UNAVAILABLE_WITHOUT_CONTAINER_NAME behavior. + +These tests verify two things: + +1. **Backwards compat** — when neither ``CONTAINER_NAME`` nor + ``UNAVAILABLE_WITHOUT_CONTAINER_NAME`` is set, the server continues to + operate in local development mode (no auth required, requests succeed). + +2. **New behavior** — when ``UNAVAILABLE_WITHOUT_CONTAINER_NAME`` is truthy + and ``CONTAINER_NAME`` is unset, requests are rejected with the status + code configured by + ``UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE`` (default 503) + rather than being allowed through. +""" + +from __future__ import annotations + +import pytest + +try: + from fastapi.testclient import TestClient + + from computer_server.main import _unavailable_status_code, app +except Exception as import_error: # pragma: no cover - environment-dependent + pytest.skip( + f"computer_server.main unavailable in this environment: {import_error}", + allow_module_level=True, + ) + + +@pytest.fixture +def clean_env(monkeypatch): + """Remove all env vars that influence auth availability for a clean baseline.""" + for var in ( + "CONTAINER_NAME", + "UNAVAILABLE_WITHOUT_CONTAINER_NAME", + "UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", + ): + monkeypatch.delenv(var, raising=False) + return monkeypatch + + +@pytest.fixture +def client(): + return TestClient(app) + + +class TestUnavailableStatusCode: + """Unit tests for the `_unavailable_status_code` helper.""" + + def test_returns_none_when_both_unset(self, clean_env): + assert _unavailable_status_code() is None + + def test_returns_none_when_container_name_set(self, clean_env): + clean_env.setenv("CONTAINER_NAME", "vm-abc") + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + # CONTAINER_NAME being set overrides the unavailable flag. + assert _unavailable_status_code() is None + + def test_returns_default_503_when_flag_truthy_and_container_missing(self, clean_env): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + assert _unavailable_status_code() == 503 + + @pytest.mark.parametrize("value", ["1", "true", "True", "YES", "y", "on"]) + def test_accepts_various_truthy_values(self, clean_env, value): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", value) + assert _unavailable_status_code() == 503 + + @pytest.mark.parametrize("value", ["0", "false", "no", "", "random"]) + def test_rejects_falsy_values(self, clean_env, value): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", value) + assert _unavailable_status_code() is None + + def test_custom_status_code(self, clean_env): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "1") + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "418") + assert _unavailable_status_code() == 418 + + def test_invalid_status_code_falls_back_to_503(self, clean_env): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "1") + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "not-a-number") + assert _unavailable_status_code() == 503 + + +class TestCmdEndpoint: + """Integration tests for the POST /cmd endpoint.""" + + def test_backwards_compat_local_dev_allows_requests(self, clean_env, client): + # No CONTAINER_NAME, no availability flag — old "local dev" behavior. + resp = client.post("/cmd", json={"command": "version", "params": {}}) + assert resp.status_code == 200, resp.text + assert "success" in resp.text + + def test_unavailable_flag_rejects_with_default_503(self, clean_env, client): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + resp = client.post("/cmd", json={"command": "version", "params": {}}) + assert resp.status_code == 503 + assert "CONTAINER_NAME" in resp.json()["detail"] + + def test_unavailable_flag_with_custom_status_code(self, clean_env, client): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "1") + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "418") + resp = client.post("/cmd", json={"command": "version", "params": {}}) + assert resp.status_code == 418 + + def test_container_name_set_bypasses_unavailable_flag(self, clean_env, client): + # CONTAINER_NAME being set means auth is required — but the unavailable + # flag should NOT apply. Without valid creds, this should 401, not 503. + clean_env.setenv("CONTAINER_NAME", "vm-xyz") + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + resp = client.post("/cmd", json={"command": "version", "params": {}}) + assert resp.status_code == 401 + + +class TestPtyEndpointAuthGate: + """Integration tests for PTY endpoints (via `_require_auth`).""" + + def test_backwards_compat_local_dev_allows_access(self, clean_env, client): + # Use a non-existent PID — we just want to verify we get past the auth gate. + # If auth passes, we get 404 (PTY not found); if not, we get 401/503. + resp = client.get("/pty/999999") + assert resp.status_code == 404, resp.text + + def test_unavailable_flag_rejects_with_default_503(self, clean_env, client): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + resp = client.get("/pty/999999") + assert resp.status_code == 503 + + def test_unavailable_flag_with_custom_status_code(self, clean_env, client): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "599") + resp = client.get("/pty/999999") + assert resp.status_code == 599 + + +class TestPlaywrightExecEndpoint: + def test_backwards_compat_local_dev_accepts_auth(self, clean_env, client): + # Browser manager may fail for other reasons, but it should NOT be 503/401. + resp = client.post("/playwright_exec", json={"command": "noop", "params": {}}) + assert resp.status_code not in (401, 503) + + def test_unavailable_flag_rejects(self, clean_env, client): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + resp = client.post("/playwright_exec", json={"command": "noop", "params": {}}) + assert resp.status_code == 503 + + +class TestStatusEndpointMiddlewareGating: + """The middleware applies uniformly — /status is reachable when the flag is off.""" + + def test_status_accessible_in_local_dev_mode(self, clean_env, client): + resp = client.get("/status") + assert resp.status_code == 200 + assert resp.json()["status"] == "ok" + + def test_status_rejected_by_middleware_when_flag_set(self, clean_env, client): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + resp = client.get("/status") + assert resp.status_code == 503 + + def test_status_accessible_when_container_name_set(self, clean_env, client): + clean_env.setenv("CONTAINER_NAME", "vm-abc") + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + resp = client.get("/status") + assert resp.status_code == 200 + + +class TestWebSocketEndpoint: + def test_backwards_compat_local_dev_allows_commands(self, clean_env, client): + with client.websocket_connect("/ws") as ws: + ws.send_json({"command": "version", "params": {}}) + data = ws.receive_json() + assert data["success"] is True + + def test_unavailable_flag_closes_with_error(self, clean_env, client): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + with client.websocket_connect("/ws") as ws: + data = ws.receive_json() + assert data["success"] is False + assert data["status_code"] == 503 + assert "CONTAINER_NAME" in data["error"] + + def test_unavailable_flag_reports_custom_status_code(self, clean_env, client): + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true") + clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "599") + with client.websocket_connect("/ws") as ws: + data = ws.receive_json() + assert data["success"] is False + assert data["status_code"] == 599