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