For reading and verification only. You may not copy, use, modify or redistribute this code. Full notice on the main page.

src/all_mcp/surfaces.py

SHA-256: edfd470d4b4b4c3903946f963ef7e5bb3df8e5d46cb3e12056633a502280369d
865 lines · 47569 bytes · View on 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
3	One 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
17	Every 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
18	read 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
20	the wrappers below put traffic in front of the app in production.
21
22	Why 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
23	another process, which would then receive ChatGPT's tokens in clear. So, before the terminator forwards the first ChatGPT byte, it sends on that same
24	connection 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
27	The relay lane has the same identity proof (relay_client.RelayClient, the dialer, runs `prove_listener` against it before it trusts relay-lane.json).
28
29	The 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
30	runtime.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
31	lifespan 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
32	concurrency ceiling, and an h11 head cap (X-Forwarded-For is never believed: the client of a request is its socket's peer).
33	"""
34	from __future__ import annotations
35
36	import asyncio
37	import base64
38	import contextlib
39	import contextvars
40	import hashlib
41	import hmac
42	import json
43	import logging
44	import os
45	import re
46	import secrets
47	import socket
48	import stat
49	import time
50	from dataclasses import dataclass
51	from pathlib import Path
52	from typing import Any, Callable
53
54	from . import e2e_policy as policy
55	from . import state_io
56	from .config import runtime_values, state_dir
57
58	_LOG = logging.getLogger(__name__)
59
60	LOOPBACK = "127.0.0.1"
61	SURFACES_FILE = "surfaces.json"
62	LANE_FILE = "relay-lane.json"
63	# What a wrapper writes into scope["state"], and the two values it can have.
64	SURFACE_KEY = "all_mcp_surface"
65	LOCAL = "local"
66	E2E = "e2e"
67	RELAY = "relay"
68	# Both surfaces that carry remote traffic: what the authentication must never treat as the computer's own tools.
69	REMOTE = 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).
72	E2E_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).
76	SWITCH = "ALL_MCP_E2E"
77	REMOTE_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.
79	RELAY_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).
82	LEGACY_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.
87	CERT_ONLY = "ALL_E2E_CERT_ONLY"
88	CANARY_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).
92	LABEL_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).
96	E2E_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.
104	LANE_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).
111	LOCAL_CONCURRENCY = 512
112	E2E_CONCURRENCY_FACTOR = 2
113	GRACEFUL_SHUTDOWN_SECONDS = 10
114	START_TIMEOUT_SECONDS = 10.0
115	WATCH_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
121	def 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
127	def in_e2e() -> bool:
128	return _current.get() == E2E
129
130
131	def 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
136	def surface_of(scope: dict) -> str \| None:
137	return (scope.get("state") or {}).get(SURFACE_KEY)
138
139
140	# ------------------------------------------------------------------------------------------------------------------------ is the E2E listener wanted?
141
142	def _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)
147	class E2EConfig:
148	label: str
149	host: str # the device's full name, exactly what a caller must write in Host
150
151
152	def 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
165	def 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
173	def 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
180	def 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
188	def 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
195	def 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
207	def 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
217	def _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
224	def _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
231	def b64u(raw: bytes) -> str:
232	return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")
233
234
235	def _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
242	def new_secret() -> str:
243	return b64u(secrets.token_bytes(32))
244
245
246	def 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
251	def 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
257	def 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
266	def 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
275	def 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
281	class 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
289	async 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)
319	class SurfaceTarget:
320	pid: int
321	host: str
322	port: int
323	secret: str
324	started_at: float
325
326
327	def surfaces_path() -> Path:
328	return state_dir() / SURFACES_FILE
329
330
331	def lane_path() -> Path:
332	return state_dir() / LANE_FILE
333
334
335	def 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
340	def 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
351	def _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
366	def 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
392	def _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
400	class 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
417	async 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
426	def _values(scope, name: bytes) -> list[bytes]:
427	return [value for key, value in scope.get("headers", []) if key == name]
428
429
430	class 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
560	class 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
577	def 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
583	def 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
594	def _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
615	def _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
628	def _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
644	class _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
653	class 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
863	def 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())