diff --git a/.github/workflows/server_tests.yml b/.github/workflows/server_tests.yml new file mode 100644 index 00000000..0ba67302 --- /dev/null +++ b/.github/workflows/server_tests.yml @@ -0,0 +1,33 @@ +name: Interpreter Server Tests + +on: + pull_request: + +permissions: + contents: read + +jobs: + test: + name: Interpreter Server Tests + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Checkout repository + uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 + + - name: Parse .tool-versions + uses: wistia/parse-tool-versions@32f568a4ffd4bfa7720ebf93f171597d1ebc979a # v2.1.1 + with: + filename: '.tool-versions' + uppercase: 'true' + prefix: 'tool_version_' + + - name: Install uv + uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6.8.0 + with: + version: '${{ env.TOOL_VERSION_UV }}' + python-version: '${{ env.TOOL_VERSION_PYTHON }}' + enable-cache: true + + - name: Test interpreter server + run: uv run --with-requirements template/requirements-test.txt pytest template/tests diff --git a/template/README.md b/template/README.md index 4f64fabe..209f5a9a 100644 --- a/template/README.md +++ b/template/README.md @@ -122,3 +122,28 @@ Inside a running sandbox you can also inspect things directly: journalctl -u jupyter -u code-interpreter systemctl status code-interpreter ``` + +## Interpreter authentication + +The template requires envd 0.9.1 or later with the local token delivery socket. +The server receives a SHA-512 token verifier over `/run/e2b-auth/token.sock` and +validates `X-Access-Token` locally. It rejects requests before provisioning and +after disconnecting. `/health` needs no bearer token, but remains unavailable +until the auth connection is established. Tokenless sandboxes cannot execute code. + +The systemd service creates `/etc/e2b/code-interpreter-auth` when first started. +This marker contains no credentials; it makes envd wait for token delivery to be +acknowledged during initialization, including resume and fork. The auth connection +is also reestablished after an interpreter or envd restart. + +Deploy envd and SDK context-auth updates before rebuilding this template. Local +Docker development explicitly sets `E2B_LOCAL=true`, which disables auth; do not +use that setting in deployed templates. Authentication protects the interpreter +API on port `49999`; it does not change Jupyter's listener or secure direct +access to Jupyter. Jupyter isolation is a separate transport change. + +Run the authentication tests locally with: + +```sh +uv run --with-requirements template/requirements-test.txt pytest template/tests/test_auth.py +``` diff --git a/template/pytest.ini b/template/pytest.ini new file mode 100644 index 00000000..e2378df8 --- /dev/null +++ b/template/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +pythonpath = server +asyncio_mode = auto +testpaths = tests diff --git a/template/requirements-test.txt b/template/requirements-test.txt new file mode 100644 index 00000000..134baa42 --- /dev/null +++ b/template/requirements-test.txt @@ -0,0 +1,5 @@ +-r server/requirements.txt +pytest==9.0.3 +pytest-asyncio==1.3.0 +jupyter-server==2.20.0 +ipykernel==6.31.0 diff --git a/template/server/auth.py b/template/server/auth.py new file mode 100644 index 00000000..574ecc83 --- /dev/null +++ b/template/server/auth.py @@ -0,0 +1,94 @@ +import asyncio +import hashlib +import hmac +import logging +from contextlib import suppress + +from starlette.requests import Request +from starlette.responses import PlainTextResponse + +TOKEN_SOCKET = "/run/e2b-auth/token.sock" +logger = logging.getLogger(__name__) + + +class AccessToken: + def __init__(self): + self.digest = None + self.connected = False + + def accepts(self, token: str) -> bool: + return bool( + self.digest + and token + and hmac.compare_digest( + self.digest, hashlib.sha512(token.encode("utf-8")).digest() + ) + ) + + async def subscribe(self): + while True: + writer = None + try: + reader, writer = await asyncio.open_unix_connection( + TOKEN_SOCKET, limit=129 + ) + while True: + line = await reader.readline() + if line == b"\n": + self.digest = None + elif len(line) == 129 and line.endswith(b"\n"): + digest = bytes.fromhex(line[:-1].decode("ascii")) + if len(digest) != 64: + raise ValueError("Invalid token verifier") + self.digest = digest + else: + raise ValueError("Invalid token verifier") + self.connected = True + writer.write(b"+") + await writer.drain() + except (OSError, ValueError, UnicodeError): + logger.warning("Local authentication connection unavailable") + finally: + self.digest = None + self.connected = False + if writer is not None: + writer.close() + with suppress(OSError): + await writer.wait_closed() + await asyncio.sleep(1) + + +class EnvdAuthMiddleware: + def __init__(self, app, access_token: AccessToken, local: bool = False): + self.app = app + self.access_token = access_token + self.local = local + + async def __call__(self, scope, receive, send): + if ( + self.local + or scope["type"] not in ("http", "websocket") + or ( + scope["type"] == "http" + and scope["method"] == "GET" + and scope["path"] == "/health" + ) + ): + await self.app(scope, receive, send) + return + + if scope["type"] == "websocket": + await send({"type": "websocket.close", "code": 1008}) + return + + token = Request(scope).headers.get("X-Access-Token", "") + if not self.access_token.accepts(token): + status = 401 if self.access_token.digest else 503 + response = PlainTextResponse( + "Unauthorized" if status == 401 else "Authentication unavailable", + status_code=status, + ) + 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..b12ebb5f 100644 --- a/template/server/envs.py +++ b/template/server/envs.py @@ -3,18 +3,25 @@ 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 ) - return response.json() + response.raise_for_status() + envs = response.json() + if not isinstance(envs, dict) or not all( + isinstance(key, str) and isinstance(value, str) + for key, value in envs.items() + ): + raise ValueError("Invalid environment variables from envd") + return envs diff --git a/template/server/main.py b/template/server/main.py index 59215f57..2a818bba 100644 --- a/template/server/main.py +++ b/template/server/main.py @@ -1,10 +1,11 @@ +import asyncio import logging import sys import httpx from typing import Dict, Union, Literal, List -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, suppress from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import PlainTextResponse @@ -14,6 +15,8 @@ from api.models.execution_request import ExecutionRequest from consts import JUPYTER_BASE_URL from contexts import create_context, normalize_language +from auth import AccessToken, EnvdAuthMiddleware +from envs import LOCAL from messaging import ContextWebSocket from stream import StreamingListJsonResponse from utils.locks import LockedMap @@ -27,6 +30,7 @@ websockets: Dict[Union[str, Literal["default"]], ContextWebSocket] = {} default_websockets = LockedMap() global client +access_token = AccessToken() @asynccontextmanager @@ -34,6 +38,7 @@ async def lifespan(app: FastAPI): global client client = httpx.AsyncClient() + auth_task = None if LOCAL else asyncio.create_task(access_token.subscribe()) try: python_context = await create_context( client, websockets, "python", "/home/user" @@ -57,10 +62,16 @@ async def lifespan(app: FastAPI): except Exception as e: logger.error(f"Failed to initialize default context: {e}") raise + finally: + if auth_task is not None: + auth_task.cancel() + with suppress(asyncio.CancelledError): + await auth_task app = FastAPI(lifespan=lifespan) +app.add_middleware(EnvdAuthMiddleware, access_token=access_token, local=LOCAL) app.add_middleware( CORSMiddleware, allow_origins=["*"], @@ -73,6 +84,8 @@ async def lifespan(app: FastAPI): @app.get("/health") async def get_health(): + if not LOCAL and not access_token.connected: + return PlainTextResponse("Authentication unavailable", status_code=503) return "OK" @@ -122,7 +135,7 @@ async def post_execute(request: Request, exec_request: ExecutionRequest): ws.execute( exec_request.code, env_vars=exec_request.env_vars, - access_token=request.headers.get("X-Access-Token", None), + access_token=request.headers.get("X-Access-Token"), ) ) diff --git a/template/server/messaging.py b/template/server/messaging.py index c51f8b21..08aa0e16 100644 --- a/template/server/messaging.py +++ b/template/server/messaging.py @@ -319,7 +319,7 @@ async def execute( self, code: Union[str, StrictStr], env_vars: Dict[StrictStr, str], - access_token: str, + access_token: Optional[str], ): if self._ws is None: raise Exception("WebSocket not connected") diff --git a/template/systemd/code-interpreter.service b/template/systemd/code-interpreter.service index 04c165b4..ba04808c 100644 --- a/template/systemd/code-interpreter.service +++ b/template/systemd/code-interpreter.service @@ -9,6 +9,8 @@ StartLimitBurst=0 [Service] Type=simple WorkingDirectory=/root/.server +ExecStartPre=/usr/bin/install -d -m 0755 /etc/e2b +ExecStartPre=/usr/bin/touch /etc/e2b/code-interpreter-auth ExecStartPre=/root/.jupyter/jupyter-healthcheck.sh ExecStart=/root/.server/.venv/bin/uvicorn main:app --host 0.0.0.0 --port 49999 --workers 1 --no-access-log --no-use-colors --timeout-keep-alive 640 Restart=on-failure diff --git a/template/tests/test_auth.py b/template/tests/test_auth.py new file mode 100644 index 00000000..d95a63dd --- /dev/null +++ b/template/tests/test_auth.py @@ -0,0 +1,119 @@ +import asyncio +import hashlib +from contextlib import suppress + +import httpx +import pytest +from anyio import fail_after +from starlette.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +import auth +import main + + +@pytest.fixture +async def client(monkeypatch): + monkeypatch.setattr(main.access_token, "digest", None) + monkeypatch.setattr(main.access_token, "connected", False) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main.app), base_url="http://interpreter" + ) as client: + yield client + + +@pytest.mark.parametrize( + "method,path", + [ + ("POST", "/execute"), + ("POST", "/contexts"), + ("GET", "/contexts"), + ("DELETE", "/contexts/missing"), + ("POST", "/contexts/missing/restart"), + ("GET", "/docs"), + ], +) +async def test_endpoints_require_current_token(client, method, path): + response = await client.request(method, path) + assert response.status_code == 503 + main.access_token.digest = hashlib.sha512(b"sandbox-token").digest() + for token in ("", "incorrect", main.access_token.digest.hex()): + response = await client.request(method, path, headers={"X-Access-Token": token}) + assert response.status_code == 401 + + +async def test_health_requires_delivery_connection_but_no_bearer(client): + assert (await client.get("/health")).status_code == 503 + main.access_token.connected = True + assert (await client.get("/health")).status_code == 200 + assert (await client.get("/contexts")).status_code == 503 + + +def test_public_websockets_are_not_exposed(): + with pytest.raises(WebSocketDisconnect) as exc: + with TestClient(main.app).websocket_connect("/api/kernels/kernel/channels"): + pytest.fail("Jupyter must not be exposed through the public server") + assert exc.value.code == 1008 + + +async def test_socket_updates_revoke_old_token_and_disconnect_fails_closed( + client, tmp_path, monkeypatch +): + path = str(tmp_path / "auth.sock") + monkeypatch.setattr(auth, "TOKEN_SOCKET", path) + connections = asyncio.Queue() + + async def connected(reader, writer): + await connections.put((reader, writer)) + + server = await asyncio.start_unix_server(connected, path=path) + task = asyncio.create_task(main.access_token.subscribe()) + try: + reader, writer = await asyncio.wait_for(connections.get(), 2) + for current, previous in ( + ("parent-token", "wrong"), + ("child-token", "parent-token"), + ): + writer.write(hashlib.sha512(current.encode()).hexdigest().encode() + b"\n") + await writer.drain() + assert await asyncio.wait_for(reader.readexactly(1), 2) == b"+" + assert ( + await client.get("/contexts", headers={"X-Access-Token": current}) + ).status_code == 200 + assert ( + await client.get("/contexts", headers={"X-Access-Token": previous}) + ).status_code == 401 + + writer.close() + await writer.wait_closed() + with fail_after(2): + while main.access_token.connected: + await asyncio.sleep(0.01) + assert ( + await client.get("/contexts", headers={"X-Access-Token": "child-token"}) + ).status_code == 503 + + reader, writer = await asyncio.wait_for(connections.get(), 3) + writer.write(hashlib.sha512(b"child-token").hexdigest().encode() + b"\n") + await writer.drain() + assert await asyncio.wait_for(reader.readexactly(1), 2) == b"+" + assert ( + await client.get("/contexts", headers={"X-Access-Token": "child-token"}) + ).status_code == 200 + + writer.write(b"\n") + await writer.drain() + assert await asyncio.wait_for(reader.readexactly(1), 2) == b"+" + assert (await client.get("/contexts")).status_code == 503 + writer.write(b"invalid verifier\n") + await writer.drain() + assert await asyncio.wait_for(reader.read(), 2) == b"" + assert (await client.get("/health")).status_code == 503 + writer.close() + await writer.wait_closed() + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + server.close() + await server.wait_closed()