feat: authenticate interpreter requests through envd

Co-Authored-By: mish@e2b.dev <mish@e2b.dev>

Author: devin-ai-integration[bot]

template/server/auth.py

@@ -0,0 +1,96 @@
+import asyncio
+import hashlib
+import hmac
+import time
+
+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
+
+CACHE_TTL = 5.0
+
+
+class AccessTokenVerifier:
+    def __init__(self, client: httpx.AsyncClient):
+        self.client = client
+        self.digest = b""
+        self.validated_at = (0.0, 0.0)
+        self.lock = asyncio.Lock()
+
+    def cached(self, digest: bytes) -> bool:
+        monotonic, wall = self.validated_at
+        return (
+            hmac.compare_digest(self.digest, digest)
+            and 0 <= time.monotonic() - monotonic < CACHE_TTL
+            and 0 <= time.time() - wall < CACHE_TTL
+        )
+
+    async def verify(self, token: str) -> int:
+        try:
+            encoded = token.encode("ascii")
+        except UnicodeEncodeError:
+            return 401
+        if not encoded:
+            return 401
+        digest = hashlib.sha256(encoded).digest()
+        if self.cached(digest):
+            return 204
+
+        async with self.lock:
+            if self.cached(digest):
+                return 204
+
+            # Start the lease before I/O so a slow response cannot extend it.
+            validated_at = (time.monotonic(), time.time())
+            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:
+                self.digest = digest
+                self.validated_at = validated_at
+                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)

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()

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=["*"],

template/tests/test_auth.py

@@ -0,0 +1,199 @@
+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,
+        monotonic=100.0,
+        wall=1000.0,
+        latency=0,
+    )
+    monkeypatch.setattr(
+        auth,
+        "time",
+        SimpleNamespace(monotonic=lambda: state.monotonic, time=lambda: state.wall),
+    )
+    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)
+        state.monotonic += state.latency
+        state.wall += state.latency
+        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_cache_expires_without_sliding_and_replaces_rotated_token(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"
+
+    state.monotonic += 4
+    state.wall += 4
+    assert (await client.get("/contexts", headers=headers)).status_code == 200
+    assert len(state.calls) == 1
+
+    state.monotonic += 1
+    state.wall += 1
+    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
+
+
+async def test_concurrent_requests_share_validation_and_wrong_tokens_cannot_hit_cache(
+    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) == 1
+    for _ in range(2):
+        assert (
+            await client.get("/contexts", headers={"X-Access-Token": "wrong"})
+        ).status_code == 401
+    assert len(state.calls) == 3
+    assert (
+        await client.get("/contexts", headers={"X-Access-Token": state.token})
+    ).status_code == 200
+    assert len(state.calls) == 3
+
+
+@pytest.mark.parametrize(
+    "failure", [200, 302, 404, 500, 503, httpx.ConnectError("offline")]
+)
+async def test_expired_cache_fails_closed_and_recovers(server, failure):
+    client, state = server
+    headers = {"X-Access-Token": state.token}
+    assert (await client.get("/contexts", headers=headers)).status_code == 200
+    state.monotonic += 5
+    state.wall += 5
+    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
+
+
+@pytest.mark.parametrize(
+    "clock,delta", [("wall", 10), ("wall", -10), ("monotonic", -10)]
+)
+async def test_clock_jumps_invalidate_cached_validation(server, clock, delta):
+    client, state = server
+    headers = {"X-Access-Token": state.token}
+    assert (await client.get("/contexts", headers=headers)).status_code == 200
+    setattr(state, clock, getattr(state, clock) + delta)
+    state.token = "rotated-token"
+    assert (await client.get("/contexts", headers=headers)).status_code == 401
+    assert len(state.calls) == 2
+
+
+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_slow_response_does_not_extend_cache_lifetime(server):
+    client, state = server
+    state.latency = 5
+    headers = {"X-Access-Token": state.token}
+    assert (await client.get("/contexts", headers=headers)).status_code == 200
+    state.token = "rotated-token"
+    assert (await client.get("/contexts", headers=headers)).status_code == 401
+    assert len(state.calls) == 2
+
+
+async def test_malformed_headers_cannot_use_cached_validation(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