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

src/all_mcp/e2e_terminator.py

SHA-256: 1c155ffd6a5fb34ce06b896469d222d2105cee18adc5d0131fc685409dd09129
839 lines · 43192 bytes · View on GitHub

1	"""The TLS terminator of ALL 2.0 (plan sections 2, 3 and 5): where ChatGPT's HTTPS ends on this computer.
2
3	The tunnel node forwards bytes it cannot read. Every stream the node opens for this computer (``OPEN``, then the device's ``ATTACH``) is a TLS connection that
4	ChatGPT made to ``<label>.<zone>``; this module wraps it in a server-side TLS with the device's own certificate and hands the plain HTTP/1.1 to the agent's E2E
5	listener, and nothing else:
6
7	* TLS 1.3 at least, ALPN ``http/1.1``, and the SNI must be this computer's own name (a handshake for any other name, or with none, ends with an alert);
8	the certificate is the one ``tls/current`` names at the moment of the handshake (``tls_store.ServingContext``), so a renewal drops no stream.
9	* The terminator sees the HTTP in clear and the client's address in the OPEN (the TLS it terminates hides it from the agent), so the limits that need an
10	address live here: streams in all, per address, per second, and per endpoint (``/register``, ``/authorize``, ``/token``, ``/pair``) per minute, with a part
11	of the streams kept for class ``openai`` so that a flood of strangers cannot starve ChatGPT. The agent never sees an address and never needs one.
12	* Nothing is rewritten. The request head is checked and forwarded byte for byte: it must arrive whole within ``terminator.request_head_s`` and within
13	``terminator.request_head_max_bytes``, ends at CRLFCRLF (a bare LF, an obsolete line fold, a repeated or non-numeric ``Content-Length``, a ``Content-Length``
14	together with ``Transfer-Encoding``, or an ``Upgrade`` is refused: h11 in the agent would unfold, merge or accept them, which is how a request is smuggled).
15	The body is followed by its framing (length or chunks), so that the NEXT request on the same connection is a head again and is counted too.
16	* The destination is the ONE address in ``surfaces.json`` (never LOCAL), read again for every stream, its pid alive, and before the first byte of ChatGPT's
17	goes anywhere the listener proves, on that same connection, that it is the agent (the HMAC challenge of surfaces.prove_listener): a port the agent freed
18	and another process bound receives the challenge and nothing more.
19	* Idle time and age come from the policy the node offered clamped to this device's own (``e2e_tunnel.Windows``); a stream is closed when it is silent for the
20	first or has lived the second.
21
22	No address is written to a log (the plan: nothing persistent names a client).
23	"""
24	from __future__ import annotations
25
26	import asyncio
27	import contextlib
28	import ipaddress
29	import logging
30	import re
31	import socket
32	import ssl
33	import time
34	from collections import OrderedDict, deque
35	from dataclasses import dataclass, field
36	from typing import Any, Callable
37
38	from . import e2e_policy as policy
39	from . import surfaces, tls_policy
40
41	log = logging.getLogger("all_mcp.e2e_terminator")
42
43	TLS_HANDSHAKE_BYTE = 0x16
44	# Constants of this module that the policy does not define (it defines the windows, the limits and the sizes): how long a local dial may take, how much of the
45	# agent's answer is read at once, and how the keepalive of a TCP connection is shaped.
46	LOCAL_DIAL_SECONDS = 5.0
47	READ_BYTES = 65536
48	KEEPALIVE_IDLE_SECONDS = 30
49	KEEPALIVE_INTERVAL_SECONDS = 10
50	KEEPALIVE_PROBES = 3
51	RATE_WINDOW_SECONDS = 1.0
52	BUDGET_WINDOW_SECONDS = 60.0
53	CHUNK_LINE_MAX = 64
54	TRAILER_MAX_BYTES = 16384
55	_TOKEN = re.compile(rb"[!#$%&'*+.^_`\|~0-9A-Za-z-]+")
56	_REQUEST_LINE = re.compile(rb"([!#$%&'+.^_`\|~0-9A-Za-z-]{1,32}) (/[\x21-\x7e]) HTTP/1\.([01])")
57	_DIGITS = re.compile(rb"[0-9]{1,18}")
58	_CHUNK_SIZE = re.compile(rb"[0-9a-fA-F]{1,15}")
59	# The end-points that are budgeted per address: the keys of policy.terminator.endpoint_budgets, which are also the paths the E2E surface answers, except
60	# ANONYMOUS_MCP, which is not a path: it is /mcp asked for without a credential the terminator has seen work (no bearer, or a bearer from an address that has
61	# never completed an authenticated call), the first call of every connector (the 401 that starts OAuth discovery) and what a scanner sends.
62	ANONYMOUS_MCP = "mcp_anonymous"
63	BUDGETED = tuple(name for name in policy.value("terminator.endpoint_budgets") if name != ANONYMOUS_MCP)
64	AUTHENTICATED_PATHS = frozenset({"/mcp", "/mcp/", "/token"})
65	# How long an address that completed an authenticated call is remembered as a client of this computer (ChatGPT's own addresses keep renewing it).
66	KNOWN_CLIENT_SECONDS = 3600.0
67
68
69	# ------------------------------------------------------------------------------------------------------------------------ small helpers
70
71	def ip_group(ip: str) -> str:
72	"""What a client address is counted as: an IPv4 address by itself, an IPv6 address by its /64 (a household or an ISP lease owns the whole /64), an IPv4-mapped
73	IPv6 address as the IPv4 it carries. Anything that is not an address is one shared group, so the limits still apply to it."""
74	try:
75	address = ipaddress.ip_address(ip)
76	except ValueError:
77	return "invalid"
78	if isinstance(address, ipaddress.IPv6Address):
79	if address.ipv4_mapped is not None:
80	return str(address.ipv4_mapped)
81	return str(ipaddress.ip_network((str(address), policy.value("node_limits.ipv6_prefix_bits")), strict=False))
82	return str(address)
83
84
85	def keepalive(writer: asyncio.StreamWriter) -> None:
86	"""TCP keepalive on the connection (a half-open one, after a network change, must not live for ever)."""
87	sock = writer.get_extra_info("socket")
88	if sock is None:
89	return
90	try:
91	sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
92	if hasattr(socket, "TCP_KEEPIDLE"):
93	sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, KEEPALIVE_IDLE_SECONDS)
94	elif hasattr(socket, "TCP_KEEPALIVE"): # macOS
95	sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, KEEPALIVE_IDLE_SECONDS)
96	if hasattr(socket, "TCP_KEEPINTVL"):
97	sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, KEEPALIVE_INTERVAL_SECONDS)
98	if hasattr(socket, "TCP_KEEPCNT"):
99	sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, KEEPALIVE_PROBES)
100	except OSError:
101	pass
102
103
104	async def close_writer(writer: asyncio.StreamWriter \| None) -> None:
105	if writer is None:
106	return
107	with contextlib.suppress(Exception):
108	writer.close()
109	with contextlib.suppress(Exception):
110	await asyncio.wait_for(writer.wait_closed(), 2.0)
111
112
113	def http_response(status: int, reason: str, text: str, *, extra: tuple[tuple[str, str], ...] = ()) -> bytes:
114	"""A complete small HTTP/1.1 answer the terminator gives itself (a refusal, never a rewrite of what the agent says): plain text, closes the connection."""
115	body = (text + "\n").encode("ascii")
116	lines = [f"HTTP/1.1 {status} {reason}", "Content-Type: text/plain", f"Content-Length: {len(body)}", "Cache-Control: no-store", "Connection: close"]
117	lines.extend(f"{name}: {value}" for name, value in extra)
118	return ("\r\n".join(lines) + "\r\n\r\n").encode("ascii") + body
119
120
121	# ------------------------------------------------------------------------------------------------------------------------ TLS over a stream that is not a socket
122
123	class TLSStream:
124	"""Server-side TLS over any asyncio byte stream (here: the data connection to the tunnel node, which is itself TLS). ``ssl.MemoryBIO`` lets TLS run over a stream
125	that is not a socket, which asyncio's own start_tls cannot do on top of another TLS transport."""
126
127	def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, context: ssl.SSLContext, *, initial: bytes = b""):
128	self._reader = reader
129	self._writer = writer
130	self._incoming = ssl.MemoryBIO()
131	self._outgoing = ssl.MemoryBIO()
132	self._ssl = context.wrap_bio(self._incoming, self._outgoing, server_side=True)
133	if initial:
134	self._incoming.write(initial)
135	self._eof = False
136	self._decrypted = b"" # bytes already decrypted when a read was cancelled while it flushed a key update: given back by the next read, never lost
137
138	async def flush(self) -> None:
139	data = self._outgoing.read()
140	if data:
141	self._writer.write(data)
142	await self._writer.drain()
143
144	async def _feed(self) -> bool:
145	if self._eof:
146	return False
147	data = await self._reader.read(READ_BYTES)
148	if not data:
149	self._incoming.write_eof()
150	self._eof = True
151	return False
152	self._incoming.write(data)
153	return True
154
155	async def handshake(self, timeout: float) -> None:
156	async def go() -> None:
157	while True:
158	try:
159	self._ssl.do_handshake()
160	return
161	except ssl.SSLWantReadError:
162	await self.flush()
163	if not await self._feed():
164	raise ConnectionResetError("the client closed during the TLS handshake")
165	try:
166	await asyncio.wait_for(go(), timeout)
167	finally:
168	with contextlib.suppress(Exception):
169	await self.flush() # the final flight, or the alert that explains a refusal
170
171	def selected_alpn(self) -> str \| None:
172	return self._ssl.selected_alpn_protocol()
173
174	def version(self) -> str \| None:
175	return self._ssl.version()
176
177	async def read(self, size: int = READ_BYTES) -> bytes:
178	"""Decrypted bytes; b"" when the client closed (close_notify or a plain close). Safe to cancel (a timeout): nothing that was decrypted is lost."""
179	if self._decrypted:
180	data, self._decrypted = self._decrypted, b""
181	return data
182	while True:
183	try:
184	data = self._ssl.read(size)
185	except ssl.SSLWantReadError:
186	await self.flush()
187	if not await self._feed():
188	return b""
189	continue
190	except (ssl.SSLZeroReturnError, ssl.SSLEOFError):
191	return b""
192	if self._outgoing.pending:
193	self._decrypted = data
194	await self.flush() # TLS 1.3 key updates and the like
195	self._decrypted = b""
196	return data
197
198	async def write(self, data: bytes) -> None:
199	view = memoryview(data)
200	while view:
201	written = self._ssl.write(view[:16384])
202	view = view[written:]
203	if self._outgoing.pending >= READ_BYTES or not view:
204	await self.flush()
205
206	async def close(self) -> None:
207	"""Say close_notify and go; the peer's own close_notify is not waited for."""
208	with contextlib.suppress(ssl.SSLError, OSError, ValueError):
209	self._ssl.unwrap()
210	with contextlib.suppress(Exception):
211	await self.flush()
212
213
214	# ------------------------------------------------------------------------------------------------------------------------ limits
215
216	@dataclass(eq=False)
217	class Ticket:
218	"""One admitted stream: give it back with ``Admission.release``."""
219	group: str
220	cls: str
221	released: bool = False
222
223
224	class Admission:
225	"""Who may hold a stream, decided from the OPEN (the client's address and the node's class for it) before the device dials the node back.
226
227	The streams in all are ``terminator.max_streams``; ``node_limits.openai_reserved_pct`` of them are reserved to class ``openai`` (strangers and hinted
228	addresses together can hold only the rest, so a flood from anywhere never leaves ChatGPT without a slot). A stranger or hinted address is also held to
229	``per_ip_concurrent`` streams and ``per_ip_new_per_s`` new ones a second (an IPv6 address by its /64); class ``openai`` is not counted per address, because
230	a handful of addresses carry all of ChatGPT. Every table is bounded (``surface.limiter_lru_max``), the oldest entry is the first to go."""
231
232	def __init__(self, *, max_streams: int \| None = None, reserved_pct: int \| None = None, per_ip_concurrent: int \| None = None,
233	per_ip_new_per_s: int \| None = None, table_max: int \| None = None, clock: Callable[[], float] = time.monotonic):
234	self.total_max = policy.value("terminator.max_streams") if max_streams is None else max_streams
235	percent = policy.value("node_limits.openai_reserved_pct") if reserved_pct is None else reserved_pct
236	self.reserved = -(-self.total_max * percent // 100) # rounded up: a quarter of 3 streams is still one stream
237	self.per_ip = policy.value("node_limits.per_ip_concurrent") if per_ip_concurrent is None else per_ip_concurrent
238	self.per_ip_rate = policy.value("node_limits.per_ip_new_per_s") if per_ip_new_per_s is None else per_ip_new_per_s
239	self._table_max = policy.value("surface.limiter_lru_max") if table_max is None else table_max
240	self._clock = clock
241	self.active = 0
242	self.active_other = 0
243	self._concurrent: dict[str, int] = {}
244	self._arrivals: OrderedDict[str, deque[float]] = OrderedDict()
245
246	@property
247	def other_max(self) -> int:
248	return self.total_max - self.reserved
249
250	def admit(self, ip: str, cls: str) -> tuple[Ticket \| None, str]:
251	"""(a ticket, "") when the stream may go on, else (None, the word that says why)."""
252	group = ip_group(ip)
253	openai = cls == "openai"
254	if self.active >= self.total_max:
255	return None, "full"
256	if not openai:
257	if self.active_other >= self.other_max:
258	return None, "full_other"
259	if self._concurrent.get(group, 0) >= self.per_ip:
260	return None, "ip_concurrent"
261	now = self._clock()
262	window = self._arrivals.get(group)
263	if window is None:
264	window = self._arrivals[group] = deque()
265	while len(self._arrivals) > self._table_max:
266	self._arrivals.popitem(last=False)
267	else:
268	self._arrivals.move_to_end(group)
269	while window and now - window[0] >= RATE_WINDOW_SECONDS:
270	window.popleft()
271	if len(window) >= self.per_ip_rate:
272	return None, "ip_rate"
273	window.append(now)
274	self._concurrent[group] = self._concurrent.get(group, 0) + 1
275	self.active_other += 1
276	self.active += 1
277	return Ticket(group, cls), ""
278
279	def release(self, ticket: Ticket) -> None:
280	if ticket.released:
281	return
282	ticket.released = True
283	self.active -= 1
284	if ticket.cls != "openai":
285	self.active_other -= 1
286	left = self._concurrent.get(ticket.group, 1) - 1
287	if left > 0:
288	self._concurrent[ticket.group] = left
289	else:
290	self._concurrent.pop(ticket.group, None)
291
292
293	class EndpointBudgets:
294	"""The budgets per address and per endpoint class, per minute (``terminator.endpoint_budgets``): the endpoints that anybody who can read the certificate
295	transparency logs can reach without a token (``/register``, ``/authorize``, ``/token``, ``/pair``). One address using up its budget costs the others nothing."""
296
297	def __init__(self, *, budgets: dict[str, dict[str, int]] \| None = None, table_max: int \| None = None, clock: Callable[[], float] = time.monotonic):
298	self._budgets = policy.value("terminator.endpoint_budgets") if budgets is None else budgets
299	self._table_max = policy.value("surface.limiter_lru_max") if table_max is None else table_max
300	self._clock = clock
301	self._windows: OrderedDict[tuple[str, str], deque[float]] = OrderedDict()
302	self._known: OrderedDict[str, float] = OrderedDict() # address group -> when it last completed an authenticated call
303
304	@staticmethod
305	def classify(path: str, *, credential_seen: bool = True) -> str \| None:
306	"""The budgeted endpoint a request path belongs to (its query is not part of it), or None. `/mcp` without a credential the terminator has seen work is the
307	anonymous class of the policy; with one, it is not budgeted (ChatGPT's addresses carry the traffic of every user of this computer)."""
308	name = path.split("?", 1)[0].rstrip("/")
309	if name == "/mcp" and not credential_seen:
310	return ANONYMOUS_MCP
311	return name[1:] if name[1:] in BUDGETED else None
312
313	def note_client(self, ip: str) -> None:
314	"""This address just completed an authenticated call: it is a client of this computer, and its `/mcp` calls are no longer counted as anonymous."""
315	key = ip_group(ip)
316	self._known[key] = self._clock()
317	self._known.move_to_end(key)
318	while len(self._known) > self._table_max:
319	self._known.popitem(last=False)
320
321	def is_client(self, ip: str) -> bool:
322	key = ip_group(ip)
323	since = self._known.get(key)
324	if since is None:
325	return False
326	if self._clock() - since >= KNOWN_CLIENT_SECONDS:
327	del self._known[key]
328	return False
329	return True
330
331	def allow(self, endpoint: str, ip: str) -> tuple[bool, float]:
332	"""(True, 0) and the request is counted, or (False, seconds until it would be allowed)."""
333	limit = self._budgets[endpoint]["per_ip_per_min"]
334	key = (endpoint, ip_group(ip))
335	now = self._clock()
336	window = self._windows.get(key)
337	if window is None:
338	window = self._windows[key] = deque()
339	while len(self._windows) > self._table_max:
340	self._windows.popitem(last=False)
341	else:
342	self._windows.move_to_end(key)
343	while window and now - window[0] >= BUDGET_WINDOW_SECONDS:
344	window.popleft()
345	if len(window) >= limit:
346	return False, max(1.0, BUDGET_WINDOW_SECONDS - (now - window[0]))
347	window.append(now)
348	return True, 0.0
349
350
351	# ------------------------------------------------------------------------------------------------------------------------ the request head and its framing
352
353	class BadRequest(Exception):
354	"""A request the terminator refuses. `status` is what the client is told (when it can still be told), `word` a fixed word for the counters."""
355
356	def __init__(self, status: int, word: str):
357	super().__init__(word)
358	self.status = status
359	self.word = word
360
361
362	@dataclass(frozen=True)
363	class RequestHead:
364	method: str
365	target: str
366	path: str
367	headers: tuple[tuple[bytes, bytes], ...] # lower-case name, value as sent (without the surrounding blanks)
368	body: str # "none" \| "length" \| "chunked"
369	length: int
370	raw_size: int
371
372	def header(self, name: bytes) -> bytes \| None:
373	for key, value in self.headers:
374	if key == name:
375	return value
376	return None
377
378
379	def parse_head(raw: bytes) -> RequestHead:
380	"""Check one request head (the bytes up to and including the CRLFCRLF) and say how its body is framed. Raises BadRequest. Nothing is changed: the caller forwards `raw`."""
381	if len(raw) > policy.value("terminator.request_head_max_bytes"):
382	raise BadRequest(431, "head_too_large")
383	if not raw.endswith(b"\r\n\r\n"):
384	raise BadRequest(400, "bad_head")
385	body = raw[:-4]
386	if b"\n" in body.replace(b"\r\n", b"") or b"\r" in body.replace(b"\r\n", b""):
387	raise BadRequest(400, "bare_line_ending") # a lone LF or CR: what a lenient parser reads as a line end and a strict one does not
388	lines = body.split(b"\r\n")
389	found = _REQUEST_LINE.fullmatch(lines[0])
390	if found is None:
391	raise BadRequest(400, "bad_request_line")
392	headers: list[tuple[bytes, bytes]] = []
393	for line in lines[1:]:
394	if line[:1] in (b" ", b"\t"):
395	raise BadRequest(400, "obs_fold") # h11 unfolds an obsolete line fold into the previous header: the agent would read what this check never saw
396	name, colon, value = line.partition(b":")
397	if not colon or _TOKEN.fullmatch(name) is None:
398	raise BadRequest(400, "bad_header")
399	value = value.strip(b" \t")
400	if any(byte < 0x20 and byte != 0x09 or byte == 0x7F for byte in value):
401	raise BadRequest(400, "bad_header_value")
402	if len(value) > policy.value("surface.header_value_max_bytes"):
403	raise BadRequest(431, "header_too_large")
404	headers.append((name.lower(), value))
405	if len(headers) > policy.value("surface.headers_max"):
406	raise BadRequest(431, "too_many_headers")
407	lengths = [value for name, value in headers if name == b"content-length"]
408	encodings = [value for name, value in headers if name == b"transfer-encoding"]
409	if len(lengths) > 1:
410	raise BadRequest(400, "repeated_content_length") # even identical: h11 merges them, a proxy in front of another parser may not
411	if len(encodings) > 1 or (encodings and encodings[0].lower() != b"chunked") or (encodings and lengths):
412	raise BadRequest(400, "bad_framing")
413	for name, value in headers:
414	if name == b"upgrade" or (name == b"connection" and b"upgrade" in value.lower()):
415	raise BadRequest(400, "upgrade") # the E2E surface is plain HTTP: there is no protocol to switch to
416	length, kind = 0, "none"
417	if lengths:
418	if _DIGITS.fullmatch(lengths[0]) is None:
419	raise BadRequest(400, "bad_content_length")
420	length, kind = int(lengths[0]), "length"
421	if length > policy.value("surface.body_max_bytes"):
422	raise BadRequest(413, "body_too_large")
423	elif encodings:
424	kind = "chunked"
425	target = found.group(2).decode("ascii")
426	path = target.split("?", 1)[0]
427	if "%" in path:
428	raise BadRequest(400, "encoded_path") # the agent's server decodes the path before it routes (/%72egister is /register): the budgets below would read another name
429	return RequestHead(found.group(1).decode("ascii"), target, path, tuple(headers), kind, length, len(raw))
430
431
432	@dataclass
433	class Segment:
434	"""A slice of what the client sent, in order: a head (with its parse) or part of a body (or of its chunk framing). Together they are exactly the bytes received."""
435	kind: str # "head" \| "body"
436	data: bytes
437	head: RequestHead \| None = None
438
439
440	class RequestFramer:
441	"""Splits the client's byte stream into request heads and body, so that every request of a keep-alive connection (not only the first) is a head again. A head is
442	held back until it is whole and has passed ``parse_head`` (an incomplete head is never forwarded); body bytes are passed on as they arrive. Raises BadRequest."""
443
444	def __init__(self) -> None:
445	self._buffer = bytearray()
446	self._state = "head"
447	self._remaining = 0
448	self._total = 0
449	self._trailer = 0
450	self._chunked_body = False
451
452	@property
453	def awaiting_head(self) -> bool:
454	return self._state == "head"
455
456	@property
457	def partial_head(self) -> bool:
458	return self._state == "head" and bool(self._buffer)
459
460	def feed(self, data: bytes) -> list[Segment]:
461	self._buffer += data
462	out: list[Segment] = []
463	while self._buffer:
464	if self._state == "head":
465	end = self._buffer.find(b"\r\n\r\n")
466	if end < 0:
467	if len(self._buffer) > policy.value("terminator.request_head_max_bytes"):
468	raise BadRequest(431, "head_too_large")
469	break
470	raw = bytes(self._buffer[:end + 4])
471	head = parse_head(raw)
472	del self._buffer[:end + 4]
473	out.append(Segment("head", raw, head))
474	self._total = 0
475	if head.body == "length" and head.length:
476	self._state, self._remaining = "body", head.length
477	elif head.body == "chunked":
478	self._state = "chunk_size"
479	elif self._state == "body":
480	take = min(self._remaining, len(self._buffer))
481	out.append(Segment("body", bytes(self._buffer[:take])))
482	del self._buffer[:take]
483	self._remaining -= take
484	if not self._remaining:
485	self._state = "head"
486	elif self._state == "chunk_size":
487	end = self._buffer.find(b"\r\n")
488	if end < 0:
489	if len(self._buffer) > CHUNK_LINE_MAX:
490	raise BadRequest(400, "bad_chunk")
491	break
492	if end > CHUNK_LINE_MAX or _CHUNK_SIZE.fullmatch(bytes(self._buffer[:end])) is None:
493	raise BadRequest(400, "bad_chunk") # no chunk extensions: a size is hexadecimal digits and nothing else
494	size = int(self._buffer[:end], 16)
495	out.append(Segment("body", bytes(self._buffer[:end + 2])))
496	del self._buffer[:end + 2]
497	self._total += size
498	if self._total > policy.value("surface.body_max_bytes"):
499	raise BadRequest(413, "body_too_large")
500	self._state, self._remaining, self._trailer = ("trailer", 0, 0) if size == 0 else ("chunk_data", size, 0)
501	elif self._state == "chunk_data":
502	take = min(self._remaining, len(self._buffer))
503	out.append(Segment("body", bytes(self._buffer[:take])))
504	del self._buffer[:take]
505	self._remaining -= take
506	if not self._remaining:
507	self._state = "chunk_end"
508	elif self._state == "chunk_end":
509	if len(self._buffer) < 2:
510	break
511	if bytes(self._buffer[:2]) != b"\r\n":
512	raise BadRequest(400, "bad_chunk")
513	out.append(Segment("body", b"\r\n"))
514	del self._buffer[:2]
515	self._state = "chunk_size"
516	else: # trailer: header lines up to the empty one
517	end = self._buffer.find(b"\r\n")
518	if end < 0:
519	if len(self._buffer) + self._trailer > TRAILER_MAX_BYTES:
520	raise BadRequest(431, "head_too_large")
521	break
522	line = bytes(self._buffer[:end + 2])
523	self._trailer += len(line)
524	if self._trailer > TRAILER_MAX_BYTES or b"\n" in line[:-2] or line[:1] in (b" ", b"\t"):
525	raise BadRequest(400, "bad_chunk")
526	out.append(Segment("body", line))
527	del self._buffer[:end + 2]
528	if end == 0:
529	self._state = "head"
530	return out
531
532
533	# ------------------------------------------------------------------------------------------------------------------------ the terminator
534
535	@dataclass
536	class TerminatorStats:
537	streams: int = 0
538	active: int = 0
539	tls_streams: int = 0
540	handshake_failures: int = 0
541	no_certificate: int = 0
542	refused: dict[str, int] = field(default_factory=dict) # admission refusals, by word
543	bad_requests: dict[str, int] = field(default_factory=dict)
544	budget_refused: int = 0
545	no_surface: int = 0
546	identity_failures: int = 0
547	agent_unreachable: int = 0
548	head_timeouts: int = 0
549	idle_closes: int = 0
550	age_closes: int = 0
551	bytes_up: int = 0 # client -> agent
552	bytes_down: int = 0 # agent -> client
553
554	def as_dict(self) -> dict[str, Any]:
555	return dict(self.__dict__)
556
557	def count(self, table: dict[str, int], word: str) -> None:
558	table[word] = table.get(word, 0) + 1
559
560
561	class Terminator:
562	"""Serves the streams the node opens: TLS, the checks above, the dial to the agent and the pipe between them.
563
564	`serving` is a tls_store.ServingContext (or anything with ``listener_context()``, ``served()`` and ``number``); `windows` is a callable that returns the
565	current ``e2e_tunnel.Windows`` (idle and age); `host` is the device's own name, the only SNI it answers; `on_good(ip)` is told when a client address just
566	completed an authenticated call (the tunnel sends it to the node as a HINT). `load_target` and `prove` are the agent's surface file and the identity
567	proof (surfaces.load_target, surfaces.prove_listener): tests replace them, nothing else does."""
568
569	def __init__(self, *, host: str, serving: Any, windows: Callable[[], Any], budgets: EndpointBudgets \| None = None,
570	on_good: Callable[[str], None] \| None = None, load_target: Callable[[], Any] \| None = None, prove: Callable[..., Any] \| None = None,
571	local_port: int \| None = None, handshake_timeout: float \| None = None, head_timeout: float \| None = None,
572	clock: Callable[[], float] = time.monotonic):
573	self.host = host
574	self.serving = serving
575	self.windows = windows
576	self.budgets = budgets if budgets is not None else EndpointBudgets(clock=clock)
577	self.on_good = on_good
578	self._load_target = load_target or surfaces.load_target
579	self._prove = prove or surfaces.prove_listener
580	self._local_port = local_port
581	self.handshake_timeout = policy.value("terminator.tls_handshake_s") if handshake_timeout is None else handshake_timeout
582	self.head_timeout = policy.value("terminator.request_head_s") if head_timeout is None else head_timeout
583	self._clock = clock
584	self.stats = TerminatorStats()
585	self._context: ssl.SSLContext \| None = None
586	self._served_number: int \| None = None
587	self.last_error = ""
588
589	# -- TLS
590	def tls_ready(self) -> bool:
591	"""TLS 1.3 is what the terminator will not serve without (plan section 3): a Python whose OpenSSL has none must say so instead of serving less
592	(tls_policy decides; the status file shows ``degraded`` / ``tls13_required``)."""
593	return tls_policy.tls13_available()
594
595	def _listener_context(self) -> ssl.SSLContext \| None:
596	"""The one context that answers every handshake and follows the certificate pointer by itself; None while there is no certificate yet."""
597	if self._context is not None:
598	return self._context
599	try:
600	context = self.serving.listener_context()
601	except Exception as exc: # noqa: BLE001 - no certificate (or a store that cannot be read) is "no certificate", never a crash
602	self.last_error = getattr(exc, "word", type(exc).__name__.lower())
603	return None
604	if context.minimum_version < ssl.TLSVersion.TLSv1_3:
605	self.last_error = "tls13_required"
606	return None
607	wanted = self.host
608	follow = context.sni_callback
609
610	def only_my_name(connection, server_name, original):
611	if not server_name or server_name.lower().rstrip(".") != wanted:
612	return ssl.ALERT_DESCRIPTION_UNRECOGNIZED_NAME # no SNI, or another name: the handshake ends with an alert
613	return follow(connection, server_name, original) if follow is not None else None
614
615	context.sni_callback = only_my_name
616	self._context = context
617	self.last_error = ""
618	return context
619
620	async def _handshake(self, reader, writer, first: bytes) -> TLSStream \| None:
621	ready = self.tls_ready()
622	context = self._listener_context() if ready else None
623	if context is None:
624	if ready:
625	self.stats.no_certificate += 1
626	else:
627	self.last_error = "tls13_required" # never silent: the tunnel reports tls13=False and the runner turns it into `degraded`
628	self.stats.handshake_failures += 1
629	return None
630	tls = TLSStream(reader, writer, context, initial=first)
631	try:
632	await tls.handshake(self.handshake_timeout)
633	except (ssl.SSLError, OSError, ConnectionError, asyncio.TimeoutError, ValueError) as exc:
634	self.stats.handshake_failures += 1
635	log.debug("TLS handshake failed (%s)", type(exc).__name__)
636	return None
637	if tls.selected_alpn() not in (None, policy.value("terminator.tls_alpn")):
638	self.stats.handshake_failures += 1
639	return None
640	self.stats.tls_streams += 1
641	number = getattr(self.serving, "number", None)
642	if number is not None and number != self._served_number:
643	self._served_number = number
644	with contextlib.suppress(Exception):
645	await asyncio.to_thread(self.serving.served) # a certificate that completed a handshake lets the previous generation go
646	return tls
647
648	# -- the agent
649	async def _dial_agent(self) -> tuple[asyncio.StreamReader, asyncio.StreamWriter] \| None:
650	"""The agent's E2E listener, proven. Never LOCAL: surfaces.json names the E2E port (and is read again now: a restart of the agent moved it), the pid must
651	be alive, and the listener must answer the challenge on this very connection before anything of ChatGPT's is written to it."""
652	try:
653	target = self._load_target()
654	except surfaces.SurfaceIdentityError:
655	self.stats.no_surface += 1
656	return None
657	if target.port == self._local_surface_port():
658	self.stats.no_surface += 1 # a surfaces.json that names the LOCAL port is not an E2E listener (the bootstrap bearer works there)
659	return None
660	try:
661	reader, writer = await asyncio.wait_for(asyncio.open_connection(target.host, target.port, limit=1 << 16), LOCAL_DIAL_SECONDS)
662	except (OSError, asyncio.TimeoutError):
663	self.stats.agent_unreachable += 1
664	return None
665	try:
666	await self._prove(reader, writer, host=target.host, port=target.port, secret=target.secret, pid=target.pid, timeout=LOCAL_DIAL_SECONDS)
667	except surfaces.SurfaceIdentityError:
668	self.stats.identity_failures += 1
669	await close_writer(writer)
670	return None
671	return reader, writer
672
673	def _local_surface_port(self) -> int:
674	if self._local_port is not None:
675	return self._local_port
676	import os
677	return int(os.environ.get("ALL_MCP_PORT", "4410"))
678
679	# -- one stream
680	async def serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, *, ip: str, first: bytes,
681	on_good: Callable[[str], None] \| None = None) -> None:
682	"""One stream the node attached: `first` is what has already been read (the ClientHello); `on_good` replaces the terminator's own when this stream came
683	through a particular node (its HINT goes to that node). Returns when it ends; never raises."""
684	self.stats.streams += 1
685	self.stats.active += 1
686	tls: TLSStream \| None = None
687	agent_writer: asyncio.StreamWriter \| None = None
688	try:
689	tls = await self._handshake(reader, writer, first)
690	if tls is None:
691	return
692	agent_writer = await self._exchange(tls, ip, on_good or self.on_good)
693	except asyncio.CancelledError:
694	raise
695	except Exception as exc: # noqa: BLE001 - a stream must end quietly
696	log.debug("stream ended: %s", type(exc).__name__)
697	finally:
698	self.stats.active -= 1
699	await close_writer(agent_writer)
700	if tls is not None:
701	await tls.close()
702
703	async def _exchange(self, tls: TLSStream, ip: str, on_good: Callable[[str], None] \| None) -> asyncio.StreamWriter \| None:
704	"""The pipe between the client (through `tls`) and the agent, until either ends or a window runs out. Returns the agent's writer, which the caller closes."""
705	windows = self.windows()
706	framer = RequestFramer()
707	started = self._clock()
708	head_since = started # the first request has head_timeout from the end of the handshake, a later one from its first byte
709	last = started # the last byte that moved in either direction
710	forwarded = 0 # request heads passed on
711	first_good = False
712	first_status_seen = False
713	agent_writer: asyncio.StreamWriter \| None = None
714	agent_task: asyncio.Task \| None = None
715	agent_done: asyncio.Future = asyncio.get_running_loop().create_future()
716
717	async def refuse(status: int, reason: str, text: str, extra=()) -> None:
718	with contextlib.suppress(Exception):
719	await tls.write(http_response(status, reason, text, extra=extra))
720
721	async def agent_to_client(agent_reader: asyncio.StreamReader) -> None:
722	nonlocal last, first_status_seen
723	while True:
724	data = await agent_reader.read(READ_BYTES)
725	if not data:
726	return
727	if not first_status_seen:
728	first_status_seen = True
729	if first_good and data[:9] == b"HTTP/1.1 " and data[9:10] == b"2":
730	# a stranger cannot make the agent answer 2xx to an authenticated call: this address is a client of this device
731	self.budgets.note_client(ip)
732	if on_good is not None:
733	with contextlib.suppress(Exception):
734	on_good(ip)
735	last = self._clock()
736	self.stats.bytes_down += len(data)
737	await tls.write(data)
738
739	async def client_to_agent() -> None:
740	nonlocal last, head_since, forwarded, first_good, agent_writer, agent_task
741	while True:
742	now = self._clock()
743	if now - started >= windows.stream_max_age:
744	raise _Over("age")
745	if now - last >= windows.stream_idle:
746	raise _Over("idle")
747	waits = [windows.stream_idle - (now - last), windows.stream_max_age - (now - started)]
748	if framer.awaiting_head and (framer.partial_head or not forwarded):
749	if now - head_since >= self.head_timeout:
750	raise _Over("head")
751	waits.append(self.head_timeout - (now - head_since))
752	try:
753	data = await asyncio.wait_for(tls.read(), min(waits))
754	except asyncio.TimeoutError:
755	continue # the deadlines are looked at again above: traffic the other way may have moved the idle clock
756	if not data:
757	return
758	if forwarded and framer.awaiting_head and not framer.partial_head:
759	head_since = self._clock()
760	last = self._clock()
761	self.stats.bytes_up += len(data)
762	try:
763	segments = framer.feed(data)
764	except BadRequest as bad:
765	self.stats.count(self.stats.bad_requests, bad.word)
766	if not forwarded:
767	await refuse(bad.status, "Bad Request" if bad.status == 400 else "Request Rejected", bad.word)
768	return
769	for segment in segments:
770	if segment.kind == "head":
771	head = segment.head
772	bearer = (head.header(b"authorization") or b"")[:7].lower() == b"bearer "
773	endpoint = self.budgets.classify(head.path, credential_seen=bearer and self.budgets.is_client(ip))
774	if endpoint is not None:
775	allowed, retry = self.budgets.allow(endpoint, ip)
776	if not allowed:
777	self.stats.budget_refused += 1
778	if not forwarded:
779	await refuse(429, "Too Many Requests", "slow down", extra=(("Retry-After", str(int(retry) + 1)),))
780	return # on a connection that is already carrying a request, closing is the only answer that cannot be woven into its response
781	if not forwarded:
782	dialed = await self._dial_agent()
783	if dialed is None:
784	await refuse(503, "Service Unavailable", "ALL is not answering here")
785	return
786	agent_reader, agent_writer = dialed
787	keepalive(agent_writer)
788	agent_task = asyncio.create_task(agent_to_client(agent_reader))
789	agent_task.add_done_callback(lambda task: agent_done.done() or agent_done.set_result(task))
790	first_good = head.path in AUTHENTICATED_PATHS and (head.path == "/token" or bearer)
791	forwarded += 1
792	agent_writer.write(segment.data)
793	await agent_writer.drain()
794
795	up = asyncio.create_task(client_to_agent())
796	try:
797	done, _ = await asyncio.wait({up, agent_done}, return_when=asyncio.FIRST_COMPLETED)
798	if up in done:
799	over = up.exception()
800	if isinstance(over, _Over):
801	{"age": self._count_age, "idle": self._count_idle, "head": self._count_head}[over.word]()
802	return agent_writer
803	if over is not None:
804	raise over
805	# the client finished its side (close_notify, or an end of the stream): say so to the agent and keep delivering what it is still answering, for a while
806	if agent_task is not None and not agent_task.done():
807	with contextlib.suppress(Exception):
808	if agent_writer.can_write_eof():
809	agent_writer.write_eof()
810	with contextlib.suppress(asyncio.TimeoutError, Exception):
811	await asyncio.wait_for(asyncio.shield(agent_task), policy.value("timing.half_close_s"))
812	return agent_writer
813	task = agent_done.result() # the agent ended its answer, or the connection: the stream is over
814	if not task.cancelled() and task.exception() is not None and not isinstance(task.exception(), (ConnectionError, OSError)):
815	raise task.exception()
816	return agent_writer
817	finally:
818	for task in (up, agent_task):
819	if task is not None and not task.done():
820	task.cancel()
821	await asyncio.gather(*(t for t in (up, agent_task) if t is not None), return_exceptions=True)
822	await close_writer(agent_writer) # also when the client's side ended with an exception: the caller never gets the writer then
823
824	def _count_age(self) -> None:
825	self.stats.age_closes += 1
826
827	def _count_idle(self) -> None:
828	self.stats.idle_closes += 1
829
830	def _count_head(self) -> None:
831	self.stats.head_timeouts += 1
832
833
834	class _Over(Exception):
835	"""A window ran out (the stream's age, its silence, or the time a request head had): the stream is closed. `word` says which."""
836
837	def __init__(self, word: str):
838	super().__init__(word)
839	self.word = word