ALL

Solo para leer y verificar. No puedes copiar, usar, modificar ni redistribuir este código. Aviso completo en la página principal.

src/all_mcp/surfaces.py

SHA-256: edfd470d4b4b4c3903946f963ef7e5bb3df8e5d46cb3e12056633a502280369d
865 líneas · 47569 bytes · Ver en GitHub

1"""The agent's surfaces (ALL 2.0, end-to-end encryption): who may reach what is decided by WHICH SOCKET a request came in on, never by what it says.
2
3One inner app (the MCP stack, OAuth, the pause and tray routes) is served by up to three uvicorn servers in one event loop:
4
5 * LOCAL 127.0.0.1:4410 (ALL_MCP_PORT). The ALL icon, the command line, the installer and the QA scripts (and the legacy relay's client while
6 the lane is off). Nothing about the request changes anything: the bootstrap bearer works, the routes of the icon exist, and the paywall gate is not
7 here (the old relay gates for the computers it still serves, because it can still read the request). It is the ONLY exempt surface.
8 * E2E 127.0.0.1:<ephemeral port>, written to state_dir()/surfaces.json (0600, with the agent's pid and a secret that is new at every start). Only
9 the computer's own TLS terminator dials it, with what ChatGPT sent. So everything that arrives here is remote, whatever it claims to be:
10 it is closed by default (http only, one Host, an exact list of paths, size caps, no forwarding headers, no bootstrap bearer), the trial's
11 paywall and the activation funnel always apply, and they fail closed (no answer from the control plane, or no device id: blocked).
12 * RELAY 127.0.0.1:<ephemeral port>, written to state_dir()/relay-lane.json (0600, same shape as surfaces.json). Only the legacy relay's client dials it (ALL 2.0
13 WP11; `ALL_MCP_RELAY_LANE`, on unless it says off). It carries what the old relay forwards, which is remote traffic that the relay terminated and can read:
14 so it gets the relay's own allowlist (the paths of its DEVICE_PATH), one Host, size caps, no forwarding headers and no bootstrap bearer, and nothing else
15 of LOCAL (no pause, tray, permissions or pairing-start route, no /public). The paywall gate is NOT here: the relay gates what it serves.
16
17Every wrapper writes the surface into `scope["state"]` (the listener marks LOCAL, E2E and RELAY; no header can: `x-all-surface: local` is stripped on E2E and RELAY and
18read nowhere). The authentication (middleware.AuthMiddleware) refuses the bootstrap bearer on E2E and RELAY like an unknown token, and the gate
19(middleware.DeviceGateMiddleware) enforces on E2E and passes the others. A request that carries no mark (a test, a developer's `build_app()`) is LOCAL-like: only
20the wrappers below put traffic in front of the app in production.
21
22Why the identity proof: the terminator reads where to dial from surfaces.json, and a port that the agent freed (it crashed, it restarted) can be bound by
23another process, which would then receive ChatGPT's tokens in clear. So, before the terminator forwards the first ChatGPT byte, it sends on that same
24connection a reserved GET with an HMAC challenge built from the per-start secret; only the agent holds the secret and answers
25(`prove_listener`, the terminator's half); anything else, a 404 or silence, is refused and nothing of ChatGPT's is sent.
26
27The relay lane has the same identity proof (relay_client.RelayClient, the dialer, runs `prove_listener` against it before it trusts relay-lane.json).
28
29The E2E listener is started lazily and supervised (SurfaceRunner): only when the rollback switch ALL_MCP_E2E is not off, ALL_MCP_REMOTE_SURFACE is not 0 and
30runtime.env names an enrolled label; inside its own exception boundary, so a failure there (it cannot bind, it crashes) leaves the LOCAL surface serving; its
31lifespan is off (the inner app's lifespan runs once, under LOCAL); and uvicorn runs on both with `proxy_headers` off, no trusted forwarders, a keep-alive and a
32concurrency ceiling, and an h11 head cap (X-Forwarded-For is never believed: the client of a request is its socket's peer).
33"""
34from __future__ import annotations
35
36import asyncio
37import base64
38import contextlib
39import contextvars
40import hashlib
41import hmac
42import json
43import logging
44import os
45import re
46import secrets
47import socket
48import stat
49import time
50from dataclasses import dataclass
51from pathlib import Path
52from typing import Any, Callable
53
54from . import e2e_policy as policy
55from . import state_io
56from .config import runtime_values, state_dir
57
58_LOG = logging.getLogger(__name__)
59
60LOOPBACK = "127.0.0.1"
61SURFACES_FILE = "surfaces.json"
62LANE_FILE = "relay-lane.json"
63# What a wrapper writes into scope["state"], and the two values it can have.
64SURFACE_KEY = "all_mcp_surface"
65LOCAL = "local"
66E2E = "e2e"
67RELAY = "relay"
68# Both surfaces that carry remote traffic: what the authentication must never treat as the computer's own tools.
69REMOTE = frozenset({E2E, RELAY})
70# The E2E wrapper also writes the name it was started for (the device's own, from runtime.env): the origin of an E2E request is the one of that name, so the issuer,
71# the resource and the OAuth store follow the listener and never the Host header a caller wrote (origins.OriginSet.for_surface).
72E2E_HOST_KEY = "all_mcp_surface_host"
73
74# The rollback switch (one name; ALL_E2E is only an alias of the developer's dev runner) and the switch that keeps the E2E listener from
75# starting at all (the update candidate probe sets it so that the probe is exactly the 1.1 probe).
76SWITCH = "ALL_MCP_E2E"
77REMOTE_SURFACE = "ALL_MCP_REMOTE_SURFACE"
78# The hardened lane of the legacy relay (WP11). ON unless it says off; ALL_MCP_REMOTE_SURFACE=0 (the update candidate probe) also keeps it from starting.
79RELAY_LANE = "ALL_MCP_RELAY_LANE"
80# The old route's own socket (the relay client's WebSocket to the relay). ON whenever the variable is absent or empty: a 1.0 or 1.1 runtime.env that has no encrypted-path
81# keys, updated to 2.0, stays connected and answering. Only an explicit off value stops it (a later 2.x, on the owner's order, or "Cerrar la ruta antigua" on one computer).
82LEGACY_RELAY = "ALL_MCP_LEGACY_RELAY"
83# The canary device's mode, "certificate manager only" (the tunnel runbook, section 8): the e2e child keeps this computer's certificate issued and renewed as usual but opens no tunnel
84# connection of its own, because the canary device app holds the connection to the nodes under the same label (two clients of one label evict each other for ever). The opposite
85# default of the rollback switches: it is ON only when it says on, so a computer that never set it is an ordinary one. CANARY_STATUS is the canary's own --status-file, where the child
86# reads whether its links are up (e2e_canary.py). Both are kept by every rewrite of runtime.env (config.E2E_RUNTIME_KEYS) and neither is ever written by an install or an update.
87CERT_ONLY = "ALL_E2E_CERT_ONLY"
88CANARY_STATUS = "ALL_E2E_CANARY_STATUS"
89_OFF = frozenset({"0", "off", "false", "no"})
90_ON = frozenset({"1", "on", "true", "yes"})
91# The runtime.env file, written by the installer when the computer enrolls: the label and the host the control plane gave it (and, optionally, the zone it must be in).
92LABEL_KEY, HOST_KEY, ZONE_KEY = "ALL_E2E_LABEL", "ALL_E2E_HOST", "ALL_E2E_ZONE"
93
94# What the E2E surface answers, by exact path. Everything else is a 404: the pause, resume, pairing start, tray,
95# permissions and public routes are not here even when the app has them, and `/mcp/` is served as `/mcp` (no redirect, ChatGPT keeps the slash it was given).
96E2E_PATHS = frozenset({
97 "/mcp", "/mcp/", "/register", "/authorize", "/token", "/pair", "/healthz",
98 "/.well-known/oauth-protected-resource", "/.well-known/oauth-protected-resource/mcp",
99 "/.well-known/oauth-authorization-server", "/.well-known/openid-configuration",
100})
101# What the lane answers: the paths of the relay's DEVICE_PATH (the relay index module), which is what the old relay forwards; the relay sends "/" + the path after
102# /d/<device>, so a bare /d/<device> and /d/<device>/ arrive as "/". The lane is the E2E list plus "/": no pause, tray, permissions, pairing-start or /public path is on
103# it, even though the app has them. The relay lane test suite reads the relay's regex and checks both lists against each other.
104LANE_PATHS = E2E_PATHS | {"/"}
105# Headers that say who the caller is or how it got here: a remote caller writes them as it likes, so none of them reaches the app.
106_STRIPPED_PREFIXES = (b"x-forwarded-", b"cf-", b"x-all-")
107_STRIPPED_NAMES = frozenset({b"forwarded"})
108
109# Ceilings of the two uvicorn servers that are not in the policy table: LOCAL carries the legacy relay's one connection per request, so it is far more
110# generous than E2E (which gets twice the terminator's stream cap: its streams plus the identity probes in flight).
111LOCAL_CONCURRENCY = 512
112E2E_CONCURRENCY_FACTOR = 2
113GRACEFUL_SHUTDOWN_SECONDS = 10
114START_TIMEOUT_SECONDS = 10.0
115WATCH_SECONDS = 5.0 # how often the runner looks again at whether the E2E listener should be running
116
117_current: contextvars.ContextVar[str | None] = contextvars.ContextVar("all_mcp_surface", default=None)
118_serving = {"e2e": False}
119
120
121def current() -> str | None:
122 """The surface of the request being served in this task (E2E, LOCAL), or None outside a request: what the funnel asks, because it is called deep inside
123 the OAuth route and the gate, far from the ASGI scope."""
124 return _current.get()
125
126
127def in_e2e() -> bool:
128 return _current.get() == E2E
129
130
131def e2e_serving() -> bool:
132 """True while this process has an E2E listener up: what keeps the paywall verdict fresh (device_access.start_refresher)."""
133 return _serving["e2e"]
134
135
136def surface_of(scope: dict) -> str | None:
137 return (scope.get("state") or {}).get(SURFACE_KEY)
138
139
140# ------------------------------------------------------------------------------------------------------------------------ is the E2E listener wanted?
141
142def _setting(values: dict[str, str], name: str) -> str:
143 return (os.environ.get(name) or values.get(name, "")).strip().lower()
144
145
146@dataclass(frozen=True)
147class E2EConfig:
148 label: str
149 host: str # the device's full name, exactly what a caller must write in Host
150
151
152def enrolled_host(values: dict[str, str] | None = None) -> E2EConfig | None:
153 """The label and host this computer enrolled with, when runtime.env holds a well-formed pair (label as the policy writes it, host one label under a zone,
154 and under ALL_E2E_ZONE when that is set), else None. A computer without them has no encrypted name, so nothing listens for one."""
155 values = runtime_values() if values is None else values
156 label, host, zone = values.get(LABEL_KEY, ""), values.get(HOST_KEY, ""), values.get(ZONE_KEY, "")
157 if (not policy.anchored("identity.label_regex").fullmatch(label) or len(host) > policy.value("identity.fqdn_max_length")
158 or not policy.anchored("identity.fqdn_regex").fullmatch(host) or not host.startswith(label + ".")):
159 return None
160 if zone and not policy.host_pattern(zone).match(host):
161 return None
162 return E2EConfig(label, host)
163
164
165def e2e_config(values: dict[str, str] | None = None) -> E2EConfig | None:
166 """What the E2E listener should serve right now, or None when it should not run: the rollback switch is off, the remote surface is off, or the runtime has
167 no enrolled label. The ONE rollback variable of the migration is ALL_MCP_E2E: it stops the listener (and the tunnel, whose own code reads the same name);
168 it never makes the paywall easier, because a computer that has no E2E surface has nothing the paywall would guard."""
169 values = runtime_values() if values is None else values
170 return None if e2e_switched_off(values) else enrolled_host(values)
171
172
173def e2e_switched_off(values: dict[str, str] | None = None) -> bool:
174 """True while the rollback switch (ALL_MCP_E2E) or the remote surface (ALL_MCP_REMOTE_SURFACE) is off: no E2E listener runs, so nothing that hands out the encrypted address
175 (the reconnect flow, the tray's status) may either."""
176 values = runtime_values() if values is None else values
177 return _setting(values, SWITCH) in _OFF or _setting(values, REMOTE_SURFACE) in _OFF
178
179
180def relay_lane_wanted(values: dict[str, str] | None = None) -> bool:
181 """Whether the agent serves the relay lane and the relay's client dials it: ALL_MCP_RELAY_LANE is not off (it is ON when absent) and ALL_MCP_REMOTE_SURFACE is not
182 off. The agent and the legacy relay's client both ask THIS, over the same runtime.env, so they never disagree (a lane the agent does not serve is a lane the client
183 does not wait for). Off leaves the legacy relay on the LOCAL surface exactly as 1.1 had it."""
184 values = runtime_values() if values is None else values
185 return _setting(values, RELAY_LANE) not in _OFF and _setting(values, REMOTE_SURFACE) not in _OFF
186
187
188def legacy_relay_wanted(values: dict[str, str] | None = None) -> bool:
189 """Whether the relay client keeps the old route's WebSocket: ALL_MCP_LEGACY_RELAY is not an explicit off value. Absent, empty or anything else means ON, so that the
190 absence of the 2.0 keys can never switch a computer's only working route off."""
191 values = runtime_values() if values is None else values
192 return _setting(values, LEGACY_RELAY) not in _OFF
193
194
195def cert_only(values: dict[str, str] | None = None) -> bool:
196 """Whether this computer is a canary's: ALL_E2E_CERT_ONLY says on (1, on, true or yes; absent, empty or anything else is off). The e2e child then runs the enroll, the certificate
197 manager and the self-check of the public name exactly as always, and starts no tunnel (e2e_runner.ChildStages). The rollback switch still wins: ALL_MCP_E2E=off stops all of it."""
198 values = runtime_values() if values is None else values
199 return _setting(values, CERT_ONLY) in _ON
200
201
202# ------------------------------------------------------------------------------------------------------------------------ the transport security of the inner app
203
204_transport_settings: list[Any] = []
205
206
207def adopt_transport_security(settings: Any, values: dict[str, str] | None = None) -> None:
208 """The inner app's DNS-rebinding settings (mcp TransportSecuritySettings, built by mcp_server.build_app) learn the E2E host: the Host the E2E surface lets
209 through must also be one the MCP transport accepts. Called when the app is built, and again by the runner when the listener starts (a computer that
210 enrolls while the agent runs)."""
211 _transport_settings.append(settings)
212 found = enrolled_host(values)
213 if found is not None:
214 _allow_host(settings, found.host)
215
216
217def _allow_host(settings: Any, host: str) -> None:
218 for field, entry in (("allowed_hosts", host), ("allowed_origins", f"https://{host}")):
219 listed = getattr(settings, field)
220 if entry not in listed:
221 listed.append(entry)
222
223
224def _allow_everywhere(host: str) -> None:
225 for settings in list(_transport_settings):
226 _allow_host(settings, host)
227
228
229# ------------------------------------------------------------------------------------------------------------------------ the identity proof (T14)
230
231def b64u(raw: bytes) -> str:
232 return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")
233
234
235def _key(secret: str) -> bytes:
236 raw = base64.urlsafe_b64decode(secret + "=" * (-len(secret) % 4))
237 if len(raw) != 32 or not policy.anchored("identity.attach_secret_regex").fullmatch(secret):
238 raise ValueError("surface secret")
239 return raw
240
241
242def new_secret() -> str:
243 return b64u(secrets.token_bytes(32))
244
245
246def challenge_mac(secret: str, nonce: str) -> str:
247 """What the terminator puts on its challenge: only a holder of the per-start secret can answer it."""
248 return b64u(hmac.new(_key(secret), b"all-surface-challenge-v1\n" + nonce.encode("ascii"), hashlib.sha256).digest())
249
250
251def surface_proof(secret: str, nonce: str, pid: int, port: int) -> str:
252 """What the listener answers: it proves it holds the secret of the process surfaces.json names, for THIS nonce, pid and port."""
253 message = b"all-surface-proof-v1\n" + nonce.encode("ascii") + b"\n" + str(pid).encode() + b"\n" + str(port).encode()
254 return b64u(hmac.new(_key(secret), message, hashlib.sha256).digest())
255
256
257def verify_challenge(secret: str, nonce: Any, mac: Any) -> bool:
258 """The agent's check of a challenge: a nonce the policy allows, and the MAC only a holder of the secret can make (compared in constant time)."""
259 try:
260 return (isinstance(nonce, str) and policy.anchored("identity.nonce_regex").fullmatch(nonce) is not None and isinstance(mac, str)
261 and hmac.compare_digest(challenge_mac(secret, nonce).encode(), mac.encode("ascii")))
262 except (ValueError, UnicodeError):
263 return False
264
265
266def verify_proof(secret: str, nonce: Any, pid: int, port: int, proof: Any) -> bool:
267 """The terminator's check of an answer: bound to this nonce, this pid and this port, so a proof for another challenge, process or port fails."""
268 try:
269 return (isinstance(nonce, str) and policy.anchored("identity.nonce_regex").fullmatch(nonce) is not None and isinstance(proof, str)
270 and hmac.compare_digest(surface_proof(secret, nonce, pid, port).encode(), proof.encode("ascii")))
271 except (ValueError, UnicodeError):
272 return False
273
274
275def identity_request(host: str, port: int, nonce: str, mac: str) -> str:
276 """The exact request head the terminator sends on the connection it is about to hand to ChatGPT, before the first ChatGPT byte."""
277 return (f"GET {policy.value('surface.identity_path')} HTTP/1.1\r\nHost: {host}:{port}\r\n{policy.value('surface.identity_challenge_header')}: {nonce}\r\n"
278 f"{policy.value('surface.identity_mac_header')}: {mac}\r\nConnection: keep-alive\r\n\r\n")
279
280
281class SurfaceIdentityError(Exception):
282 """The listener could not prove it is the agent. `word` is a fixed word (never text from the network)."""
283
284 def __init__(self, word: str):
285 super().__init__(word)
286 self.word = word
287
288
289async def prove_listener(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, *, host: str, port: int, secret: str, pid: int,
290 timeout: float = 5.0) -> None:
291 """The terminator's half of the proof: on a freshly opened connection to the surface, challenge it and check its answer; return only when it is the agent
292 that holds `secret` for `pid` and `port`, with the answer consumed (the connection is kept alive for ChatGPT's bytes). Raises SurfaceIdentityError, and
293 the caller then closes the connection and sends nothing of ChatGPT's."""
294 nonce = secrets.token_urlsafe(policy.value("protocol.nonce_bytes"))
295 try:
296 writer.write(identity_request(host, port, nonce, challenge_mac(secret, nonce)).encode("ascii"))
297 await asyncio.wait_for(writer.drain(), timeout)
298 head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout)
299 except (OSError, asyncio.TimeoutError, asyncio.IncompleteReadError, asyncio.LimitOverrunError, ValueError):
300 raise SurfaceIdentityError("no_answer") from None
301 lines = head[:-4].split(b"\r\n")
302 if lines[0] != b"HTTP/1.1 200 OK" or len(head) > 2048:
303 raise SurfaceIdentityError("not_the_agent")
304 seen: dict[bytes, bytes] = {}
305 for line in lines[1:]:
306 name, _, value = line.partition(b":")
307 if name.lower() in seen:
308 raise SurfaceIdentityError("not_the_agent")
309 seen[name.lower()] = value.strip()
310 if seen.get(b"content-length") != b"0":
311 raise SurfaceIdentityError("not_the_agent")
312 if not verify_proof(secret, nonce, pid, port, seen.get(policy.value("surface.identity_proof_header").encode(), b"").decode("latin-1")):
313 raise SurfaceIdentityError("bad_proof")
314
315
316# ------------------------------------------------------------------------------------------------------------------------ surfaces.json
317
318@dataclass(frozen=True)
319class SurfaceTarget:
320 pid: int
321 host: str
322 port: int
323 secret: str
324 started_at: float
325
326
327def surfaces_path() -> Path:
328 return state_dir() / SURFACES_FILE
329
330
331def lane_path() -> Path:
332 return state_dir() / LANE_FILE
333
334
335def write_surfaces_file(*, pid: int, port: int, secret: str, path: Path | None = None) -> None:
336 """state_dir()/surfaces.json (0600: the secret in it is what proves the listener is the agent). Replaces whatever an earlier start left."""
337 state_io.atomic_json(path or surfaces_path(), {"v": 1, "pid": pid, "host": LOOPBACK, "port": port, "secret": secret, "started_at": int(time.time())})
338
339
340def remove_surfaces_file(*, pid: int, secret: str, path: Path | None = None) -> None:
341 """Remove surfaces.json, but only while it still names this listener (a newer start may have written its own)."""
342 path = path or surfaces_path()
343 try:
344 data = json.loads(path.read_text(encoding="utf-8"))
345 if data.get("pid") == pid and data.get("secret") == secret:
346 path.unlink()
347 except (OSError, ValueError, AttributeError):
348 pass
349
350
351def _pid_alive(pid: int) -> bool:
352 try:
353 import psutil
354 return bool(psutil.pid_exists(pid))
355 except ImportError:
356 pass
357 try:
358 os.kill(pid, 0)
359 except ProcessLookupError:
360 return False
361 except OSError:
362 return True
363 return True
364
365
366def load_target(path: Path | None = None, *, require_alive: bool = True) -> SurfaceTarget:
367 """What the terminator dials: the one validated address of the E2E listener. Raises SurfaceIdentityError("no_surface") for anything that is not a
368 well-formed surfaces.json of a live agent: a missing or oversized file, a link, one other accounts can read, a bad field, a pid that is gone."""
369 path = path or surfaces_path()
370 try:
371 info = path.lstat()
372 if not stat.S_ISREG(info.st_mode) or info.st_size > 4096 or (os.name != "nt" and info.st_mode & 0o077):
373 raise SurfaceIdentityError("no_surface")
374 data = json.loads(path.read_text(encoding="utf-8"))
375 except (OSError, ValueError):
376 raise SurfaceIdentityError("no_surface") from None
377 try:
378 pid, port, started = data["pid"], data["port"], data["started_at"]
379 ok = (isinstance(data, dict) and set(data) == {"v", "pid", "host", "port", "secret", "started_at"} and data["v"] == 1 and data["host"] == LOOPBACK
380 and type(pid) is int and 1 <= pid <= 0xFFFFFFFF and type(port) is int and 1024 <= port <= 65535
381 and isinstance(data["secret"], str) and policy.anchored("identity.attach_secret_regex").fullmatch(data["secret"]) is not None
382 and isinstance(started, (int, float)) and not isinstance(started, bool))
383 except (KeyError, TypeError):
384 ok = False
385 if not ok or (require_alive and not _pid_alive(pid)):
386 raise SurfaceIdentityError("no_surface")
387 return SurfaceTarget(pid, LOOPBACK, port, data["secret"], float(started))
388
389
390# ------------------------------------------------------------------------------------------------------------------------ the wrappers
391
392def _marked(scope: dict, surface: str, **changes: Any) -> dict:
393 state = {**(scope.get("state") or {}), SURFACE_KEY: surface}
394 state.pop(E2E_HOST_KEY, None) # only the E2E wrapper writes it, below: whatever came in with the scope is not believed
395 if "host" in changes:
396 state[E2E_HOST_KEY] = changes.pop("host")
397 return {**scope, **changes, "state": state}
398
399
400class LocalSurface:
401 """The LOCAL listener's wrapper: marks every request LOCAL, otherwise changes nothing (the lifespan, which the inner app runs once, passes through)."""
402
403 def __init__(self, app):
404 self.app = app
405
406 async def __call__(self, scope, receive, send):
407 if scope["type"] == "lifespan":
408 await self.app(scope, receive, send)
409 return
410 token = _current.set(LOCAL)
411 try:
412 await self.app(_marked(scope, LOCAL), receive, send)
413 finally:
414 _current.reset(token)
415
416
417async def _respond(send, status: int, body: bytes = b"", *, extra=(), head: bool = False) -> None:
418 headers = [(b"cache-control", b"no-store"), (b"content-length", str(len(body)).encode())]
419 if body:
420 headers.append((b"content-type", b"application/json"))
421 headers.extend(extra)
422 await send({"type": "http.response.start", "status": status, "headers": headers})
423 await send({"type": "http.response.body", "body": b"" if head else body})
424
425
426def _values(scope, name: bytes) -> list[bytes]:
427 return [value for key, value in scope.get("headers", []) if key == name]
428
429
430class E2ESurface:
431 """The E2E listener's wrapper: everything that arrives here is remote, so it is closed by default and the only things let through are what ChatGPT needs.
432 (RelaySurface, below, is the same wrapper for the legacy relay's lane: its own mark, its own list of paths, and the Host of its own socket.)
433
434 In order: a lifespan or websocket scope never reaches the app; the head is bounded (headers, their size, the whole head, one body length, no Upgrade); the identity
435 proof answers its reserved path and nothing else does; the request names this computer, once, by its exact name (421); the path is on the list (404, and
436 `/mcp/` is served as `/mcp`); `/healthz` says only that the agent is up; the headers that claim an origin are removed; the body is capped while it streams
437 (413); and the request is marked E2E for the authentication and the gate. `host` is the device's full name; `secret`, `pid` and `port` are what surfaces.json
438 holds."""
439
440 surface = E2E
441 paths = E2E_PATHS
442 _answers_healthz = True # E2E says only that the agent is up; the relay lane lets the app answer, as the old relay always forwarded it
443
444 def __init__(self, app, *, host: str, secret: str, pid: int, port: int):
445 self.app = app
446 self.host = host.lower()
447 self._secret = secret
448 self._pid = pid
449 self._port = port
450 _key(secret) # a secret that cannot key the proof is a start that must fail now, not at the first challenge
451 self._body_max = policy.value("surface.body_max_bytes")
452 self._headers_max = policy.value("surface.headers_max")
453 self._value_max = policy.value("surface.header_value_max_bytes")
454 self._head_max = head_cap()
455 self._proof_path = policy.value("surface.identity_path")
456
457 async def __call__(self, scope, receive, send):
458 kind = scope["type"]
459 if kind == "lifespan": # E2E runs with the lifespan off; if one ever arrives it is answered here, so the inner app's lifespan never runs twice
460 while True:
461 message = await receive()
462 if message["type"] == "lifespan.startup":
463 await send({"type": "lifespan.startup.complete"})
464 elif message["type"] == "lifespan.shutdown":
465 await send({"type": "lifespan.shutdown.complete"})
466 return
467 if kind != "http": # a websocket (or anything else): the E2E surface has no scope but http, so nothing skips the authentication
468 if kind == "websocket":
469 await receive()
470 await send({"type": "websocket.close", "code": 1008})
471 return
472 headers = scope.get("headers", [])
473 if (len(headers) > self._headers_max or any(len(value) > self._value_max for _, value in headers)
474 or len(scope.get("raw_path") or b"") + len(scope.get("query_string") or b"") + sum(len(key) + len(value) + 4 for key, value in headers) > self._head_max):
475 await _respond(send, 431, b'{"error":"request_header_fields_too_large"}', extra=[(b"connection", b"close")])
476 return
477 if _values(scope, b"upgrade"): # no websocket and no h2c upgrade: a request that asks for one is refused, not served as a plain one
478 await _respond(send, 400, b'{"error":"upgrade_not_supported"}', extra=[(b"connection", b"close")])
479 return
480 lengths = _values(scope, b"content-length")
481 if len(lengths) > 1 or (lengths and not lengths[0].isdigit()) or (lengths and _values(scope, b"transfer-encoding")):
482 await _respond(send, 400, b'{"error":"invalid_request"}', extra=[(b"connection", b"close")]) # the framing a smuggler plays with
483 return
484 path = scope.get("path", "")
485 if path == self._proof_path:
486 await self._prove(scope, send)
487 return
488 hosts = _values(scope, b"host")
489 if len(hosts) != 1 or hosts[0].lower() != self.host.encode("latin-1"):
490 await _respond(send, 421, b'{"error":"misdirected_request"}')
491 return
492 if path not in self.paths:
493 await _respond(send, 404, b'{"error":"not_found"}')
494 return
495 if path == "/healthz" and self._answers_healthz:
496 if scope["method"] not in ("GET", "HEAD"):
497 await _respond(send, 405, b'{"error":"method_not_allowed"}', extra=[(b"allow", b"GET, HEAD")])
498 return
499 await _respond(send, 200, b'{"ok":true}', head=scope["method"] == "HEAD")
500 return
501 if lengths and int(lengths[0]) > self._body_max:
502 await _respond(send, 413, b'{"error":"request_too_large"}', extra=[(b"connection", b"close")])
503 return
504 changes: dict[str, Any] = {"headers": [(key, value) for key, value in headers
505 if not key.startswith(_STRIPPED_PREFIXES) and key not in _STRIPPED_NAMES]}
506 if path == "/mcp/":
507 changes["path"], changes["raw_path"] = "/mcp", b"/mcp"
508 if self.surface == E2E: # only the E2E listener names an encrypted origin; the relay lane's requests belong to the primary one
509 changes["host"] = self.host
510 await self._serve(_marked(scope, self.surface, **changes), receive, send)
511
512 async def _prove(self, scope, send) -> None:
513 """The reserved identity path: answered only to a valid challenge (a 404 like any other path otherwise, so a stranger learns nothing)."""
514 nonces = _values(scope, policy.value("surface.identity_challenge_header").encode())
515 macs = _values(scope, policy.value("surface.identity_mac_header").encode())
516 nonce = nonces[0].decode("latin-1") if len(nonces) == 1 and len(macs) == 1 else ""
517 if scope["method"] != "GET" or not verify_challenge(self._secret, nonce, macs[0].decode("latin-1") if macs else ""):
518 await _respond(send, 404, b'{"error":"not_found"}')
519 return
520 proof = surface_proof(self._secret, nonce, self._pid, self._port).encode()
521 await _respond(send, 200, b"", extra=[(policy.value("surface.identity_proof_header").encode(), proof)])
522
523 async def _serve(self, scope, receive, send) -> None:
524 total = 0
525 exceeded = False
526 started = False
527
528 async def capped_receive():
529 nonlocal total, exceeded
530 if exceeded:
531 return {"type": "http.disconnect"}
532 message = await receive()
533 if message["type"] == "http.request":
534 total += len(message.get("body", b""))
535 if total > self._body_max: # counted as it streams, whatever Content-Length said (or did not say)
536 exceeded = True
537 return {"type": "http.disconnect"}
538 return message
539
540 async def guarded_send(message):
541 nonlocal started
542 if exceeded:
543 return
544 if message["type"] == "http.response.start":
545 started = True
546 await send(message)
547
548 token = _current.set(self.surface)
549 try:
550 await self.app(scope, capped_receive, guarded_send)
551 except Exception:
552 if not exceeded:
553 raise
554 finally:
555 _current.reset(token)
556 if exceeded and not started:
557 await _respond(send, 413, b'{"error":"request_too_large"}', extra=[(b"connection", b"close")])
558
559
560class RelaySurface(E2ESurface):
561 """The relay lane's wrapper (WP11): what the legacy relay's client dials. It is the E2E wrapper with the relay's own allowlist (`LANE_PATHS`), the mark RELAY, and one
562 Host: the lane's own socket, `127.0.0.1:<port>`, the only name the client writes (the relay's Host is dropped and the client's HTTP library sets this one); any other
563 Host is a request that did not come from the client (421). The app answers `/healthz` itself here: the installers and the site read the old relay's health through it,
564 and that report is the same one 1.1 gave. What this surface never carries, wherever a request says it came from: the routes of the computer's own tools, the
565 bootstrap bearer (middleware.AuthMiddleware), the paywall gate (the relay gates what it serves) and a websocket."""
566
567 surface = RELAY
568 paths = LANE_PATHS
569 _answers_healthz = False
570
571 def __init__(self, app, *, secret: str, pid: int, port: int):
572 super().__init__(app, host=f"{LOOPBACK}:{port}", secret=secret, pid=pid, port=port)
573
574
575# ------------------------------------------------------------------------------------------------------------------------ the listeners
576
577def head_cap() -> int:
578 """The most a request head may weigh: one header as long as the policy allows, plus an ordinary head. h11 enforces it only on a head that arrives incomplete (a
579 head that comes whole in one read is parsed whatever its size), so the E2E wrapper counts it again."""
580 return policy.value("surface.header_value_max_bytes") + policy.value("terminator.request_head_max_bytes")
581
582
583def uvicorn_config(app, *, log_level: str, concurrency: int, ws: str = "auto", lifespan: str = "auto", **options: Any):
584 """The settings both servers share (uvicorn is imported here so that importing this module stays cheap). `proxy_headers` is OFF and no address is trusted to
585 forward: uvicorn would otherwise believe X-Forwarded-For from 127.0.0.1, which is every caller there is (the tunnel and the terminator are on loopback), and the
586 rate-limit key and the client of a request would be whatever the remote caller wrote. Access logs are off (the request line carries OAuth and pairing
587 parameters)."""
588 import uvicorn
589 return uvicorn.Config(app, interface="asgi3", http="h11", log_level=log_level, access_log=False, proxy_headers=False, forwarded_allow_ips="", ws=ws, lifespan=lifespan,
590 timeout_keep_alive=policy.value("surface.keep_alive_s"), limit_concurrency=concurrency,
591 h11_max_incomplete_event_size=head_cap(), **options)
592
593
594def _local_server_class():
595 import uvicorn
596
597 class LocalServer(uvicorn.Server):
598 """The LOCAL server. When it starts to shut down (a signal) it first tells the runner, so the E2E listener and surfaces.json go before the process does:
599 uvicorn raises the signal again as soon as `serve` is done, which ends the process before anything after it could run."""
600
601 def __init__(self, config, before_shutdown: Callable[[], Any]):
602 super().__init__(config)
603 self._before_shutdown = before_shutdown
604
605 async def shutdown(self, sockets=None):
606 try:
607 await self._before_shutdown()
608 except Exception: # noqa: BLE001
609 _LOG.warning("e2e surface did not stop cleanly", exc_info=False)
610 await super().shutdown(sockets=sockets)
611
612 return LocalServer
613
614
615def _server_class():
616 import uvicorn
617
618 class QuietServer(uvicorn.Server):
619 """The E2E server: it does not take over SIGINT and SIGTERM (the LOCAL server owns them, and stops this one)."""
620
621 @contextlib.contextmanager
622 def capture_signals(self):
623 yield
624
625 return QuietServer
626
627
628def _bind_loopback() -> socket.socket:
629 """A listening socket on 127.0.0.1 and an ephemeral port. On Windows it is bound exclusively (SO_EXCLUSIVEADDRUSE), so nothing can share the port."""
630 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
631 try:
632 if hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
633 sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
634 sock.bind((LOOPBACK, 0))
635 sock.listen(socket.SOMAXCONN)
636 sock.setblocking(False)
637 except BaseException:
638 sock.close()
639 raise
640 return sock
641
642
643@dataclass
644class _Listener:
645 config: E2EConfig | None # what the E2E listener serves; None for the relay lane, which has nothing to follow
646 server: Any
647 task: "asyncio.Task[None]"
648 pid: int
649 port: int
650 secret: str
651
652
653class SurfaceRunner:
654 """Serves the app on the LOCAL surface and, while they are wanted, on the E2E one and on the relay lane: up to three uvicorn servers in one loop. `values`, `bind`
655 and `server_class` exist so the tests can drive it without a runtime.env, a real port or a real uvicorn; the defaults are the real ones (`lane_server_class` and
656 `lane_path` are the lane's own, so a test that breaks the E2E server does not break the lane with it)."""
657
658 def __init__(self, app, *, port: int, log_level: str = "info", values: Callable[[], dict[str, str]] = runtime_values,
659 bind: Callable[[], socket.socket] = _bind_loopback, server_class: Any = None, state_path: Path | None = None, watch: float = WATCH_SECONDS,
660 lane_path: Path | None = None, lane_server_class: Any = None):
661 self.app = app
662 self.port = port
663 self.log_level = log_level
664 self._values = values
665 self._bind = bind
666 self._server_class = server_class
667 self._path = state_path
668 self._watch_seconds = watch
669 self._lane_file = lane_path
670 self._lane_server_class = lane_server_class
671 self._listener: _Listener | None = None
672 self._failures = 0
673 self._retry_at = 0.0
674 self._lane: _Listener | None = None
675 self._lane_failures = 0
676 self._lane_retry_at = 0.0
677 self._stopping = False
678 self.local: Any = None
679
680 # ---- the whole thing
681
682 async def serve(self) -> None:
683 """Block until the LOCAL server stops (a signal, or it could not start: then SystemExit, as `uvicorn.run` had it)."""
684 self.local = self.make_local()
685 watcher = asyncio.create_task(self._watch())
686 try:
687 await self.local.serve()
688 finally:
689 watcher.cancel()
690 await asyncio.gather(watcher, return_exceptions=True)
691 await self.stop_lane()
692 await self.stop_e2e()
693
694 def make_local(self):
695 config = uvicorn_config(LocalSurface(self.app), log_level=self.log_level, concurrency=LOCAL_CONCURRENCY, host=LOOPBACK, port=self.port)
696 return _local_server_class()(config, self._begin_shutdown)
697
698 async def _begin_shutdown(self) -> None:
699 self._stopping = True
700 await self.stop_lane()
701 await self.stop_e2e()
702
703 async def _watch(self) -> None:
704 while True:
705 try:
706 await self.reconcile()
707 except Exception: # noqa: BLE001 - nothing the E2E side does may take the agent down
708 _LOG.warning("surface supervisor error", exc_info=False)
709 await asyncio.sleep(self._watch_seconds)
710
711 # ---- the E2E listener
712
713 async def reconcile(self) -> None:
714 """Make the E2E listener and the relay lane match what is wanted now. Each has its own boundary: one that fails never keeps the other from being looked at."""
715 if self._stopping:
716 return
717 try:
718 await self._reconcile_e2e()
719 finally:
720 await self._reconcile_lane()
721
722 async def _reconcile_e2e(self) -> None:
723 """Make the E2E listener match what is wanted now: start it, stop it, or start it again after it died (with a backoff)."""
724 wanted = e2e_config(self._values())
725 listener = self._listener
726 if listener is not None and listener.task.done():
727 _LOG.warning("e2e surface stopped by itself")
728 await self.stop_e2e()
729 self._failures += 1
730 self._retry_at = time.monotonic() + min(policy.value("timing.backoff_max_s"), policy.value("timing.backoff_min_s") * 2 ** (self._failures - 1))
731 listener = None
732 elif listener is not None:
733 self._failures = 0
734 if listener is not None and listener.config != wanted:
735 await self.stop_e2e()
736 listener = None
737 if wanted is not None and listener is None and time.monotonic() >= self._retry_at:
738 try:
739 await self.start_e2e(wanted)
740 except Exception: # noqa: BLE001 - one boundary: a listener that cannot start is a missing remote surface, not a dead agent
741 _LOG.warning("e2e surface could not start", exc_info=False)
742 await self.stop_e2e()
743 self._failures += 1
744 self._retry_at = time.monotonic() + min(policy.value("timing.backoff_max_s"), policy.value("timing.backoff_min_s") * 2 ** (self._failures - 1))
745
746 async def start_e2e(self, wanted: E2EConfig) -> None:
747 sock = self._bind()
748 try:
749 port = sock.getsockname()[1]
750 secret = new_secret()
751 pid = os.getpid()
752 surface = E2ESurface(self.app, host=wanted.host, secret=secret, pid=pid, port=port)
753 config = uvicorn_config(surface, log_level=self.log_level, concurrency=policy.value("terminator.max_streams") * E2E_CONCURRENCY_FACTOR,
754 lifespan="off", ws="none", timeout_graceful_shutdown=GRACEFUL_SHUTDOWN_SECONDS)
755 server = (self._server_class or _server_class())(config)
756 task = asyncio.create_task(self._guard(server, sock))
757 except BaseException:
758 sock.close()
759 raise
760 self._listener = _Listener(wanted, server, task, pid, port, secret)
761 deadline = time.monotonic() + START_TIMEOUT_SECONDS
762 while not getattr(server, "started", False) and not task.done():
763 if time.monotonic() > deadline:
764 raise RuntimeError("e2e surface did not start")
765 await asyncio.sleep(0.02)
766 if task.done() or self._stopping or self._listener is None:
767 raise RuntimeError("e2e surface did not start")
768 _allow_everywhere(wanted.host) # a computer that enrolled while the agent ran: the MCP transport must accept the name too
769 write_surfaces_file(pid=pid, port=port, secret=secret, path=self._path)
770 _serving["e2e"] = True
771
772 # ---- the relay lane
773
774 async def _reconcile_lane(self) -> None:
775 """The same loop for the lane the legacy relay's client dials (WP11): up while `relay_lane_wanted`, down when it is not, up again after it died (with the
776 backoff of the E2E listener). It needs no label: a computer that has never enrolled is exactly who the old relay serves."""
777 if self._stopping:
778 return
779 wanted = relay_lane_wanted(self._values())
780 listener = self._lane
781 if listener is not None and listener.task.done():
782 _LOG.warning("relay lane stopped by itself")
783 await self.stop_lane()
784 self._lane_failures += 1
785 self._lane_retry_at = time.monotonic() + min(policy.value("timing.backoff_max_s"), policy.value("timing.backoff_min_s") * 2 ** (self._lane_failures - 1))
786 listener = None
787 elif listener is not None:
788 self._lane_failures = 0
789 if listener is not None and not wanted:
790 await self.stop_lane()
791 listener = None
792 if wanted and listener is None and time.monotonic() >= self._lane_retry_at:
793 try:
794 await self.start_lane()
795 except Exception: # noqa: BLE001 - one boundary: a lane that cannot start is a legacy route that answers "unavailable" and is retried, not a dead agent
796 _LOG.warning("relay lane could not start", exc_info=False)
797 await self.stop_lane()
798 self._lane_failures += 1
799 self._lane_retry_at = time.monotonic() + min(policy.value("timing.backoff_max_s"),
800 policy.value("timing.backoff_min_s") * 2 ** (self._lane_failures - 1))
801
802 async def start_lane(self) -> None:
803 sock = self._bind()
804 try:
805 port = sock.getsockname()[1]
806 secret = new_secret()
807 pid = os.getpid()
808 surface = RelaySurface(self.app, secret=secret, pid=pid, port=port)
809 config = uvicorn_config(surface, log_level=self.log_level, concurrency=policy.value("terminator.max_streams") * E2E_CONCURRENCY_FACTOR,
810 lifespan="off", ws="none", timeout_graceful_shutdown=GRACEFUL_SHUTDOWN_SECONDS)
811 server = (self._lane_server_class or _server_class())(config)
812 task = asyncio.create_task(self._guard(server, sock))
813 except BaseException:
814 sock.close()
815 raise
816 self._lane = _Listener(None, server, task, pid, port, secret)
817 deadline = time.monotonic() + START_TIMEOUT_SECONDS
818 while not getattr(server, "started", False) and not task.done():
819 if time.monotonic() > deadline:
820 raise RuntimeError("relay lane did not start")
821 await asyncio.sleep(0.02)
822 if task.done() or self._stopping or self._lane is None:
823 raise RuntimeError("relay lane did not start")
824 write_surfaces_file(pid=pid, port=port, secret=secret, path=self._lane_file or lane_path())
825
826 async def stop_lane(self) -> None:
827 listener, self._lane = self._lane, None
828 if listener is None:
829 return
830 remove_surfaces_file(pid=listener.pid, secret=listener.secret, path=self._lane_file or lane_path()) # first: the client must not dial a port that is going away
831 await self._stop(listener)
832
833 async def _stop(self, listener: _Listener) -> None:
834 listener.server.should_exit = True
835 done, _ = await asyncio.wait({listener.task}, timeout=GRACEFUL_SHUTDOWN_SECONDS + 5)
836 if not done:
837 listener.server.force_exit = True
838 listener.task.cancel()
839 await asyncio.gather(listener.task, return_exceptions=True)
840
841 async def _guard(self, server, sock: socket.socket) -> None:
842 """The E2E server's boundary. uvicorn exits with SystemExit when it cannot start, and asyncio would carry that out of the event loop and end the agent."""
843 try:
844 await server.serve(sockets=[sock])
845 except (Exception, SystemExit) as exc: # noqa: BLE001
846 _LOG.warning("e2e surface failed: %s", type(exc).__name__)
847 finally:
848 # A server that died before its own shutdown (cancelled, crashed) still has its listening socket registered with the loop: close the asyncio servers
849 # first, or the next listener that is given the same file descriptor number would be served by this dead one.
850 for listening in getattr(server, "servers", None) or []:
851 listening.close()
852 sock.close()
853
854 async def stop_e2e(self) -> None:
855 listener, self._listener = self._listener, None
856 _serving["e2e"] = False
857 if listener is None:
858 return
859 remove_surfaces_file(pid=listener.pid, secret=listener.secret, path=self._path) # first: nothing may dial a port that is going away
860 await self._stop(listener)
861
862
863def run(app, *, port: int, log_level: str = "info") -> None:
864 """`uvicorn.run(app, host="127.0.0.1", port=port)` of the agent, now with the two surfaces. Blocks."""
865 asyncio.run(SurfaceRunner(app, port=port, log_level=log_level).serve())