diff --git a/template/server/auth.py b/template/server/auth.py new file mode 100644 index 00000000..46ccbd02 --- /dev/null +++ b/template/server/auth.py @@ -0,0 +1,62 @@ +import httpx +from starlette.datastructures import Headers +from starlette.responses import PlainTextResponse +from starlette.types import ASGIApp, Receive, Scope, Send + +from envs import ENVD_PORT + + +class AccessTokenVerifier: + def __init__(self, client: httpx.AsyncClient): + self.client = client + + async def verify(self, token: str) -> int: + if not token or not token.isascii(): + return 401 + try: + response = await self.client.get( + f"http://127.0.0.1:{ENVD_PORT}/auth", + headers={"X-Access-Token": token}, + ) + except httpx.HTTPError: + return 503 + + if response.status_code == 204: + return 204 + return 401 if response.status_code == 401 else 503 + + +class EnvdAuthMiddleware: + def __init__(self, app: ASGIApp, local: bool = False): + self.app = app + self.local = local + + async def __call__(self, scope: Scope, receive: Receive, send: Send): + if scope["type"] == "websocket" and not self.local: + await send({"type": "websocket.close", "code": 1008}) + return + + if ( + scope["type"] != "http" + or self.local + or (scope["method"] == "GET" and scope["path"] == "/health") + ): + await self.app(scope, receive, send) + return + + tokens = Headers(scope=scope).getlist("x-access-token") + status = 401 + if len(tokens) == 1: + verifier = getattr(scope["app"].state, "access_token_verifier", None) + status = await verifier.verify(tokens[0]) if verifier else 503 + + if status != 204: + response = PlainTextResponse( + "Unauthorized" if status == 401 else "Authentication unavailable", + status_code=status, + headers={"Cache-Control": "no-store"}, + ) + await response(scope, receive, send) + return + + await self.app(scope, receive, send) diff --git a/template/server/envs.py b/template/server/envs.py index 8752e7b9..9f3fa8b5 100644 --- a/template/server/envs.py +++ b/template/server/envs.py @@ -3,18 +3,19 @@ import httpx -LOCAL = os.getenv("E2B_LOCAL", False) +LOCAL = os.getenv("E2B_LOCAL", "").lower() == "true" ENVD_PORT = 49983 async def get_envs(access_token: Optional[str]) -> dict: if LOCAL: return {"E2B_TEST_VARIABLE": "true"} - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(trust_env=False) as client: headers = {} if access_token: headers["X-Access-Token"] = f"{access_token}" response = await client.get( f"http://localhost:{ENVD_PORT}/envs", headers=headers ) + response.raise_for_status() return response.json() diff --git a/template/server/main.py b/template/server/main.py index a8d43b74..1a48e562 100644 --- a/template/server/main.py +++ b/template/server/main.py @@ -12,8 +12,10 @@ from api.models.context import Context from api.models.create_context import CreateContext from api.models.execution_request import ExecutionRequest +from auth import AccessTokenVerifier, EnvdAuthMiddleware from consts import JUPYTER_BASE_URL, JUPYTER_SOCKET_PATH from contexts import create_context, normalize_language +from envs import LOCAL from messaging import ContextWebSocket from stream import StreamingListJsonResponse from utils.locks import LockedMap @@ -35,6 +37,8 @@ async def lifespan(app: FastAPI): client = httpx.AsyncClient( transport=httpx.AsyncHTTPTransport(uds=JUPYTER_SOCKET_PATH), trust_env=False ) + auth_client = httpx.AsyncClient(trust_env=False, timeout=5.0) + app.state.access_token_verifier = AccessTokenVerifier(auth_client) try: python_context = await create_context( @@ -59,10 +63,14 @@ async def lifespan(app: FastAPI): except Exception as e: logger.error(f"Failed to initialize default context: {e}") raise + finally: + del app.state.access_token_verifier + await auth_client.aclose() app = FastAPI(lifespan=lifespan) +app.add_middleware(EnvdAuthMiddleware, local=LOCAL) app.add_middleware( CORSMiddleware, allow_origins=["*"], diff --git a/template/tests/test_auth.py b/template/tests/test_auth.py new file mode 100644 index 00000000..c774af57 --- /dev/null +++ b/template/tests/test_auth.py @@ -0,0 +1,157 @@ +import asyncio +from types import SimpleNamespace + +import httpx +import pytest +from starlette.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +import auth +import main + + +@pytest.fixture +async def server(monkeypatch): + state = SimpleNamespace(token="valid-token", calls=[], status=None) + monkeypatch.setattr(main, "websockets", {}) + + async def envd(request): + assert str(request.url) == "http://127.0.0.1:49983/auth" + state.calls.append(request) + await asyncio.sleep(0) + if isinstance(state.status, Exception): + raise state.status + status = state.status + if status is None: + status = 204 if request.headers["X-Access-Token"] == state.token else 401 + return httpx.Response(status) + + async with httpx.AsyncClient(transport=httpx.MockTransport(envd)) as upstream: + monkeypatch.setattr( + main.app.state, + "access_token_verifier", + auth.AccessTokenVerifier(upstream), + raising=False, + ) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main.app), base_url="http://server" + ) as client: + yield client, state + + +@pytest.mark.parametrize( + "method,path", + [ + ("POST", "/execute"), + ("POST", "/contexts"), + ("GET", "/contexts"), + ("POST", "/contexts/id/restart"), + ("DELETE", "/contexts/id"), + ("GET", "/docs"), + ], +) +async def test_requests_require_auth_before_route_handling(server, method, path): + client, state = server + for headers in ({}, {"X-Access-Token": "wrong"}): + response = await client.request(method, path, headers=headers) + assert response.status_code == 401 + assert response.headers["cache-control"] == "no-store" + assert len(state.calls) == 1 + + +async def test_each_request_is_validated(server): + client, state = server + for _ in range(2): + response = await client.get( + "/contexts", headers={"X-Access-Token": state.token} + ) + assert response.status_code == 200 + assert response.json() == [] + assert len(state.calls) == 2 + + +async def test_token_rotation_and_revocation_apply_to_next_request(server): + client, state = server + headers = {"X-Access-Token": state.token} + response = await client.get("/contexts", headers=headers) + assert response.status_code == 200 + assert response.json() == [] + state.token = "rotated-token" + + assert (await client.get("/contexts", headers=headers)).status_code == 401 + assert len(state.calls) == 2 + assert ( + await client.get("/contexts", headers={"X-Access-Token": state.token}) + ).status_code == 200 + assert len(state.calls) == 3 + assert (await client.get("/contexts", headers=headers)).status_code == 401 + headers = {"X-Access-Token": state.token} + state.token = None + assert (await client.get("/contexts", headers=headers)).status_code == 401 + assert len(state.calls) == 5 + + +async def test_concurrent_requests_are_validated_individually(server): + client, state = server + responses = await asyncio.gather( + *[ + client.get("/contexts", headers={"X-Access-Token": state.token}) + for _ in range(10) + ] + ) + assert all(response.status_code == 200 for response in responses) + assert len(state.calls) == 10 + + +@pytest.mark.parametrize( + "failure", [200, 302, 404, 500, 503, httpx.ConnectError("offline")] +) +async def test_envd_failures_reject_previously_valid_token_and_recover(server, failure): + client, state = server + headers = {"X-Access-Token": state.token} + assert (await client.get("/contexts", headers=headers)).status_code == 200 + state.status = failure + for _ in range(2): + assert (await client.get("/contexts", headers=headers)).status_code == 503 + assert len(state.calls) == 3 + state.status = None + assert (await client.get("/contexts", headers=headers)).status_code == 200 + assert len(state.calls) == 4 + + +async def test_health_and_cors_preflight_do_not_require_auth(server): + client, state = server + assert (await client.get("/health")).status_code == 200 + assert (await client.get("/health/")).status_code == 401 + response = await client.options( + "/contexts", + headers={ + "Origin": "https://example.com", + "Access-Control-Request-Method": "GET", + "Access-Control-Request-Headers": "X-Access-Token", + }, + ) + assert response.status_code == 200 + assert response.headers["access-control-allow-origin"] == "*" + assert not state.calls + + +async def test_malformed_headers_are_rejected_without_calling_envd(server): + client, state = server + assert ( + await client.get("/contexts", headers={"X-Access-Token": state.token}) + ).status_code == 200 + for headers in ( + [("X-Access-Token", state.token), ("X-Access-Token", "wrong")], + [(b"X-Access-Token", b"\xff")], + {"X-Access-Token": ""}, + ): + assert (await client.get("/contexts", headers=headers)).status_code == 401 + assert len(state.calls) == 1 + + +def test_public_websockets_are_rejected(): + with pytest.raises(WebSocketDisconnect) as exc: + with TestClient(main.app).websocket_connect("/contexts"): + pytest.fail("public websocket accepted") + assert exc.value.code == 1008