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_tunnel.py
| 1 | """ALL 2.0, DEVICE side of the tunnel, protocol v2 (plan sections 2, 3 and 5; the contract, sections 2 and 3): ChatGPT speaks ordinary HTTPS to a name that |
| 2 | is this computer. |
| 3 | |
| 4 | A tunnel node (a small Go program on a public IP) reads only the SNI of a client's ClientHello and copies bytes. For every client of this computer it says ``OPEN`` on |
| 5 | the control connection this module keeps to it; the device dials the node back (``ATTACH``) and the raw stream, the client's own TLS, is handed to the terminator |
| 6 | (e2e_terminator.py), which ends it here with the certificate whose key only this computer has. |
| 7 | |
| 8 | What this module is, in the order a connection lives: |
| 9 | |
| 10 | * ``E2ETunnel`` keeps ONE control connection per node of the list the control plane gave this computer at enroll (``NodeLink``). Each one dials the node on its port with |
| 11 | TLS 1.3, ALPN ``all-tunnel/2`` and the node's public-CA certificate (certifi's bundle, the node's name, and an optional set of SPKI pins), and says |
| 12 | ``HELLO 2 <session>``; the session is a v2 session of audience ``all-tunnel`` that ``device_credentials`` mints (and re-mints, with the control plane's time, when this |
| 13 | computer's clock is off). The node answers ``CHALLENGE``; the device answers ``PROOF`` (a DPoP proof over the NODE's URL, so the proof is good nowhere else); the node |
| 14 | asks the control plane and answers ``OK``: the host this computer serves (it must be the one the control plane gave THIS computer at its own enroll: a node |
| 15 | cannot rename a device), the status (active, gated, paused), when the connection ends (``valid_until``) and the windows it offers (ping, dead peer, idle, age). |
| 16 | The device clamps every window to its own policy and never trusts an offer beyond it. |
| 17 | * Before ``valid_until`` minus ``timing.renew_before_valid_until_s`` it sends ``RENEW`` with a NEW session (and the node asks the control plane again: this is how an |
| 18 | ended trial becomes ``gated`` without waiting for an event). ``PING`` every ``timing.ping_interval_s``; a node that is silent for ``timing.dead_peer_s`` is |
| 19 | dead. ``DRAIN`` makes the device prefer the other node; ``STATEMENT`` (the control plane's signed access statement, forwarded) is verified, must not be older than |
| 20 | the last one taken, and is handed to ``device_access.ingest``. A refused connection is retried with full-jitter back-off from ``timing.backoff_min_s`` to |
| 21 | ``timing.backoff_max_s``; a wake from sleep or a change of network drops the connection at once and reconnects without back-off |
| 22 | (``timing.reconnect_after_wake_s`` is the bound the tests hold it to). |
| 23 | * Streams: the OPEN passes ``Admission`` (e2e_terminator.py) before the device does anything for it, then ``ATTACH`` on a new connection to the same node, then the |
| 24 | terminator. A good source address (one that just completed an authenticated call) is told to the node with ``HINT``. |
| 25 | |
| 26 | The tunnel never trusts the node with anything but availability: a node can cut or delay, and cannot rename the device, extend its own grant past what the control plane |
| 27 | says, make it dial anything but the agent's one proven port, or lift a paywall (only signed statements do). Nothing of a client (an address, a name) is ever logged. |
| 28 | |
| 29 | Nothing here runs on its own: all_mcp.e2e_runner starts it in a supervised child process of the relay service, after the enroll and the wave flag said so. |
| 30 | """ |
| 31 | from __future__ import annotations |
| 32 | |
| 33 | import asyncio |
| 34 | import base64 |
| 35 | import contextlib |
| 36 | import hashlib |
| 37 | import ipaddress |
| 38 | import json |
| 39 | import logging |
| 40 | import random |
| 41 | import re |
| 42 | import socket |
| 43 | import ssl |
| 44 | import time |
| 45 | from dataclasses import dataclass, field |
| 46 | from typing import Any, Callable |
| 47 | |
| 48 | import certifi |
| 49 | from cryptography import x509 |
| 50 | from cryptography.hazmat.primitives import serialization |
| 51 | |
| 52 | from . import e2e_policy as policy |
| 53 | from .e2e_terminator import Admission, EndpointBudgets, TLS_HANDSHAKE_BYTE, Terminator, close_writer, keepalive |
| 54 | |
| 55 | log = logging.getLogger("all_mcp.e2e_tunnel") |
| 56 | |
| 57 | # Constants of this module that the policy does not define: how long a dial and the wait for an OK may take, how long a connection must have lived to count as |
| 58 | # stable (a node that accepts and closes in a loop does not reset the back-off), the least time between two RENEWs, and how the clock and the network are watched. |
| 59 | DIAL_SECONDS = 10.0 |
| 60 | OK_WAIT_SECONDS = 30.0 |
| 61 | STABLE_SECONDS = 30.0 |
| 62 | RENEW_MIN_GAP_SECONDS = 5.0 |
| 63 | RENEW_RETRY_SECONDS = (5.0, 15.0) |
| 64 | CLOCK_WATCH_SECONDS = 1.0 |
| 65 | NETWORK_WATCH_SECONDS = 2.0 |
| 66 | WALL_JUMP_SECONDS = 3.0 |
| 67 | STALL_WAKE_SECONDS = 15.0 |
| 68 | RESYNC_SECONDS = 20.0 # the limit of one background read of the control plane's time |
| 69 | REPLACED_WAIT_SECONDS = (5.0, 30.0) |
| 70 | DRAIN_RECHECK_SECONDS = 1.0 # how often a node that said DRAIN is looked at again to see whether the other node still serves |
| 71 | STATEMENTS_IN_FLIGHT = 4 |
| 72 | # What a node can say (CLOSE <reason>) that means the control plane's answer about this computer has changed: the child reads its enroll again at once instead of |
| 73 | # waiting for its next turn (the kill switch must not take half an hour to reach a computer that is connected). |
| 74 | ATTENTION_WORDS = frozenset({"killed", "revoked", "released"}) |
| 75 | HINT_RESEND_DIVISOR = 3 |
| 76 | TUNNEL_PATH = policy.value("protocol.proof_htu_path") |
| 77 | _STATUSES = frozenset(policy.value("enums.authorize_status")) |
| 78 | _NODE_ID = re.compile(r"[a-z0-9]{1,16}") |
| 79 | _SECONDS = re.compile(r"(?:0|[1-9][0-9]{0,5})") |
| 80 | _STATEMENT = re.compile(r"[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+") |
| 81 | |
| 82 | |
| 83 | # ------------------------------------------------------------------------------------------------------------------------ errors |
| 84 | |
| 85 | class TunnelError(Exception): |
| 86 | """A connection that has to end. `word` is a fixed word (never text from the network) that the status file and the heartbeat may carry.""" |
| 87 | |
| 88 | def __init__(self, word: str): |
| 89 | super().__init__(word) |
| 90 | self.word = word |
| 91 | |
| 92 | |
| 93 | class ProtocolError(TunnelError): |
| 94 | """A line that is not in the grammar (or arrives where it should not).""" |
| 95 | |
| 96 | |
| 97 | class NodeRefused(TunnelError): |
| 98 | """The node answered ERR <code> (or closed with CLOSE <reason>).""" |
| 99 | |
| 100 | |
| 101 | class Kicked(Exception): |
| 102 | """The connection was dropped on purpose (a wake from sleep, a change of network): reconnect at once, no back-off.""" |
| 103 | |
| 104 | |
| 105 | # ------------------------------------------------------------------------------------------------------------------------ the lines |
| 106 | |
| 107 | def valid_ip(text: Any) -> bool: |
| 108 | """An address as the contract writes it: IPv4 dotted quad without leading zeros, or IPv6 in RFC 4291 text form, lower case, no zone, no brackets.""" |
| 109 | if not isinstance(text, str) or text != text.lower() or any(char in text for char in "%[] ") or not text.isascii(): |
| 110 | return False |
| 111 | try: |
| 112 | ipaddress.ip_address(text) |
| 113 | except ValueError: |
| 114 | return False |
| 115 | return True |
| 116 | |
| 117 | |
| 118 | def _bounded_json(text: str) -> Any: |
| 119 | def pairs(items: list[tuple[str, Any]]) -> dict[str, Any]: |
| 120 | keys = [key for key, _ in items] |
| 121 | if len(set(keys)) != len(keys): |
| 122 | raise ValueError("duplicate key") |
| 123 | return dict(items) |
| 124 | |
| 125 | def refuse(_: str) -> Any: |
| 126 | raise ValueError("not a plain number") |
| 127 | |
| 128 | return json.loads(text, object_pairs_hook=pairs, parse_float=refuse, parse_constant=refuse) |
| 129 | |
| 130 | |
| 131 | def _epoch(value: Any) -> bool: |
| 132 | return type(value) is int and 1 <= value <= 4102444800 |
| 133 | |
| 134 | |
| 135 | def _ok_body(text: str) -> dict[str, Any]: |
| 136 | try: |
| 137 | body = _bounded_json(text) |
| 138 | except (ValueError, RecursionError): |
| 139 | raise ProtocolError("bad_command") from None |
| 140 | windows = ("ping_interval_s", "dead_peer_s", "stream_idle_s", "stream_max_age_s", "renew_before_valid_until_s") |
| 141 | ok = (isinstance(body, dict) and set(body) == {"v", "node", "host", "status", "valid_until", "server_time", "policy"} |
| 142 | and type(body["v"]) is int and body["v"] == policy.value("protocol.version") |
| 143 | and isinstance(body["node"], str) and _NODE_ID.fullmatch(body["node"]) is not None |
| 144 | and isinstance(body["host"], str) and len(body["host"]) <= policy.value("identity.fqdn_max_length") |
| 145 | and policy.anchored("identity.fqdn_regex").fullmatch(body["host"]) is not None |
| 146 | and body["status"] in _STATUSES and _epoch(body["valid_until"]) and _epoch(body["server_time"]) |
| 147 | and isinstance(body["policy"], dict) and set(body["policy"]) == set(windows) |
| 148 | and all(type(body["policy"][name]) is int and 1 <= body["policy"][name] <= 86400 for name in windows)) |
| 149 | if not ok: |
| 150 | raise ProtocolError("bad_command") |
| 151 | return body |
| 152 | |
| 153 | |
| 154 | def parse_node_line(line: str) -> dict[str, Any]: |
| 155 | """One line the node sent (the text before the LF), as {"verb": ..., <field>: ...}. Raises ProtocolError. The grammar is contracts/schemas/wire-v2.grammar.json |
| 156 | (the tunnel client's tests run every n2d vector of contracts/vectors/lines.json through this function): a line is ASCII 0x20-0x7e, one space between fields, |
| 157 | no empty field, and it has its own maximum.""" |
| 158 | if not isinstance(line, str) or not line or any(not 0x20 <= ord(char) <= 0x7E for char in line): |
| 159 | raise ProtocolError("bad_command") |
| 160 | verb, space, rest = line.partition(" ") |
| 161 | if verb == "OPEN": |
| 162 | limit = policy.value("protocol.open_line_max_bytes") |
| 163 | elif verb == "STATEMENT": |
| 164 | limit = policy.value("protocol.statement_line_max_bytes") |
| 165 | else: |
| 166 | limit = policy.value("protocol.line_max_bytes") |
| 167 | if len(line) + 1 > limit: |
| 168 | raise ProtocolError("line_too_long") |
| 169 | if space and not rest: |
| 170 | raise ProtocolError("bad_command") |
| 171 | if verb == "OK": |
| 172 | if not space: |
| 173 | raise ProtocolError("bad_command") |
| 174 | return {"verb": verb, "body": _ok_body(rest)} |
| 175 | fields = rest.split(" ") if space else [] |
| 176 | if any(not field_ for field_ in fields): |
| 177 | raise ProtocolError("bad_command") |
| 178 | out: dict[str, Any] = {"verb": verb} |
| 179 | if verb == "CHALLENGE" and len(fields) == 1 and policy.anchored("identity.nonce_regex").fullmatch(fields[0]): |
| 180 | out["nonce"] = fields[0] |
| 181 | elif verb == "ERR" and len(fields) == 1 and fields[0] in policy.value("protocol.error_codes"): |
| 182 | out["code"] = fields[0] |
| 183 | elif (verb == "OPEN" and len(fields) == 4 and policy.anchored("identity.stream_id_regex").fullmatch(fields[0]) |
| 184 | and policy.anchored("identity.attach_secret_regex").fullmatch(fields[1]) and valid_ip(fields[2]) |
| 185 | and fields[3] in policy.value("protocol.source_classes")): |
| 186 | out.update(stream_id=fields[0], secret=fields[1], peer_ip=fields[2], cls=fields[3]) |
| 187 | elif verb == "PONG" and not fields: |
| 188 | pass |
| 189 | elif verb == "STATEMENT" and len(fields) == 1 and _STATEMENT.fullmatch(fields[0]): |
| 190 | out["statement"] = fields[0] |
| 191 | elif verb == "DRAIN" and len(fields) == 1 and _SECONDS.fullmatch(fields[0]): |
| 192 | out["quiet_s"] = int(fields[0]) |
| 193 | elif verb == "CLOSE" and len(fields) <= 1 and (not fields or fields[0] in policy.value("protocol.close_reasons")): |
| 194 | out["reason"] = fields[0] if fields else "" |
| 195 | else: |
| 196 | raise ProtocolError("bad_command") |
| 197 | return out |
| 198 | |
| 199 | |
| 200 | def _token_line(verb: str, token: str) -> str: |
| 201 | if (not isinstance(token, str) or not token.isascii() or len(token) > policy.value("protocol.jwt_max_bytes") |
| 202 | or _STATEMENT.fullmatch(token) is None): |
| 203 | raise ProtocolError("bad_command") |
| 204 | return f"{verb} {token}\n" if verb != "HELLO" else f"HELLO {policy.value('protocol.version')} {token}\n" |
| 205 | |
| 206 | |
| 207 | def hello_line(session: str) -> str: |
| 208 | return _token_line("HELLO", session) |
| 209 | |
| 210 | |
| 211 | def renew_line(session: str) -> str: |
| 212 | return _token_line("RENEW", session) |
| 213 | |
| 214 | |
| 215 | def proof_line(proof: str) -> str: |
| 216 | return _token_line("PROOF", proof) |
| 217 | |
| 218 | |
| 219 | def attach_line(stream_id: str, secret: str) -> str: |
| 220 | if (policy.anchored("identity.stream_id_regex").fullmatch(stream_id) is None |
| 221 | or policy.anchored("identity.attach_secret_regex").fullmatch(secret) is None): |
| 222 | raise ProtocolError("bad_command") |
| 223 | return f"ATTACH {stream_id} {secret}\n" |
| 224 | |
| 225 | |
| 226 | def hint_line(addresses: list[str]) -> str: |
| 227 | if not 1 <= len(addresses) <= policy.value("protocol.hint_max_addresses") or not all(valid_ip(address) for address in addresses): |
| 228 | raise ProtocolError("bad_command") |
| 229 | return "HINT " + " ".join(addresses) + "\n" |
| 230 | |
| 231 | |
| 232 | def close_line(reason: str = "") -> str: |
| 233 | if reason and reason not in policy.value("protocol.close_reasons"): |
| 234 | raise ProtocolError("bad_command") |
| 235 | return f"CLOSE {reason}\n" if reason else "CLOSE\n" |
| 236 | |
| 237 | |
| 238 | # ------------------------------------------------------------------------------------------------------------------------ windows |
| 239 | |
| 240 | @dataclass(frozen=True) |
| 241 | class Windows: |
| 242 | """The windows a connection works with: this device's own policy, or what a node offered clamped to it (never above).""" |
| 243 | ping_interval: float |
| 244 | dead_peer: float |
| 245 | stream_idle: float |
| 246 | stream_max_age: float |
| 247 | renew_before: float |
| 248 | |
| 249 | @classmethod |
| 250 | def own(cls) -> "Windows": |
| 251 | return cls(float(policy.value("timing.ping_interval_s")), float(policy.value("timing.dead_peer_s")), float(policy.value("timing.stream_idle_s")), |
| 252 | float(policy.value("timing.stream_max_age_s")), float(policy.value("timing.renew_before_valid_until_s"))) |
| 253 | |
| 254 | def clamped(self, offered: dict[str, int]) -> "Windows": |
| 255 | """What the node offered in its OK, held to this device's own windows: each is the smaller of the two (and at least one second, or this device's own window |
| 256 | when that is shorter: a window of nothing would close every stream at once). A dead-peer window that is not longer than the ping interval would drop a healthy |
| 257 | connection between two pings, so it is raised to twice the ping interval, still never above this device's own.""" |
| 258 | def held(name: str, own: float) -> float: |
| 259 | return max(min(1.0, own), min(float(offered[name]), own)) |
| 260 | |
| 261 | ping = held("ping_interval_s", self.ping_interval) |
| 262 | dead = held("dead_peer_s", self.dead_peer) |
| 263 | if dead <= ping: |
| 264 | dead = min(self.dead_peer, ping * 2) |
| 265 | ping = min(ping, dead / 2) |
| 266 | return Windows(ping, dead, held("stream_idle_s", self.stream_idle), held("stream_max_age_s", self.stream_max_age), |
| 267 | held("renew_before_valid_until_s", self.renew_before)) |
| 268 | |
| 269 | |
| 270 | # ------------------------------------------------------------------------------------------------------------------------ nodes and their certificates |
| 271 | |
| 272 | @dataclass(frozen=True) |
| 273 | class NodeSpec: |
| 274 | """A node of the list the control plane gave this computer: its id, name, port and optional SPKI pins (base64url SHA-256 of the SubjectPublicKeyInfo). |
| 275 | `address` (tests and a node reached by a fixed address) is where to connect; the name is still the one the certificate must carry.""" |
| 276 | id: str |
| 277 | host: str |
| 278 | port: int |
| 279 | pins: frozenset[str] = frozenset() |
| 280 | address: str | None = None |
| 281 | |
| 282 | @property |
| 283 | def url(self) -> str: |
| 284 | """The node's own URL: what the DPoP proof is bound to (``https://<host>[:<port> unless 443]/all-tunnel/2``).""" |
| 285 | return f"https://{self.host}{'' if self.port == 443 else ':' + str(self.port)}{TUNNEL_PATH}" |
| 286 | |
| 287 | @classmethod |
| 288 | def from_json(cls, value: dict[str, Any]) -> "NodeSpec": |
| 289 | return cls(value["id"], value["host"], int(value["port"]), frozenset(value.get("pins") or ())) |
| 290 | |
| 291 | |
| 292 | def spki_pin(der: bytes) -> str: |
| 293 | """The pin of a certificate: base64url (no padding) of the SHA-256 of its SubjectPublicKeyInfo, so a renewal with the same key keeps the pin.""" |
| 294 | public = x509.load_der_x509_certificate(der).public_key() |
| 295 | spki = public.public_bytes(serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo) |
| 296 | return base64.urlsafe_b64encode(hashlib.sha256(spki).digest()).rstrip(b"=").decode("ascii") |
| 297 | |
| 298 | |
| 299 | def node_client_context() -> ssl.SSLContext: |
| 300 | """TLS to a node: its certificate must chain to a public root (certifi's bundle: fresh or lagging machines see the same roots as the rest of ALL) and carry the |
| 301 | node's name, TLS 1.3 at least, and the ALPN that says this is the tunnel.""" |
| 302 | context = ssl.create_default_context(cafile=certifi.where()) |
| 303 | context.minimum_version = ssl.TLSVersion.TLSv1_3 |
| 304 | context.set_alpn_protocols([policy.value("protocol.alpn")]) |
| 305 | return context |
| 306 | |
| 307 | |
| 308 | # ------------------------------------------------------------------------------------------------------------------------ the network, as far as a change can be seen |
| 309 | |
| 310 | # Interfaces a Mac raises and drops on its own (AirDrop's awdl/llw, VPN and iCloud Private Relay utun, the anpi/ap/bridge/gif/stf helpers): none of them is a move between |
| 311 | # networks, and counting them reconnected every tunnel each time AirDrop woke up. |
| 312 | VIRTUAL_INTERFACE_PREFIXES = ("lo", "loopback", "awdl", "llw", "utun", "anpi", "ap", "bridge", "gif", "stf") |
| 313 | |
| 314 | |
| 315 | def default_network_snapshot() -> tuple: |
| 316 | """What changes when the computer moves between networks: the address the default routes would send from (a UDP connect sends nothing) and the interfaces that are |
| 317 | up. Cheap enough to take every couple of seconds.""" |
| 318 | routes = [] |
| 319 | for family, probe in ((socket.AF_INET, "192.0.2.1"), (socket.AF_INET6, "2001:db8::1")): |
| 320 | try: |
| 321 | with socket.socket(family, socket.SOCK_DGRAM) as probe_socket: |
| 322 | probe_socket.connect((probe, 9)) |
| 323 | routes.append(probe_socket.getsockname()[0]) |
| 324 | except OSError: |
| 325 | routes.append("") |
| 326 | try: |
| 327 | import psutil |
| 328 | up = tuple(sorted(name for name, stats in psutil.net_if_stats().items() if stats.isup and not name.lower().startswith(VIRTUAL_INTERFACE_PREFIXES))) |
| 329 | except Exception: # noqa: BLE001 - the routes alone are a good enough signal |
| 330 | up = () |
| 331 | return (*routes, up) |
| 332 | |
| 333 | |
| 334 | # ------------------------------------------------------------------------------------------------------------------------ one control connection per node |
| 335 | |
| 336 | @dataclass |
| 337 | class _Conn: |
| 338 | """What one control connection is doing right now.""" |
| 339 | reader: asyncio.StreamReader |
| 340 | writer: asyncio.StreamWriter |
| 341 | windows: Windows |
| 342 | exchange: str = "hello" # which exchange is under way: "hello" (first) or "renew" |
| 343 | awaiting: str | None = "challenge" # what the node owes: "challenge", then "ok"; None between exchanges |
| 344 | exchange_deadline: float = 0.0 |
| 345 | session: str = "" # the session of the exchange under way |
| 346 | ready: bool = False |
| 347 | heard: float = 0.0 |
| 348 | ping_at: float = 0.0 |
| 349 | renew_at: float | None = None |
| 350 | valid_server: float = 0.0 # valid_until, on the control plane's clock |
| 351 | fetch: asyncio.Future | None = None |
| 352 | ready_at: float = 0.0 |
| 353 | close_reason: str = "" |
| 354 | grace_ok: bool = False # the last RENEW was answered "control_unavailable": the control plane, not the node, is what does not answer |
| 355 | |
| 356 | |
| 357 | class NodeLink: |
| 358 | """The control connection to one node, kept for as long as the tunnel runs: connect, authenticate, serve (OPEN, PING, RENEW, STATEMENT, DRAIN), and connect again |
| 359 | when it ends.""" |
| 360 | |
| 361 | def __init__(self, tunnel: "E2ETunnel", spec: NodeSpec): |
| 362 | self.tunnel = tunnel |
| 363 | self.spec = spec |
| 364 | self.state = "connecting" # connecting | ready | draining | down |
| 365 | self.status = "" # what the node last said about this computer: active | gated | paused |
| 366 | self.valid_until = 0 # the control plane's moment, from the last OK |
| 367 | self.last_error = "" |
| 368 | self.connects = 0 |
| 369 | self.failures = 0 |
| 370 | self.renewals = 0 |
| 371 | self.drain_until = 0.0 # monotonic |
| 372 | self._writer: asyncio.StreamWriter | None = None |
| 373 | self._kick = asyncio.Event() |
| 374 | self._streams: set[asyncio.Task] = set() |
| 375 | self._task: asyncio.Task | None = None |
| 376 | self._hinted: dict[str, float] = {} |
| 377 | self._stopping = False |
| 378 | |
| 379 | # -- life |
| 380 | def start(self) -> None: |
| 381 | if self._task is None: |
| 382 | self._task = asyncio.create_task(self.run(), name=f"e2e-node-{self.spec.id}") |
| 383 | |
| 384 | async def stop(self) -> None: |
| 385 | self._stopping = True |
| 386 | self._kick.set() |
| 387 | if self._writer is not None: |
| 388 | with contextlib.suppress(Exception): |
| 389 | self._writer.write(close_line("shutdown").encode("ascii")) |
| 390 | if self._task is not None: |
| 391 | self._task.cancel() |
| 392 | await asyncio.gather(self._task, return_exceptions=True) |
| 393 | for task in list(self._streams): |
| 394 | task.cancel() |
| 395 | await asyncio.gather(*self._streams, return_exceptions=True) |
| 396 | |
| 397 | def kick(self) -> None: |
| 398 | """Drop the connection now and reconnect at once (a wake from sleep or a change of network: the connection is probably dead and only a timer would say so).""" |
| 399 | self._kick.set() |
| 400 | |
| 401 | @property |
| 402 | def ready(self) -> bool: |
| 403 | return self.state in ("ready", "draining") |
| 404 | |
| 405 | # -- the reconnect loop |
| 406 | def _delay(self, attempt: int) -> float: |
| 407 | """Full-jitter back-off: somewhere between the policy's minimum and a ceiling that doubles with every failure up to its maximum.""" |
| 408 | low, high = self.tunnel.backoff |
| 409 | return random.uniform(low, max(low, min(high, low * 2 ** (attempt + 1)))) |
| 410 | |
| 411 | async def run(self) -> None: |
| 412 | attempt = 0 |
| 413 | while not self._stopping: |
| 414 | began = time.monotonic() |
| 415 | delay: float | None = None |
| 416 | try: |
| 417 | await self._connection() |
| 418 | except asyncio.CancelledError: |
| 419 | raise |
| 420 | except Kicked: |
| 421 | self.last_error = "network_changed" |
| 422 | delay, attempt = random.uniform(0.0, 1.0), 0 |
| 423 | except NodeRefused as refused: |
| 424 | self.last_error = refused.word |
| 425 | delay = self._refusal_delay(refused.word, attempt) |
| 426 | if refused.word in ATTENTION_WORDS: |
| 427 | self.tunnel.attention.set() |
| 428 | if refused.word == "clock_skew": |
| 429 | with contextlib.suppress(Exception): |
| 430 | await asyncio.to_thread(self.tunnel.credentials.sync_clock, True) # the difference with the control plane is learned again |
| 431 | except TunnelError as error: |
| 432 | self.last_error = error.word |
| 433 | except ssl.SSLCertVerificationError: |
| 434 | self.last_error = "certificate_refused" # the node's certificate is not for its name, has expired or has no public root |
| 435 | except (ssl.SSLError, ConnectionError, OSError, asyncio.TimeoutError, asyncio.IncompleteReadError, asyncio.LimitOverrunError) as error: |
| 436 | self.last_error = "timeout" if isinstance(error, asyncio.TimeoutError) else ("tls_error" if isinstance(error, ssl.SSLError) else "connect_failed") |
| 437 | except Exception as error: # noqa: BLE001 - what a credential or a storage problem says is a closed word, never its text |
| 438 | self.last_error = self.tunnel.credential_word(error) |
| 439 | delay = self._refusal_delay(self.last_error, attempt) # a revoked computer or a license that is over is not asked again every second |
| 440 | if self.last_error == "internal_error": |
| 441 | log.warning("tunnel node %s: %s", self.spec.id, type(error).__name__) |
| 442 | finally: |
| 443 | was_ready = self.state in ("ready", "draining") |
| 444 | self.state = "down" |
| 445 | self._writer = None |
| 446 | self.tunnel.link_ended(self) |
| 447 | if self._stopping: |
| 448 | return |
| 449 | self.failures += 1 |
| 450 | stable = was_ready and time.monotonic() - began >= STABLE_SECONDS |
| 451 | if stable: |
| 452 | attempt = 0 |
| 453 | if delay is None: |
| 454 | delay = self._delay(attempt) |
| 455 | attempt += 1 |
| 456 | self.state = "connecting" |
| 457 | await self._wait_to_reconnect(delay) |
| 458 | self._kick.clear() |
| 459 | |
| 460 | async def _wait_to_reconnect(self, delay: float) -> None: |
| 461 | """Wait `delay` before dialing again. A node that said DRAIN is left alone for its quiet period while the other node serves, and that is looked at again every |
| 462 | DRAIN_RECHECK_SECONDS: when the other node goes down too, this one is dialed at once rather than leaving the computer with no connection. A kick (wake, network |
| 463 | change) ends the wait at any time.""" |
| 464 | until = time.monotonic() + delay |
| 465 | while True: |
| 466 | now = time.monotonic() |
| 467 | if self.drain_until > now and self.tunnel.other_node_ready(self): |
| 468 | left = max(until, self.drain_until) - now # the node said it is going down: the other node serves while it does |
| 469 | else: |
| 470 | left = until - now |
| 471 | if left <= 0: |
| 472 | return |
| 473 | try: |
| 474 | await asyncio.wait_for(self._kick.wait(), min(left, DRAIN_RECHECK_SECONDS)) |
| 475 | return |
| 476 | except asyncio.TimeoutError: |
| 477 | continue |
| 478 | |
| 479 | def _refusal_delay(self, word: str, attempt: int) -> float | None: |
| 480 | """After a refusal: the plain back-off for a hiccup, the ceiling for an answer that will not change by itself (a revoked computer, a license that is over, a |
| 481 | node that names this computer something else) and a short wait for a node that closed because another connection of this computer replaced it.""" |
| 482 | low, high = self.tunnel.backoff |
| 483 | if word in ("device_revoked", "license_inactive", "revoked", "released", "killed", "host_mismatch"): |
| 484 | return random.uniform(high / 2, high) |
| 485 | if word == "replaced": |
| 486 | return random.uniform(*REPLACED_WAIT_SECONDS) |
| 487 | return None |
| 488 | |
| 489 | # -- dialing |
| 490 | async def dial(self, *, limit: int = 1 << 16) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: |
| 491 | """TLS to the node, verified before a single byte is sent: its certificate chains to a public root and carries its name (the context), is on the pin list when |
| 492 | there is one, and the connection is TLS 1.3 with the ALPN of the tunnel.""" |
| 493 | spec = self.spec |
| 494 | reader, writer = await asyncio.wait_for( |
| 495 | asyncio.open_connection(spec.address or spec.host, spec.port, ssl=self.tunnel.client_context(), server_hostname=spec.host, |
| 496 | ssl_handshake_timeout=DIAL_SECONDS, limit=limit), DIAL_SECONDS + 5) |
| 497 | try: |
| 498 | ssl_object = writer.get_extra_info("ssl_object") |
| 499 | if ssl_object is None or ssl_object.version() != "TLSv1.3" or ssl_object.selected_alpn_protocol() != policy.value("protocol.alpn"): |
| 500 | raise TunnelError("tls_error") |
| 501 | if spec.pins: |
| 502 | der = ssl_object.getpeercert(binary_form=True) |
| 503 | if not der or spki_pin(der) not in spec.pins: |
| 504 | raise TunnelError("pin_mismatch") # nothing was sent: the session never leaves this computer for a certificate that is not the pinned one |
| 505 | except BaseException: |
| 506 | await close_writer(writer) |
| 507 | raise |
| 508 | keepalive(writer) |
| 509 | return reader, writer |
| 510 | |
| 511 | # -- one connection |
| 512 | async def _write(self, conn: _Conn, line: str) -> None: |
| 513 | conn.writer.write(line.encode("ascii")) |
| 514 | await asyncio.wait_for(conn.writer.drain(), conn.windows.dead_peer) |
| 515 | |
| 516 | async def _connection(self) -> None: |
| 517 | tunnel = self.tunnel |
| 518 | reader, writer = await self.dial(limit=policy.value("protocol.statement_line_max_bytes") + 16) |
| 519 | self._writer = writer |
| 520 | conn = _Conn(reader, writer, tunnel.windows_own, heard=time.monotonic()) |
| 521 | try: |
| 522 | conn.session = await tunnel.session(force=False) |
| 523 | await self._write(conn, hello_line(conn.session)) |
| 524 | conn.exchange_deadline = time.monotonic() + OK_WAIT_SECONDS |
| 525 | await self._serve(conn) |
| 526 | finally: |
| 527 | if conn.fetch is not None and not conn.fetch.done(): |
| 528 | conn.fetch.cancel() |
| 529 | await close_writer(writer) |
| 530 | |
| 531 | def _timers(self, conn: _Conn, now: float) -> float: |
| 532 | """Seconds until the next thing this connection has to do on its own.""" |
| 533 | times = [conn.heard + conn.windows.dead_peer] |
| 534 | if conn.awaiting: |
| 535 | times.append(conn.exchange_deadline) |
| 536 | if conn.ready: |
| 537 | times.append(conn.ping_at) |
| 538 | if not conn.awaiting: # while a RENEW is under way the node decides (and the exchange has its own deadline) |
| 539 | times.append(self._local_valid_until(conn) + (self._grace_seconds() if conn.grace_ok else 0.0)) |
| 540 | if conn.renew_at is not None and conn.fetch is None: |
| 541 | times.append(conn.renew_at) |
| 542 | return max(0.0, min(times) - now) |
| 543 | |
| 544 | @staticmethod |
| 545 | def _grace_seconds() -> float: |
| 546 | """How long past valid_until a connection may live while the control plane does not answer: the node's own grace for a paid computer (policy authorization.grace_paid_s).""" |
| 547 | return float(policy.value("authorization.grace_paid_s")) |
| 548 | |
| 549 | def _local_valid_until(self, conn: _Conn) -> float: |
| 550 | """valid_until as a reading of the monotonic clock: the control plane's moment minus its time now, plus now.""" |
| 551 | return time.monotonic() + (conn.valid_server - self.tunnel.now()) |
| 552 | |
| 553 | async def _serve(self, conn: _Conn) -> None: |
| 554 | read_task: asyncio.Task | None = None |
| 555 | kick_task: asyncio.Task | None = None |
| 556 | try: |
| 557 | while True: |
| 558 | if read_task is None: |
| 559 | read_task = asyncio.ensure_future(conn.reader.readuntil(b"\n")) |
| 560 | if kick_task is None or kick_task.done(): |
| 561 | kick_task = asyncio.ensure_future(self._kick.wait()) |
| 562 | waiting = {read_task, kick_task} | ({conn.fetch} if conn.fetch is not None else set()) |
| 563 | now = time.monotonic() |
| 564 | done, _ = await asyncio.wait(waiting, timeout=self._timers(conn, now), return_when=asyncio.FIRST_COMPLETED) |
| 565 | now = time.monotonic() |
| 566 | if kick_task in done: |
| 567 | self._kick.clear() |
| 568 | raise Kicked |
| 569 | if self._stopping: |
| 570 | return |
| 571 | if conn.fetch is not None and conn.fetch in done: |
| 572 | fetch, conn.fetch = conn.fetch, None |
| 573 | await self._start_renew(conn, fetch) |
| 574 | if read_task in done: |
| 575 | task, read_task = read_task, None |
| 576 | try: |
| 577 | raw = task.result() |
| 578 | except asyncio.IncompleteReadError: |
| 579 | raise TunnelError("closed_by_node") from None |
| 580 | except asyncio.LimitOverrunError: |
| 581 | raise ProtocolError("line_too_long") from None |
| 582 | conn.heard = now |
| 583 | await self._line(conn, raw[:-1].decode("ascii", "replace"), now) |
| 584 | await self._tick(conn, now) |
| 585 | finally: |
| 586 | for task in (read_task, kick_task): |
| 587 | if task is not None and not task.done(): |
| 588 | task.cancel() |
| 589 | await asyncio.gather(*(t for t in (read_task, kick_task) if t is not None), return_exceptions=True) |
| 590 | |
| 591 | async def _tick(self, conn: _Conn, now: float) -> None: |
| 592 | if now - conn.heard >= conn.windows.dead_peer: |
| 593 | raise TunnelError("dead_peer") |
| 594 | if conn.awaiting and now >= conn.exchange_deadline: |
| 595 | if not conn.ready: |
| 596 | raise TunnelError("handshake_timeout") |
| 597 | self._renew_failed(conn, now) # a RENEW the node never answered: the connection is still good until valid_until |
| 598 | if not conn.ready: |
| 599 | return |
| 600 | over = now - self._local_valid_until(conn) |
| 601 | if over >= 0 and not conn.awaiting and not (conn.grace_ok and over < self._grace_seconds()): |
| 602 | # The grant is over. A computer whose last RENEW got "control_unavailable" keeps its connection for the paid grace (the node decides who gets it: a trial or gated |
| 603 | # computer is closed by the node at valid_until whatever this side does, and a refusal of the control plane is never "the control plane is down"): closing here, at |
| 604 | # valid_until, made the grace useless: the device redialled, the node needed a good authorize for the new HELLO and the control plane was down. |
| 605 | raise TunnelError("expired") |
| 606 | if now >= conn.ping_at: |
| 607 | await self._write(conn, "PING\n") |
| 608 | conn.ping_at = now + conn.windows.ping_interval |
| 609 | if conn.renew_at is not None and now >= conn.renew_at and conn.fetch is None and not conn.awaiting: |
| 610 | conn.renew_at = None |
| 611 | conn.fetch = asyncio.ensure_future(self.tunnel.session(force=True)) |
| 612 | |
| 613 | async def _start_renew(self, conn: _Conn, fetch: asyncio.Future) -> None: |
| 614 | try: |
| 615 | session = fetch.result() |
| 616 | except asyncio.CancelledError: |
| 617 | raise |
| 618 | except Exception as error: # noqa: BLE001 |
| 619 | self.last_error = self.tunnel.credential_word(error) |
| 620 | conn.grace_ok = self.last_error == "control_unreachable" # the control plane did not answer this computer's own request for a session: that is the outage the grace is for (a refusal is not) |
| 621 | self._renew_failed(conn, time.monotonic()) |
| 622 | return |
| 623 | conn.session, conn.exchange, conn.awaiting = session, "renew", "challenge" |
| 624 | conn.exchange_deadline = time.monotonic() + OK_WAIT_SECONDS |
| 625 | await self._write(conn, renew_line(session)) |
| 626 | |
| 627 | def _renew_failed(self, conn: _Conn, now: float) -> None: |
| 628 | """A renewal that did not come through: the old valid_until still holds, so try again soon, and well before it.""" |
| 629 | conn.awaiting = None |
| 630 | left = self._local_valid_until(conn) - now |
| 631 | retry = random.uniform(*RENEW_RETRY_SECONDS) |
| 632 | conn.renew_at = now + (min(retry, max(0.5, left / 2)) if left > 0 else retry) # past valid_until (the grace) the retry keeps its normal pace, never a storm |
| 633 | |
| 634 | async def _line(self, conn: _Conn, line: str, now: float) -> None: |
| 635 | message = parse_node_line(line) |
| 636 | verb = message["verb"] |
| 637 | if verb == "CHALLENGE": |
| 638 | if conn.awaiting != "challenge": |
| 639 | raise ProtocolError("unexpected_challenge") |
| 640 | proof = await self.tunnel.proof(conn.session, message["nonce"], self.spec.url) |
| 641 | conn.awaiting = "ok" |
| 642 | conn.exchange_deadline = now + OK_WAIT_SECONDS |
| 643 | await self._write(conn, proof_line(proof)) |
| 644 | elif verb == "OK": |
| 645 | if conn.awaiting != "ok": |
| 646 | raise ProtocolError("unexpected_ok") |
| 647 | self._on_ok(conn, message["body"], now) |
| 648 | elif verb == "ERR": |
| 649 | if conn.ready: |
| 650 | self.last_error = message["code"] |
| 651 | if message["code"] == "control_unavailable": |
| 652 | conn.grace_ok = True # "the control plane did not answer" is what earns the paid grace |
| 653 | elif message["code"] != "busy": |
| 654 | conn.grace_ok = False # any other word is an answer of the control plane (a refusal): no grace. "busy" (the node's own pacing) says nothing about the control plane |
| 655 | self._renew_failed(conn, now) # a refused RENEW keeps the connection and its old valid_until |
| 656 | else: |
| 657 | raise NodeRefused(message["code"]) |
| 658 | elif verb == "OPEN": |
| 659 | if conn.ready: |
| 660 | self.tunnel.on_open(self, message) |
| 661 | elif verb == "PONG": |
| 662 | return |
| 663 | elif verb == "STATEMENT": |
| 664 | self.tunnel.on_statement(message["statement"]) |
| 665 | elif verb == "DRAIN": |
| 666 | self.drain_until = now + min(message["quiet_s"], policy.value("timing.drain_quiet_s")) # the grammar allows far more than the policy |
| 667 | if self.state == "ready": |
| 668 | self.state = "draining" |
| 669 | elif verb == "CLOSE": |
| 670 | raise NodeRefused(message["reason"] or "closed_by_node") |
| 671 | else: |
| 672 | raise ProtocolError("bad_command") # a verb the device only SENDS (HELLO, PING...) coming from the node |
| 673 | |
| 674 | def _on_ok(self, conn: _Conn, body: dict[str, Any], now: float) -> None: |
| 675 | if body["host"] != self.tunnel.host: |
| 676 | raise TunnelError("host_mismatch") # the node says this computer is somebody else's name: it is not believed, nothing is served for it |
| 677 | server_now = self.tunnel.now() |
| 678 | limit = server_now + policy.value("authorization.valid_until_paid_s") + policy.value("authorization.node_skew_max_s") |
| 679 | valid = min(float(body["valid_until"]), limit) |
| 680 | if valid <= server_now: |
| 681 | raise TunnelError("grant_expired") |
| 682 | conn.windows = self.tunnel.windows_own.clamped(body["policy"]) |
| 683 | conn.valid_server = valid |
| 684 | conn.grace_ok = False |
| 685 | conn.awaiting = None |
| 686 | self.status, self.valid_until = body["status"], int(valid) |
| 687 | if conn.exchange == "renew": |
| 688 | self.renewals += 1 |
| 689 | if not conn.ready: |
| 690 | conn.ready, conn.ready_at = True, now |
| 691 | # A draining node refuses a new HELLO and sends DRAIN right after the OK of one that was already authenticating (tunnel/server.go followUp), so |
| 692 | # an OK on a new connection is a node that is not draining now, whatever the previous connection was told: a DRAIN that follows sets it again. |
| 693 | self.drain_until = 0.0 |
| 694 | self.state = "ready" |
| 695 | self.connects += 1 |
| 696 | self.last_error = "" |
| 697 | conn.ping_at = now + conn.windows.ping_interval |
| 698 | self.tunnel.link_ready(self) |
| 699 | self.tunnel.windows_changed(conn.windows) |
| 700 | wait = valid - server_now - conn.windows.renew_before |
| 701 | conn.renew_at = now + max(wait, min(RENEW_MIN_GAP_SECONDS, (valid - server_now) / 2)) + random.uniform(0, min(10.0, max(0.0, wait) * 0.05)) |
| 702 | |
| 703 | # -- what streams and statements ask of the control connection |
| 704 | def send_hint(self, addresses: list[str]) -> None: |
| 705 | """Tell the node about client addresses that just completed authenticated calls (each at most once in a third of the node's memory of them).""" |
| 706 | writer = self._writer |
| 707 | if writer is None or self.state not in ("ready", "draining"): |
| 708 | return |
| 709 | now = time.monotonic() |
| 710 | fresh = [address for address in addresses if now - self._hinted.get(address, float("-inf")) >= policy.value("node_limits.hint_ttl_s") / HINT_RESEND_DIVISOR] |
| 711 | for start in range(0, len(fresh), policy.value("protocol.hint_max_addresses")): |
| 712 | batch = fresh[start:start + policy.value("protocol.hint_max_addresses")] |
| 713 | try: |
| 714 | writer.write(hint_line(batch).encode("ascii")) |
| 715 | except (ProtocolError, OSError): |
| 716 | return |
| 717 | for address in batch: |
| 718 | self._hinted[address] = now |
| 719 | while len(self._hinted) > policy.value("surface.limiter_lru_max"): |
| 720 | self._hinted.pop(next(iter(self._hinted))) |
| 721 | |
| 722 | def track(self, task: asyncio.Task) -> None: |
| 723 | self._streams.add(task) |
| 724 | task.add_done_callback(self._streams.discard) |
| 725 | |
| 726 | def snapshot(self) -> dict[str, Any]: |
| 727 | return {"id": self.spec.id, "state": self.state, "status": self.status, "valid_until": self.valid_until or None, "last_error": self.last_error, |
| 728 | "connects": self.connects, "failures": self.failures, "renewals": self.renewals, "streams": len(self._streams)} |
| 729 | |
| 730 | |
| 731 | # ------------------------------------------------------------------------------------------------------------------------ the tunnel |
| 732 | |
| 733 | def _statement_seq(token: str) -> int | None: |
| 734 | """The `seq` of a statement that was already verified (the signature is what makes it believable; this only reads it).""" |
| 735 | try: |
| 736 | payload = token.split(".")[1] |
| 737 | claims = json.loads(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))) |
| 738 | seq = claims.get("seq") |
| 739 | return seq if type(seq) is int and seq >= 1 else None |
| 740 | except (IndexError, ValueError, UnicodeError, AttributeError): |
| 741 | return None |
| 742 | |
| 743 | |
| 744 | @dataclass |
| 745 | class TunnelStats: |
| 746 | streams_opened: int = 0 |
| 747 | streams_refused: dict[str, int] = field(default_factory=dict) |
| 748 | attach_failures: int = 0 |
| 749 | statements_taken: int = 0 |
| 750 | statements_refused: int = 0 |
| 751 | wakes: int = 0 |
| 752 | network_changes: int = 0 |
| 753 | |
| 754 | def as_dict(self) -> dict[str, Any]: |
| 755 | return dict(self.__dict__) |
| 756 | |
| 757 | |
| 758 | class E2ETunnel: |
| 759 | """The device's side of the tunnel: a control connection to every node, the terminator for every stream, and the watchers of the clock and the network. |
| 760 | |
| 761 | `host` is the name this computer serves (from its own enroll); `nodes` the list the control plane gave it; `credentials` is a device_credentials.DeviceCredentials (or |
| 762 | anything with ``tunnel_session``, ``tunnel_proof``, ``now``, ``sync_clock``, ``verify_access`` and ``clock``); `serving` a tls_store.ServingContext. `ingest` takes a |
| 763 | verified statement (device_access.ingest). Everything else exists so that tests can drive it without a network, a clock or an agent.""" |
| 764 | |
| 765 | def __init__(self, *, host: str, nodes: list[NodeSpec], credentials: Any, serving: Any, ingest: Callable[[str], bool] | None = None, |
| 766 | client_context: Callable[[], ssl.SSLContext] = node_client_context, own_windows: Windows | None = None, |
| 767 | backoff: tuple[float, float] | None = None, admission: Admission | None = None, budgets: EndpointBudgets | None = None, |
| 768 | terminator: Terminator | None = None, terminator_options: dict[str, Any] | None = None, |
| 769 | network_snapshot: Callable[[], Any] = default_network_snapshot, wall: Callable[[], float] = time.time, |
| 770 | monotonic: Callable[[], float] = time.monotonic): |
| 771 | self.host = host |
| 772 | self.credentials = credentials |
| 773 | self._client_context = client_context |
| 774 | self.windows_own = own_windows or Windows.own() |
| 775 | self._windows = self.windows_own |
| 776 | self.backoff = backoff or (float(policy.value("timing.backoff_min_s")), float(policy.value("timing.backoff_max_s"))) |
| 777 | self.admission = admission or Admission() |
| 778 | self.terminator = terminator or Terminator(host=host, serving=serving, windows=lambda: self._windows, |
| 779 | budgets=budgets, **(terminator_options or {})) |
| 780 | if ingest is None: |
| 781 | from . import device_access |
| 782 | ingest = device_access.ingest |
| 783 | self._ingest = ingest |
| 784 | self._network_snapshot = network_snapshot |
| 785 | self._wall = wall |
| 786 | self._monotonic = monotonic |
| 787 | self.stats = TunnelStats() |
| 788 | self.links: dict[str, NodeLink] = {} |
| 789 | self._last_seq = 0 |
| 790 | self._statements = 0 |
| 791 | self._stop = asyncio.Event() |
| 792 | self._watch_task: asyncio.Task | None = None |
| 793 | self._resync_task: asyncio.Task | None = None |
| 794 | self._changed = asyncio.Event() |
| 795 | self._client_ctx: ssl.SSLContext | None = None |
| 796 | self.attention = asyncio.Event() # a node said something that changes what the control plane thinks of this computer: read the enroll again now |
| 797 | self.set_nodes(nodes) |
| 798 | |
| 799 | # -- the nodes |
| 800 | def set_nodes(self, nodes: list[NodeSpec]) -> None: |
| 801 | """The list the control plane gave: a node that is new gets a connection, one that is gone loses it, one whose name, port or pins changed gets a fresh one.""" |
| 802 | wanted = {spec.id: spec for spec in nodes} |
| 803 | for node_id, link in list(self.links.items()): |
| 804 | if node_id not in wanted or link.spec != wanted[node_id]: |
| 805 | del self.links[node_id] |
| 806 | if self._watch_task is not None: # running: end it in the background |
| 807 | asyncio.get_running_loop().create_task(link.stop()) |
| 808 | for node_id, spec in wanted.items(): |
| 809 | if node_id not in self.links: |
| 810 | self.links[node_id] = NodeLink(self, spec) |
| 811 | if self._watch_task is not None: |
| 812 | self.links[node_id].start() |
| 813 | |
| 814 | def other_node_ready(self, link: NodeLink) -> bool: |
| 815 | return any(other is not link and other.state == "ready" for other in self.links.values()) |
| 816 | |
| 817 | # -- what the links ask for (blocking credential work runs in a thread) |
| 818 | def client_context(self) -> ssl.SSLContext: |
| 819 | """The TLS context for every connection to a node, built once (loading the root bundle for each stream would cost more than the stream).""" |
| 820 | if self._client_ctx is None: |
| 821 | self._client_ctx = self._client_context() |
| 822 | return self._client_ctx |
| 823 | |
| 824 | def now(self) -> float: |
| 825 | """The control plane's time (the offset the credentials learned): what valid_until is measured against.""" |
| 826 | return self.credentials.now() |
| 827 | |
| 828 | async def session(self, *, force: bool) -> str: |
| 829 | return await asyncio.to_thread(self.credentials.tunnel_session, force) |
| 830 | |
| 831 | async def proof(self, session: str, nonce: str, node_url: str) -> str: |
| 832 | return await asyncio.to_thread(self.credentials.tunnel_proof, session, nonce, node_url) |
| 833 | |
| 834 | def credential_word(self, error: BaseException) -> str: |
| 835 | """The closed word for a credential failure (never its text).""" |
| 836 | code = getattr(error, "code", "") |
| 837 | if code in ("device_revoked", "license_inactive", "device_identity_required", "device_owned_by_other_license", "invalid_license"): |
| 838 | return {"device_owned_by_other_license": "device_revoked", "invalid_license": "license_inactive"}.get(code, code) |
| 839 | if getattr(error, "retryable", False): |
| 840 | return "control_unreachable" |
| 841 | return "internal_error" if not hasattr(error, "code") else "session_refused" |
| 842 | |
| 843 | # -- links report |
| 844 | def link_ready(self, link: NodeLink) -> None: |
| 845 | self._changed.set() |
| 846 | |
| 847 | def link_ended(self, link: NodeLink) -> None: |
| 848 | self._changed.set() |
| 849 | |
| 850 | def windows_changed(self, windows: Windows) -> None: |
| 851 | self._windows = windows # the windows the newest OK clamped: streams that start from now on use them |
| 852 | |
| 853 | # -- streams |
| 854 | def on_open(self, link: NodeLink, message: dict[str, Any]) -> None: |
| 855 | if link.status == "paused": |
| 856 | return # a node does not open streams for a paused computer; one that does is not served |
| 857 | ticket, why = self.admission.admit(message["peer_ip"], message["cls"]) |
| 858 | if ticket is None: |
| 859 | refused = self.stats.streams_refused |
| 860 | refused[why] = refused.get(why, 0) + 1 |
| 861 | return # no ATTACH: the node closes the client when the attach window ends |
| 862 | self.stats.streams_opened += 1 |
| 863 | link.track(asyncio.create_task(self._stream(link, message, ticket))) |
| 864 | |
| 865 | async def _stream(self, link: NodeLink, message: dict[str, Any], ticket) -> None: |
| 866 | writer: asyncio.StreamWriter | None = None |
| 867 | try: |
| 868 | try: |
| 869 | reader, writer = await link.dial() |
| 870 | writer.write(attach_line(message["stream_id"], message["secret"]).encode("ascii")) |
| 871 | await asyncio.wait_for(writer.drain(), policy.value("protocol.attach_timeout_s")) |
| 872 | first = await asyncio.wait_for(reader.read(1 << 16), self.terminator.handshake_timeout) |
| 873 | except asyncio.CancelledError: |
| 874 | raise |
| 875 | except Exception: # noqa: BLE001 - a stream that cannot be attached is just not served |
| 876 | self.stats.attach_failures += 1 |
| 877 | return |
| 878 | if not first or first.startswith(b"ERR") or first[0] != TLS_HANDSHAKE_BYTE: |
| 879 | self.stats.attach_failures += 1 # the node refused the attach, or what came is not a TLS handshake: no client of ChatGPT starts with anything else |
| 880 | return |
| 881 | await self.terminator.serve(reader, writer, ip=message["peer_ip"], first=first, on_good=lambda ip, link=link: link.send_hint([ip])) |
| 882 | finally: |
| 883 | self.admission.release(ticket) |
| 884 | await close_writer(writer) |
| 885 | |
| 886 | # -- statements |
| 887 | def on_statement(self, token: str) -> None: |
| 888 | if self._statements >= STATEMENTS_IN_FLIGHT: |
| 889 | return # a node that floods statements gets nothing more verified than it can be served |
| 890 | self._statements += 1 |
| 891 | task = asyncio.create_task(self._take_statement(token)) |
| 892 | task.add_done_callback(lambda _: setattr(self, "_statements", self._statements - 1)) |
| 893 | |
| 894 | async def _take_statement(self, token: str) -> None: |
| 895 | """A signed statement the node forwarded: taken only when it verifies against the control plane's keys (a forgery and a statement for another computer do |
| 896 | not), carries a `seq`, and that `seq` is not lower than the highest taken so far (an old `active` replayed cannot lift a paywall). The same `seq` again is |
| 897 | harmless. device_access.ingest then decides whether it is newer than what the computer holds.""" |
| 898 | try: |
| 899 | await asyncio.to_thread(self.credentials.verify_access, token) |
| 900 | except Exception: # noqa: BLE001 - not believed |
| 901 | self.stats.statements_refused += 1 |
| 902 | return |
| 903 | seq = _statement_seq(token) |
| 904 | if seq is None or seq < self._last_seq: |
| 905 | self.stats.statements_refused += 1 |
| 906 | return |
| 907 | self._last_seq = seq |
| 908 | try: |
| 909 | taken = await asyncio.to_thread(self._ingest, token) |
| 910 | except Exception: # noqa: BLE001 |
| 911 | taken = False |
| 912 | if taken: |
| 913 | self.stats.statements_taken += 1 |
| 914 | else: |
| 915 | self.stats.statements_refused += 1 |
| 916 | |
| 917 | # -- running |
| 918 | async def run(self) -> None: |
| 919 | """Keep every node's control connection up until stop().""" |
| 920 | self._watch_task = asyncio.create_task(self._watch()) |
| 921 | for link in self.links.values(): |
| 922 | link.start() |
| 923 | try: |
| 924 | await self._stop.wait() |
| 925 | finally: |
| 926 | background = [task for task in (self._watch_task, self._resync_task) if task is not None] |
| 927 | for task in background: |
| 928 | task.cancel() |
| 929 | await asyncio.gather(*background, *(link.stop() for link in list(self.links.values())), return_exceptions=True) |
| 930 | |
| 931 | def stop(self) -> None: |
| 932 | self._stop.set() |
| 933 | |
| 934 | def kick_all(self, reason: str) -> None: |
| 935 | """Drop every control connection now (wake, network change): each reconnects within a second or so, with no back-off.""" |
| 936 | if reason == "wake": |
| 937 | self.stats.wakes += 1 |
| 938 | else: |
| 939 | self.stats.network_changes += 1 |
| 940 | for link in self.links.values(): |
| 941 | link.kick() |
| 942 | |
| 943 | def _resync_clock(self) -> None: |
| 944 | """Learn the control plane's time again in the background, never ahead of the reconnection: one at a time, each with its own limit.""" |
| 945 | if self._resync_task is not None and not self._resync_task.done(): |
| 946 | return |
| 947 | |
| 948 | async def resync() -> None: |
| 949 | with contextlib.suppress(Exception): |
| 950 | await asyncio.wait_for(asyncio.to_thread(self.credentials.sync_clock, True), RESYNC_SECONDS) |
| 951 | |
| 952 | self._resync_task = asyncio.create_task(resync(), name="e2e-clock-resync") |
| 953 | |
| 954 | async def _watch(self) -> None: |
| 955 | """Notice a wake from sleep, a step of the clock and a change of network, which no timer can: a connection that was open when the lid closed is dead and says |
| 956 | nothing. The wall clock moving more than the monotonic one (macOS and Linux do not count sleep in the monotonic clock) or the loop not running for a long |
| 957 | while (Windows does count it) is a wake; the routes and the interfaces that are up changing is a network change. Either one drops the connections.""" |
| 958 | last_wall, last_mono = self._wall(), self._monotonic() |
| 959 | last_net = await asyncio.to_thread(self._network_snapshot) |
| 960 | last_net_at = last_mono |
| 961 | while True: |
| 962 | await asyncio.sleep(CLOCK_WATCH_SECONDS) |
| 963 | wall, mono = self._wall(), self._monotonic() |
| 964 | jump = (wall - last_wall) - (mono - last_mono) |
| 965 | stalled = mono - last_mono > STALL_WAKE_SECONDS |
| 966 | last_wall, last_mono = wall, mono |
| 967 | if abs(jump) > WALL_JUMP_SECONDS or stalled: |
| 968 | # the connections start over at once (the network may not be usable yet, and a request waiting for it would keep the dead ones up); the difference |
| 969 | # between this computer's clock and the control plane's, learned before, is learned again alongside |
| 970 | self.kick_all("wake") |
| 971 | self._resync_clock() |
| 972 | last_net = await asyncio.to_thread(self._network_snapshot) |
| 973 | last_net_at = mono |
| 974 | continue |
| 975 | if mono - last_net_at >= NETWORK_WATCH_SECONDS: |
| 976 | last_net_at = mono |
| 977 | try: |
| 978 | current = await asyncio.to_thread(self._network_snapshot) |
| 979 | except Exception: # noqa: BLE001 |
| 980 | continue |
| 981 | if current != last_net: |
| 982 | last_net = current |
| 983 | self.kick_all("network") |
| 984 | self._resync_clock() |
| 985 | |
| 986 | # -- what the status file shows |
| 987 | def state(self) -> str: |
| 988 | """connecting (no node yet), ready (every node), degraded (some, not all), clock_skew (no node, and the answer was that the clock is wrong).""" |
| 989 | links = list(self.links.values()) |
| 990 | ready = [link for link in links if link.ready] |
| 991 | if links and len(ready) == len(links): |
| 992 | return "ready" |
| 993 | if ready: |
| 994 | return "degraded" |
| 995 | if any(link.last_error == "clock_skew" for link in links): |
| 996 | return "clock_skew" |
| 997 | return "connecting" |
| 998 | |
| 999 | def snapshot(self) -> dict[str, Any]: |
| 1000 | links = list(self.links.values()) |
| 1001 | return {"state": self.state(), "nodes": [link.snapshot() for link in links], "connected": sum(1 for link in links if link.ready), "configured": len(links), |
| 1002 | "last_error": next((link.last_error for link in links if link.last_error), ""), "stats": self.stats.as_dict(), |
| 1003 | "terminator": self.terminator.stats.as_dict(), "tls13": self.terminator.tls_ready(), |
| 1004 | "windows": {"ping_interval_s": self._windows.ping_interval, "dead_peer_s": self._windows.dead_peer, |
| 1005 | "stream_idle_s": self._windows.stream_idle, "stream_max_age_s": self._windows.stream_max_age}} |