"""Browser-usable gateway for DeepSeek Harness Web (DSH) sandboxes. This is the control + data plane in front of the agent-sandbox CRDs. It is a single teaching-friendly module with three jobs: * **Users** -- create a ``SandboxClaim`` over a warm pool, mint a bearer token (only its SHA-256 is stored, as a claim annotation), report derived state, and delete the claim (which cascades the sandbox and its PVC). * **Browser sessions** -- exchange a user token for an HttpOnly session cookie so the DSH UI can be opened in a browser at ``/u//``. Browser requests authenticate with that cookie; automation can keep using ``Authorization: Bearer ``. * **Reverse proxy** -- HTTP (including SSE) and WebSocket proxying to the sandbox's DSH Web server on ``:3080``. DSH itself is cookie-gated: the gateway scrapes the launch token from the pod logs, performs the ``?token=`` handshake server-side, caches the resulting ``dsh-auth-*`` cookie, and injects it upstream. The browser never sees the launch or bearer token. A background sweeper suspends sandboxes idle for ``IDLE_TIMEOUT`` seconds. Suspending deletes the pod only; the PVC (and the user's state) survives. Deliberately NOT production code: single replica, in-memory state and token cache, no TLS termination, and launch-token discovery by log scraping. Run locally with:: uvicorn gateway:create_app --factory --host 0.0.0.0 --port 8080 The factory signature (``Settings``, ``Kubernetes``, an httpx client and the websockets connector are all injectable) is what makes the test suite cluster-free. """ from __future__ import annotations import asyncio import contextlib import hashlib import html import json import os import re import secrets import time from dataclasses import dataclass, field from urllib.parse import parse_qs, parse_qsl, quote, urlencode, urlsplit, urlunsplit import anyio import httpx import websockets from kubernetes import client as k8s_client from kubernetes import config as k8s_config from kubernetes.client.exceptions import ApiException from starlette.applications import Starlette from starlette.datastructures import Headers from starlette.requests import Request from starlette.responses import ( HTMLResponse, JSONResponse, RedirectResponse, Response, StreamingResponse, ) from starlette.routing import Route, WebSocketRoute from starlette.websockets import WebSocket # --- Kubernetes / DSH constants ------------------------------------------- GROUP, VERSION = "extensions.agents.x-k8s.io", "v1beta1" CORE_GROUP = "agents.x-k8s.io" TOKEN_ANNOTATION = "aaas.example.com/token-sha256" # Ownership marker: the gateway shares its namespace with claims applied by # hand (the resource-layer walkthrough in the example README). Every claim is # named "dsh-", so without an explicit marker the sweeper cannot tell its # own users from a claim someone created manually. MANAGED_ANNOTATION = "aaas.example.com/managed-by" MANAGED_BY = "dsh-gateway" POD_NAME_ANNOTATION = "agents.x-k8s.io/pod-name" USER_LABEL = "sandbox.users.io/dsh-user" # RFC 1123 label: lowercase alphanumerics and "-", must start/end alphanumeric. # The claim name is "dsh-", so cap the user part to keep the claim a # valid Kubernetes object name. DNS1123_LABEL = re.compile(r"^[a-z0-9]([-a-z0-9]*[a-z0-9])?$") MAX_USER_LEN = 63 - len("dsh-") # Headers a proxy must not forward in either direction (RFC 7230 section 6.1). HOP_BY_HOP = frozenset( { "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "trailers", "transfer-encoding", "upgrade", } ) SAFE_METHODS = frozenset({"GET", "HEAD", "OPTIONS"}) # Headers the gateway owns: never forwarded from the browser, overwritten with # the gateway's own view of the connection instead. OWNED_REQUEST_HEADERS = HOP_BY_HOP | { "authorization", "cookie", "host", "x-forwarded-for", "x-forwarded-host", "x-forwarded-port", "x-forwarded-proto", "x-real-ip", } _DSH_WEB_LINE = re.compile(r"dsh web:\s+(?Phttps?://\S+)") class UserExists(Exception): """The claim already exists (a 409 from the API server).""" class BodyTooLarge(Exception): """The request body exceeded ``MAX_BODY_MB``.""" class HandshakeError(Exception): """DSH's launch token or auth cookie could not be obtained.""" class SandboxUnavailable(Exception): """The sandbox is not (yet) ready to serve traffic.""" class UpstreamError(Exception): """The upstream DSH connection failed.""" # --- Pure helpers (unit-testable without an app) -------------------------- def valid_user(user: str) -> bool: """A user is usable as a claim-name suffix (dns-1123 label, length-capped).""" return ( isinstance(user, str) and len(user) <= MAX_USER_LEN and DNS1123_LABEL.fullmatch(user) is not None ) def claim_name(user: str) -> str: return f"dsh-{user}" def claim_user(claim: dict) -> str: return claim["metadata"]["name"].removeprefix("dsh-") def gateway_managed(claim: dict) -> bool: """True only for claims this gateway created. The sweeper must not suspend claims it does not own: a claim applied from a resource-layer manifest (e.g. ``30-claim-alice.yaml``) has the same ``dsh-`` name but no ownership marker, and suspending it deletes a pod the operator is using. """ annotations = (claim.get("metadata") or {}).get("annotations") or {} return annotations.get(MANAGED_ANNOTATION) == MANAGED_BY def condition(obj: dict, ctype: str) -> bool: for c in (obj.get("status") or {}).get("conditions", []): if c.get("type") == ctype: return c.get("status") == "True" return False def derive_state(sandbox: dict | None) -> str: """The state machine a platform shows users (mode x conditions).""" if sandbox is None: return "Provisioning" mode = sandbox.get("spec", {}).get("operatingMode", "Running") if mode == "Suspended": return "Suspended" if condition(sandbox, "Suspended") else "Suspending" return "Ready" if condition(sandbox, "Ready") else "Waking" def token_matches(claim: dict, token: str | None) -> bool: """Timing-safe check of a user token against the stored SHA-256.""" want = (claim.get("metadata", {}).get("annotations") or {}).get( TOKEN_ANNOTATION, "" ) if not token or not want: return False got = hashlib.sha256(token.encode()).hexdigest() return secrets.compare_digest(got, want) def parse_dsh_launch_token(logs: str) -> str | None: """Return the token from the last ``dsh web: http://...?token=...`` line. DSH prints the launch URL once at startup. We scan the tail of the pod log and keep the last match so a restarted process (which prints a new token) wins over the stale one. """ token = None for line in (logs or "").splitlines(): match = _DSH_WEB_LINE.search(line) if not match: continue values = parse_qs(urlsplit(match.group("url")).query).get("token") if values and values[0]: token = values[0] return token def extract_dsh_cookies(set_cookie_headers: list[str]) -> str: """Join the ``dsh-auth-*`` pairs from Set-Cookie headers into a Cookie value.""" pairs = [] for raw in set_cookie_headers: first = raw.split(";", 1)[0] name, sep, value = first.partition("=") name = name.strip() if sep and name.startswith("dsh-auth-"): pairs.append(f"{name}={value.strip()}") return "; ".join(pairs) def same_origin(headers) -> bool: """True when Origin (or Referer) names the same authority as Host.""" host = (headers.get("host") or "").strip().lower() origin = headers.get("origin") if origin: return urlsplit(origin).netloc.lower() == host referer = headers.get("referer") if referer: return urlsplit(referer).netloc.lower() == host return False def rewrite_location(location: str, user: str) -> str: """Prefix absolute-path redirects so the browser stays under ``/u//``.""" if not location.startswith("/") or location.startswith("//"): return location parts = urlsplit(location) prefix = f"/u/{user}" if parts.path == prefix or parts.path.startswith(f"{prefix}/"): return location query = [ (key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key != "token" ] return urlunsplit(("", "", prefix + parts.path, urlencode(query), parts.fragment)) def filtered_query(query_string: str) -> str: """Drop ``?token=`` before a request is sent upstream.""" pairs = [ (key, value) for key, value in parse_qsl(query_string, keep_blank_values=True) if key != "token" ] return urlencode(pairs) def _env_int(name: str, default: int) -> int: try: return int(os.environ.get(name, default)) except ValueError: return default def _env_bool(name: str, default: bool) -> bool: raw = os.environ.get(name) if raw is None: return default return raw.strip().lower() in {"1", "true", "yes", "on"} def _env_float(name: str, default: float) -> float: try: return float(os.environ.get(name, default)) except ValueError: return default # --- Settings -------------------------------------------------------------- @dataclass(frozen=True) class Settings: """Runtime configuration, read from the environment by default. Defaults are the demo values from the example manifests. Everything is injected into ``create_app`` so tests never need the environment or a cluster. """ namespace: str = "dsh-demo" pool: str = "dsh-pool" idle_timeout: int = 60 wake_timeout: int = 120 dsh_port: int = 3080 dsh_container: str = "dsh-web" dsh_log_tail_lines: int = 2000 cookie_name: str = "dsh_gw_session" cookie_secure: bool = False cookie_max_age: int = 7 * 24 * 60 * 60 max_body_bytes: int = 32 * 1024 * 1024 http_timeout: float = 300.0 sweep_interval: float = 15.0 # Connection retries while a Ready pod is still binding its port. upstream_attempts: int = 3 upstream_retry_delay: float = 1.0 @classmethod def from_env(cls) -> Settings: return cls( namespace=os.environ.get("NAMESPACE", "dsh-demo"), pool=os.environ.get("POOL", "dsh-pool"), idle_timeout=_env_int("IDLE_TIMEOUT", 60), wake_timeout=_env_int("WAKE_TIMEOUT", 120), dsh_port=_env_int("DSH_PORT", 3080), dsh_container=os.environ.get("DSH_CONTAINER", "dsh-web"), dsh_log_tail_lines=_env_int("DSH_LOG_TAIL_LINES", 2000), cookie_name=os.environ.get("COOKIE_NAME", "dsh_gw_session"), cookie_secure=_env_bool("COOKIE_SECURE", False), cookie_max_age=_env_int("COOKIE_MAX_AGE", 7 * 24 * 60 * 60), max_body_bytes=_env_int("MAX_BODY_MB", 32) * 1024 * 1024, http_timeout=_env_float("HTTP_TIMEOUT", 300.0), sweep_interval=_env_float("SWEEP_INTERVAL", 15.0), ) # --- Kubernetes adapter ---------------------------------------------------- # # The official Kubernetes client is synchronous. Every call is pushed onto a # worker thread with anyio.to_thread so it never blocks the event loop. The # client and its credentials are loaded lazily, on the first real call: that # keeps ``GET /healthz`` (and importing this module) cluster-free. class Kubernetes: """Thin async wrapper over the sync CustomObjects/CoreV1 API clients.""" def __init__(self, namespace: str = "dsh-demo"): self.namespace = namespace self._crd: k8s_client.CustomObjectsApi | None = None self._core: k8s_client.CoreV1Api | None = None def _clients(self): if self._crd is None: try: k8s_config.load_incluster_config() except k8s_config.ConfigException: k8s_config.load_kube_config() self._crd = k8s_client.CustomObjectsApi() self._core = k8s_client.CoreV1Api() return self._crd, self._core async def _call(self, which: str, method: str, *args, **kwargs): def op(): crd, core = self._clients() target = crd if which == "crd" else core return getattr(target, method)(*args, **kwargs) return await anyio.to_thread.run_sync(op) # -- claims ------------------------------------------------------------ async def get_claim(self, user: str) -> dict | None: try: return await self._call( "crd", "get_namespaced_custom_object", GROUP, VERSION, self.namespace, "sandboxclaims", claim_name(user), ) except ApiException as exc: if exc.status == 404: return None raise async def create_claim(self, user: str, token_hash: str, pool: str) -> None: body = { "apiVersion": f"{GROUP}/{VERSION}", "kind": "SandboxClaim", "metadata": { "name": claim_name(user), "annotations": { TOKEN_ANNOTATION: token_hash, MANAGED_ANNOTATION: MANAGED_BY, }, }, "spec": { "warmPoolRef": {"name": pool}, "additionalPodMetadata": {"labels": {USER_LABEL: user}}, }, } try: await self._call( "crd", "create_namespaced_custom_object", GROUP, VERSION, self.namespace, "sandboxclaims", body, ) except ApiException as exc: if exc.status == 409: raise UserExists(user) from exc raise async def delete_claim(self, user: str) -> None: try: await self._call( "crd", "delete_namespaced_custom_object", GROUP, VERSION, self.namespace, "sandboxclaims", claim_name(user), ) except ApiException as exc: if exc.status != 404: # already gone => idempotent success raise async def list_claims(self) -> list[dict]: result = await self._call( "crd", "list_namespaced_custom_object", GROUP, VERSION, self.namespace, "sandboxclaims", ) return result.get("items", []) # -- sandboxes --------------------------------------------------------- async def get_sandbox(self, name: str) -> dict | None: try: return await self._call( "crd", "get_namespaced_custom_object", CORE_GROUP, VERSION, self.namespace, "sandboxes", name, ) except ApiException as exc: if exc.status == 404: return None raise async def patch_operating_mode(self, name: str, mode: str) -> None: await self._call( "crd", "patch_namespaced_custom_object", CORE_GROUP, VERSION, self.namespace, "sandboxes", name, {"spec": {"operatingMode": mode}}, ) # -- pods and logs ----------------------------------------------------- async def get_pod(self, name: str) -> dict | None: try: return await self._call( "core", "read_namespaced_pod", name, self.namespace ) except ApiException as exc: if exc.status == 404: return None raise async def read_pod_logs( self, name: str, container: str, tail_lines: int ) -> str: result = await self._call( "core", "read_namespaced_pod_log", name, self.namespace, container=container, tail_lines=tail_lines, ) return result or "" # --- Gateway --------------------------------------------------------------- @dataclass class GatewayState: """In-memory, per-process state (single replica by design).""" last_activity: dict[str, float] = field(default_factory=dict) in_flight: dict[str, int] = field(default_factory=dict) maintenance: dict[str, bool] = field(default_factory=dict) # (user, pod UID, browser Host) -> "dsh-auth-a=...; dsh-auth-b=..." handshake: dict[tuple[str, str, str], str] = field(default_factory=dict) lock: asyncio.Lock = field(default_factory=asyncio.Lock) class Gateway: """Route handlers plus the state machine they share.""" def __init__( self, settings: Settings, k8s, http_client: httpx.AsyncClient, ws_connect, ): self.settings = settings self.k8s = k8s self.http = http_client self.ws_connect = ws_connect self.state = GatewayState() # -- credentials ------------------------------------------------------- def credentials(self, headers, cookies) -> tuple[str | None, str | None]: """Return ``(token, cookie_user)``; bearer wins over the session cookie.""" auth = headers.get("authorization", "") if auth.lower().startswith("bearer "): return auth[7:].strip(), None raw = cookies.get(self.settings.cookie_name) if raw: user, sep, token = raw.partition(":") if sep and user and token: return token, user return None, None def _authorized( self, claim: dict, token: str | None, cookie_user: str | None, user: str, ) -> bool: if cookie_user is not None and cookie_user != user: return False # a session for another user never authorizes this one return token_matches(claim, token) @staticmethod def _mutation_from_cookie(headers, method: str, cookie_user: str | None) -> bool: return ( cookie_user is not None and method.upper() not in SAFE_METHODS and not same_origin(headers) ) def set_session(self, response: Response, user: str, token: str) -> None: response.set_cookie( self.settings.cookie_name, f"{user}:{token}", max_age=self.settings.cookie_max_age, path="/", httponly=True, samesite="lax", secure=self.settings.cookie_secure, ) # -- activity / maintenance ------------------------------------------- async def _touch(self, user: str) -> None: async with self.state.lock: self.state.last_activity[user] = time.time() async def _acquire(self, user: str) -> bool: """Reserve an in-flight slot, waiting out a concurrent suspend.""" deadline = time.monotonic() + self.settings.wake_timeout while True: async with self.state.lock: if not self.state.maintenance.get(user): self.state.in_flight[user] = ( self.state.in_flight.get(user, 0) + 1 ) self.state.last_activity[user] = time.time() return True if time.monotonic() > deadline: return False await asyncio.sleep(0.1) async def _release(self, user: str) -> None: async with self.state.lock: self.state.in_flight[user] = max( self.state.in_flight.get(user, 1) - 1, 0 ) self.state.last_activity[user] = time.time() # -- sandbox / handshake helpers -------------------------------------- async def _sandbox_of(self, claim: dict) -> dict | None: name = (claim.get("status") or {}).get("sandbox", {}).get("name") if not name: return None return await self.k8s.get_sandbox(name) async def wait_ready(self, user: str, timeout: int) -> dict | None: deadline = time.monotonic() + timeout while time.monotonic() < deadline: claim = await self.k8s.get_claim(user) sandbox = await self._sandbox_of(claim) if claim else None if sandbox is not None and derive_state(sandbox) == "Ready": return sandbox await asyncio.sleep(1) return None async def _resolve_sandbox(self, user: str, claim: dict) -> tuple[dict, str]: sandbox = await self._sandbox_of(claim) if sandbox is None: raise SandboxUnavailable("sandbox not provisioned yet") if derive_state(sandbox) != "Ready": # Wake-on-connect: safe in ANY state, including mid-suspension. self._clear_handshake(user) try: await self.k8s.patch_operating_mode( sandbox["metadata"]["name"], "Running" ) except Exception as exc: raise SandboxUnavailable("could not wake the sandbox") from exc sandbox = await self.wait_ready(user, self.settings.wake_timeout) if sandbox is None: raise SandboxUnavailable("DSH is waking up, retry shortly") await self._touch(user) fqdn = (sandbox.get("status") or {}).get("serviceFQDN") if not fqdn: raise SandboxUnavailable("sandbox has no service FQDN yet") return sandbox, fqdn def _pod_name(self, sandbox: dict) -> str: annotations = sandbox.get("metadata", {}).get("annotations") or {} return ( annotations.get(POD_NAME_ANNOTATION) or sandbox["metadata"]["name"] ) def _clear_handshake(self, user: str) -> None: for key in [k for k in self.state.handshake if k[0] == user]: self.state.handshake.pop(key, None) async def dsh_cookie( self, user: str, sandbox: dict, fqdn: str, host: str, force: bool = False, ) -> tuple[str, tuple[str, str, str], bool]: """Return ``(cookie, cache_key, was_cached)``, handshaking on a miss.""" pod_name = self._pod_name(sandbox) try: pod = await self.k8s.get_pod(pod_name) except ApiException as exc: raise HandshakeError(f"could not read pod {pod_name}: {exc}") from exc if pod is None: raise HandshakeError(f"pod {pod_name} not found") uid = (pod.get("metadata") or {}).get("uid") or pod_name key = (user, uid, host) if not force: cached = self.state.handshake.get(key) if cached: return cached, key, True try: logs = await self.k8s.read_pod_logs( pod_name, self.settings.dsh_container, self.settings.dsh_log_tail_lines, ) except ApiException as exc: raise HandshakeError( f"could not read logs for {pod_name}: {exc}" ) from exc launch_token = parse_dsh_launch_token(logs) if not launch_token: raise HandshakeError("DSH launch token not found in pod logs") url = ( f"http://{fqdn}:{self.settings.dsh_port}/" f"?token={quote(launch_token, safe='')}" ) headers = { "Host": host, "X-Forwarded-Host": host, "X-Forwarded-Proto": "http", } try: response = await self.http.get( url, headers=headers, follow_redirects=False, timeout=self.settings.http_timeout, ) except httpx.HTTPError as exc: raise HandshakeError(f"DSH handshake failed: {exc}") from exc if response.status_code >= 400: raise HandshakeError( f"DSH handshake returned {response.status_code}" ) cookie = extract_dsh_cookies(response.headers.get_list("set-cookie")) if not cookie: raise HandshakeError("DSH handshake set no dsh-auth cookie") self.state.handshake[key] = cookie return cookie, key, False # -- landing page / health -------------------------------------------- async def redirect_to_workspace(self, request: Request) -> Response: """``/u/`` -> ``/u//``; DSH's relative URLs need the slash.""" user = request.path_params["user"] if not valid_user(user): return self._error(404, "unknown user") target = f"/u/{user}/" if request.url.query: target = f"{target}?{request.url.query}" return RedirectResponse(target, status_code=308) async def healthz(self, request: Request) -> JSONResponse: return JSONResponse({"status": "ok"}) async def landing(self, request: Request) -> HTMLResponse: token, user = self.credentials(request.headers, request.cookies) session = "" if user and token: safe_user = html.escape(user) session = ( f'

Signed in as {safe_user}. ' f'Open workspace

' '
' '
' ) body = f""" DSH sandbox gateway

DeepSeek Harness sandboxes

{session}

Create a sandbox

Sign in with a token

Automation can use Authorization: Bearer <token> against the JSON API. The browser session is an HttpOnly cookie.

""" return HTMLResponse(body) # -- JSON helpers ------------------------------------------------------ @staticmethod def _error(status_code: int, message: str, headers=None) -> JSONResponse: return JSONResponse({"error": message}, status_code=status_code, headers=headers) @staticmethod def _wants_form(request: Request) -> bool: ctype = request.headers.get("content-type", "").lower() return ctype.startswith( ("application/x-www-form-urlencoded", "multipart/form-data") ) async def _read_payload(self, request: Request) -> dict: if request.headers.get("content-type", "").lower().startswith( "application/json" ): data = await self._read_body(request) try: payload = json.loads(data) except ValueError: payload = None return payload if isinstance(payload, dict) else {} form = await request.form() return {key: form.get(key) for key in form} async def _read_body(self, request: Request) -> bytes: limit = self.settings.max_body_bytes raw_length = request.headers.get("content-length") if raw_length: try: if int(raw_length) > limit: raise BodyTooLarge except ValueError: pass chunks, total = [], 0 async for chunk in request.stream(): total += len(chunk) if total > limit: raise BodyTooLarge chunks.append(chunk) return b"".join(chunks) # -- users ------------------------------------------------------------- async def signup(self, request: Request) -> Response: wants_form = self._wants_form(request) try: payload = await self._read_payload(request) except BodyTooLarge: return self._error(413, "request body too large") user = payload.get("user", "") if not valid_user(user): return self._error( 400, "body must be {'user': ''}" ) if await self.k8s.get_claim(user) is not None: return self._error(409, "user exists") token = secrets.token_urlsafe(32) token_hash = hashlib.sha256(token.encode()).hexdigest() try: await self.k8s.create_claim(user, token_hash, self.settings.pool) except UserExists: return self._error(409, "user exists") await self._touch(user) sandbox = await self.wait_ready(user, self.settings.wake_timeout) state = derive_state(sandbox) sandbox_name = (sandbox or {}).get("metadata", {}).get("name") if wants_form: response = RedirectResponse(f"/u/{user}/", status_code=303) self.set_session(response, user, token) return response response = JSONResponse( {"user": user, "token": token, "state": state, "sandbox": sandbox_name}, status_code=201 if state == "Ready" else 202, ) self.set_session(response, user, token) return response async def login(self, request: Request) -> Response: wants_form = self._wants_form(request) try: payload = await self._read_payload(request) except BodyTooLarge: return self._error(413, "request body too large") user = payload.get("user", "") token = payload.get("token", "") if not isinstance(user, str) or not isinstance(token, str): return self._error(400, "user and token are required") claim = await self.k8s.get_claim(user) if claim is None: return self._error(404, "not found") if not token_matches(claim, token): return self._error(401, "unauthorized") await self._touch(user) if wants_form: response = RedirectResponse(f"/u/{user}/", status_code=303) else: response = JSONResponse({"user": user, "state": "ok"}) self.set_session(response, user, token) return response async def logout(self, request: Request) -> Response: response = RedirectResponse("/", status_code=303) response.delete_cookie(self.settings.cookie_name, path="/") return response async def get_user(self, request: Request) -> Response: user = request.path_params["user"] if not valid_user(user): return self._error(404, "not found") claim = await self.k8s.get_claim(user) if claim is None: return self._error(404, "not found") token, cookie_user = self.credentials(request.headers, request.cookies) if not self._authorized(claim, token, cookie_user, user): return self._error(401, "unauthorized") sandbox = await self._sandbox_of(claim) return JSONResponse( { "user": user, "state": derive_state(sandbox), "sandbox": (sandbox or {}).get("metadata", {}).get("name"), } ) async def delete_user(self, request: Request) -> Response: user = request.path_params["user"] if not valid_user(user): return self._error(404, "not found") claim = await self.k8s.get_claim(user) if claim is None: return self._error(404, "not found") token, cookie_user = self.credentials(request.headers, request.cookies) if not self._authorized(claim, token, cookie_user, user): return self._error(401, "unauthorized") if self._mutation_from_cookie(request.headers, "DELETE", cookie_user): return self._error(403, "cross-origin request rejected") await self.k8s.delete_claim(user) self._clear_handshake(user) async with self.state.lock: self.state.last_activity.pop(user, None) self.state.in_flight.pop(user, None) return JSONResponse( {"user": user, "note": "claim deleted; sandbox and PVC cascade"} ) # -- HTTP proxy -------------------------------------------------------- async def proxy_http(self, request: Request) -> Response: user = request.path_params["user"] if not valid_user(user): return self._error(404, "unknown user") claim = await self.k8s.get_claim(user) if claim is None: return self._error(404, "unknown user") token, cookie_user = self.credentials(request.headers, request.cookies) if not self._authorized(claim, token, cookie_user, user): return self._error(401, "unauthorized") if self._mutation_from_cookie(request.headers, request.method, cookie_user): return self._error(403, "cross-origin request rejected") if not await self._acquire(user): return self._error( 503, "gateway busy, retry shortly", {"Retry-After": "5"} ) streaming = False try: sandbox, fqdn = await self._resolve_sandbox(user, claim) host = request.headers.get("host", "") cookie, _, was_cached = await self.dsh_cookie( user, sandbox, fqdn, host ) body = await self._read_body(request) upstream = await self._send_upstream( request, user, fqdn, cookie, body ) if upstream.status_code == 401 and was_cached: # The cached DSH cookie expired (pod restart, logout, ...): # handshake once more and replay the (already buffered) request. await upstream.aclose() cookie, _, _ = await self.dsh_cookie( user, sandbox, fqdn, host, force=True ) upstream = await self._send_upstream( request, user, fqdn, cookie, body ) response = StreamingResponse( self._relay(user, upstream), status_code=upstream.status_code, headers=self._response_headers(upstream, user), ) streaming = True return response except BodyTooLarge: return self._error(413, "request body too large") except (HandshakeError, SandboxUnavailable) as exc: return self._error( 503, f"{exc}; retry shortly", {"Retry-After": "10"}, ) except UpstreamError: return self._error(502, "upstream unavailable") finally: if not streaming: await self._release(user) async def _send_upstream( self, request: Request, user: str, fqdn: str, cookie: str, body: bytes, ) -> httpx.Response: path = request.path_params.get("path", "") url = f"http://{fqdn}:{self.settings.dsh_port}/{path}" query = filtered_query(request.url.query) if query: url = f"{url}?{query}" headers = { key: value for key, value in request.headers.items() if key.lower() not in OWNED_REQUEST_HEADERS and key.lower() != "content-length" } host = request.headers.get("host", "") headers["host"] = host headers["cookie"] = cookie # only the cached DSH auth cookie headers["x-forwarded-host"] = host headers["x-forwarded-proto"] = request.url.scheme if request.client: headers["x-forwarded-for"] = request.client.host else: headers["x-forwarded-for"] = "" port = host.rsplit(":", 1)[1] if ":" in host else "" headers["x-forwarded-port"] = port or ( "443" if request.url.scheme == "https" else "80" ) last_error: Exception | None = None for attempt in range(self.settings.upstream_attempts): try: upstream_request = self.http.build_request( request.method, url, headers=headers, content=body ) return await self.http.send(upstream_request, stream=True) except (httpx.ConnectError, httpx.ConnectTimeout) as exc: # Retry only failures that cannot have reached the upstream; # replaying a timed-out POST could duplicate a mutation. last_error = exc if attempt + 1 < self.settings.upstream_attempts: await asyncio.sleep(self.settings.upstream_retry_delay) except httpx.HTTPError as exc: raise UpstreamError("upstream request failed") from exc raise UpstreamError("upstream unavailable") from last_error @staticmethod def _response_headers(upstream: httpx.Response, user: str): headers = [] for key, value in upstream.headers.multi_items(): lowered = key.lower() if lowered in HOP_BY_HOP or lowered == "set-cookie": continue # never leak upstream cookies to the browser if lowered == "location": value = rewrite_location(value, user) headers.append((key, value)) return Headers(raw=[(k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in headers]) async def _relay(self, user: str, upstream: httpx.Response): try: async for chunk in upstream.aiter_raw(): yield chunk finally: with contextlib.suppress(Exception): await upstream.aclose() await self._release(user) # -- WebSocket proxy --------------------------------------------------- async def _ws_reject(self, websocket: WebSocket, code: int) -> None: # A close code can only be delivered after the upgrade is accepted; # before accept Starlette would turn this into an HTTP 403. with contextlib.suppress(Exception): await websocket.accept() with contextlib.suppress(Exception): await websocket.close(code=code) async def proxy_ws(self, websocket: WebSocket) -> None: user = websocket.path_params["user"] if not valid_user(user): return await self._ws_reject(websocket, 4401) claim = await self.k8s.get_claim(user) if claim is None: return await self._ws_reject(websocket, 4401) token, cookie_user = self.credentials(websocket.headers, websocket.cookies) if not self._authorized(claim, token, cookie_user, user): return await self._ws_reject(websocket, 4401) if cookie_user is not None and not same_origin(websocket.headers): return await self._ws_reject(websocket, 4401) if not await self._acquire(user): return await self._ws_reject(websocket, 1011) upstream_cm = None try: try: sandbox, fqdn = await self._resolve_sandbox(user, claim) host = websocket.headers.get("host", "") cookie, _, _ = await self.dsh_cookie(user, sandbox, fqdn, host) path = websocket.path_params.get("path", "") url = f"ws://{fqdn}:{self.settings.dsh_port}/{path}" query = filtered_query(websocket.url.query) if query: url = f"{url}?{query}" headers = { key: value for key, value in websocket.headers.items() if key.lower() not in OWNED_REQUEST_HEADERS and not key.lower().startswith("sec-websocket-") } headers["host"] = host headers["cookie"] = cookie headers["x-forwarded-host"] = host headers["x-forwarded-proto"] = websocket.url.scheme if websocket.client: headers["x-forwarded-for"] = websocket.client.host subprotocols = [ part.strip() for part in websocket.headers.get( "sec-websocket-protocol", "" ).split(",") if part.strip() ] upstream_cm = self.ws_connect( url, additional_headers=headers, subprotocols=subprotocols or None, open_timeout=self.settings.wake_timeout, max_size=None, proxy=None, ) upstream = await upstream_cm.__aenter__() except Exception: return await self._ws_reject(websocket, 1011) try: await websocket.accept( subprotocol=getattr(upstream, "subprotocol", None) ) except RuntimeError: pass await self._pump(websocket, upstream) with contextlib.suppress(Exception): await websocket.close( code=self._upstream_close_code(upstream) ) finally: if upstream_cm is not None: with contextlib.suppress(Exception): await upstream_cm.__aexit__(None, None, None) await self._release(user) @staticmethod def _upstream_close_code(upstream) -> int: code = getattr(upstream, "close_code", None) if isinstance(code, int) and 1000 <= code <= 4999 and code not in (1005, 1006): return code return 1000 async def _pump(self, websocket: WebSocket, upstream) -> None: async def client_to_upstream(): try: while True: message = await websocket.receive() if message["type"] == "websocket.disconnect": return if message.get("text") is not None: await upstream.send(message["text"]) elif message.get("bytes") is not None: await upstream.send(message["bytes"]) finally: with contextlib.suppress(Exception): await upstream.close() async def upstream_to_client(): async for message in upstream: if isinstance(message, str): await websocket.send_text(message) else: await websocket.send_bytes(message) tasks = [ asyncio.create_task(client_to_upstream()), asyncio.create_task(upstream_to_client()), ] done, pending = await asyncio.wait( tasks, return_when=asyncio.FIRST_COMPLETED ) for task in pending: task.cancel() await asyncio.gather(*done, *pending, return_exceptions=True) # -- idle sweeper ------------------------------------------------------ async def sweep_once(self) -> None: claims = await self.k8s.list_claims() for claim in claims: # Only suspend claims this gateway created. Everything else in # the namespace (resource-layer claims from the README walkthrough, # other controllers' objects) is not ours to manage. if not gateway_managed(claim): continue user = claim_user(claim) try: sandbox = await self._sandbox_of(claim) except ApiException: continue # transient/stale claim; retry next cycle if sandbox is None or derive_state(sandbox) != "Ready": continue async with self.state.lock: if self.state.maintenance.get(user): continue if self.state.in_flight.get(user, 0) > 0: continue # requests in flight => the user is active idle = time.time() - self.state.last_activity.setdefault( user, time.time() ) if idle < self.settings.idle_timeout: continue self.state.maintenance[user] = True try: print(f"suspending {user} (idle {int(idle)}s)", flush=True) await self.k8s.patch_operating_mode( sandbox["metadata"]["name"], "Suspended" ) self._clear_handshake(user) finally: async with self.state.lock: self.state.maintenance[user] = False async def sweeper_loop(self) -> None: while True: await asyncio.sleep(self.settings.sweep_interval) try: await self.sweep_once() except asyncio.CancelledError: raise except Exception as exc: print( f"sweeper: transient error, retrying next cycle: {exc}", flush=True, ) # --- Application factory --------------------------------------------------- def _lifespan(gateway: Gateway, close_http_client: httpx.AsyncClient | None = None): @contextlib.asynccontextmanager async def lifespan(app): task = asyncio.create_task(gateway.sweeper_loop()) try: yield finally: task.cancel() with contextlib.suppress(asyncio.CancelledError): await task if close_http_client is not None: with contextlib.suppress(Exception): await close_http_client.aclose() return lifespan def create_app( settings: Settings | None = None, k8s=None, http_client: httpx.AsyncClient | None = None, ws_connect=None, ) -> Starlette: """Build the ASGI app. All collaborators are injectable for tests.""" settings = settings or Settings.from_env() k8s = k8s if k8s is not None else Kubernetes(settings.namespace) owns_http_client = http_client is None if http_client is None: http_client = httpx.AsyncClient( follow_redirects=False, timeout=settings.http_timeout ) gateway = Gateway( settings, k8s, http_client, ws_connect or websockets.connect ) routes = [ Route("/healthz", gateway.healthz), Route("/", gateway.landing), Route("/users", gateway.signup, methods=["POST"]), Route("/login", gateway.login, methods=["POST"]), Route("/logout", gateway.logout, methods=["POST"]), Route("/users/{user}", gateway.get_user), Route("/users/{user}", gateway.delete_user, methods=["DELETE"]), Route("/u/{user}", gateway.redirect_to_workspace), Route( "/u/{user}/{path:path}", gateway.proxy_http, methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"], ), WebSocketRoute("/u/{user}/{path:path}", gateway.proxy_ws), ] app = Starlette( routes=routes, lifespan=_lifespan(gateway, http_client if owns_http_client else None), ) app.state.gateway = gateway return app