ALL

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

contracts/fake_node.py

SHA-256: 7c596a2d1f21a210bca7bed36cf16893686df0a55b8d84478077577e42710cc7
371 líneas · 20231 bytes · Ver en GitHub

1"""A Python fake tunnel node for protocol v2 (ALL 2.0): the node side of the contract as a pure state machine.
2
3It is what WP13 (the device's tunnel client) is tested against, and the golden transcripts of contracts/vectors/transcripts.json are its
4output: a generator outside this set feeds it the scripted inputs of each scenario and writes down what it did, and the contract test suite
5checks that it still does exactly that and that the contract checker (e2e_contracts.check_transcript, the same one the Go and TypeScript
6suites run) accepts it. No sockets here: `handle(event)` takes one input event and returns the output events, so a test drives it with a
7fake clock, and `FakeNodeServer` (below) is the thin asyncio wrapper that speaks the same lines over real sockets.
8
9Input events (what happens TO the node) Output events (what the node does)
10 recv {conn, line} a line from a device send {conn, line} a line to a device
11 cp_reply {conn, status, body} the control plane cp {conn, call, body} a call to the control plane
12 client_hello {sni, complete, peer_ip, class} tls_alert {sni} the node answers a client with a TLS alert
13 attach {conn, line} a data connection attach {conn, line, result} ok, or err (with the ERR line in `reply`)
14 tick {} time passes (send CLOSE expired, ...)
15 end {} the scenario is over
16Every event carries `t`, the node's clock in seconds. Randomness (nonce, stream id, attach secret) is injected so a scenario is deterministic.
17"""
18from __future__ import annotations
19
20import asyncio
21import os
22import time
23from collections import deque
24from typing import Any, Awaitable, Callable
25
26import e2e_contracts as c
27
28# The word the node answers a device with, per reason the contract checker found.
29_ERR_FOR = {"malformed_token": "bad_hello", "not_yet_valid": "clock_skew", "stale_proof": "clock_skew", "wrong_audience": "auth", "expired": "auth"}
30
31
32def err_for(code: str, default: str = "auth") -> str:
33 return _ERR_FOR.get(code, default)
34
35
36class FakeNode:
37 def __init__(self, *, node_id: str, node_url: str, jwks: dict[str, Any], skew: int, randoms: dict[str, list[str]]):
38 self.node_id = node_id
39 self.node_url = node_url
40 self.keys = c.parse_jwks(jwks)
41 self.skew = skew
42 # injected randomness first (a scenario is deterministic), fresh random words once it runs out (a live fake node)
43 self.nonces = deque(randoms.get("nonces", []))
44 self.stream_ids = deque(randoms.get("stream_ids", []))
45 self.secrets = deque(randoms.get("secrets", []))
46 self.conns: dict[str, dict[str, Any]] = {}
47 self.seen: set[str] = set()
48 self.streams: dict[str, dict[str, Any]] = {}
49
50 # ------------------------------------------------------------------------------------------------------------ plumbing
51
52 def _conn(self, name: str) -> dict[str, Any]:
53 return self.conns.setdefault(name, {"phase": "new", "claims": None, "session": None, "kind": "hello", "nonce": None, "waiting": False,
54 "host": None, "status": None, "valid_until": 0, "sub": None, "hints": []})
55
56 @staticmethod
57 def _send(t: int, conn: str, line: str) -> dict[str, Any]:
58 return {"t": t, "conn": conn, "ev": "send", "line": line}
59
60 def _err(self, t: int, conn: str, code: str) -> list[dict[str, Any]]:
61 return [self._send(t, conn, c.format_line("ERR", code=code))]
62
63 # ------------------------------------------------------------------------------------------------------------ the entry point
64
65 def handle(self, event: dict[str, Any]) -> list[dict[str, Any]]:
66 t = event["t"]
67 kind = event["ev"]
68 if kind == "recv":
69 return self._recv(t, event["conn"], event["line"])
70 if kind == "cp_reply":
71 return self._cp_reply(t, event["conn"], event["status"], event["body"])
72 if kind == "client_hello":
73 return self._client_hello(t, event)
74 if kind == "attach":
75 return self._attach(t, event)
76 if kind == "tick":
77 return self._tick(t)
78 return []
79
80 # ------------------------------------------------------------------------------------------------------------ device lines
81
82 def _recv(self, t: int, name: str, line: str) -> list[dict[str, Any]]:
83 conn = self._conn(name)
84 try:
85 message = c.parse_line(line, "d2n")
86 except c.ProtocolError as error:
87 return self._err(t, name, error.code)
88 verb = message["verb"]
89 if verb in ("HELLO", "RENEW"):
90 if conn["phase"] != ("new" if verb == "HELLO" else "ready"): # RENEW only on a live connection: not mid-handshake, not after a CLOSE
91 return self._err(t, name, "bad_command")
92 result = c.verify_session(message["session"], now=t, audience=c.pol("identity.aud_tunnel"), keys=self.keys, skew=self.skew)
93 if not result.ok:
94 return self._err(t, name, err_for(result.error or ""))
95 if verb == "RENEW" and result.claims["sub"] != conn["sub"]:
96 return self._err(t, name, "auth")
97 conn.update(session=message["session"], claims=result.claims, kind="hello" if verb == "HELLO" else "renew", waiting=False)
98 if verb == "HELLO":
99 conn["phase"] = "hello"
100 conn["nonce"] = self.nonces.popleft() if self.nonces else c.b64u_encode(os.urandom(c.pol("protocol.nonce_bytes")))
101 return [self._send(t, name, c.format_line("CHALLENGE", nonce=conn["nonce"]))]
102 if verb == "PROOF":
103 if conn["nonce"] is None:
104 return self._err(t, name, "bad_command")
105 result = c.verify_proof(message["proof"], now=t, session_token=conn["session"], session_claims=conn["claims"], nonce=conn["nonce"],
106 htu=self.node_url, skew=self.skew, seen=self.seen)
107 conn["nonce"] = None
108 if not result.ok:
109 return self._err(t, name, err_for(result.error or "", "bad_proof") if result.error != "replayed_jti" else "bad_proof")
110 self.seen.add(result.jti or "")
111 conn["waiting"] = True
112 body = {"v": 1, "device_id": conn["claims"]["sub"], "session_jti": conn["claims"]["jti"], "reason": conn["kind"], "node": self.node_id}
113 return [{"t": t, "conn": name, "ev": "cp", "call": "authorize", "body": body}]
114 if verb == "PING":
115 return [self._send(t, name, c.format_line("PONG"))]
116 if verb == "HINT":
117 conn["hints"] = message["addresses"]
118 return []
119 if verb == "CLOSE":
120 conn["phase"] = "closed"
121 return []
122 return self._err(t, name, "bad_command")
123
124 def _cp_reply(self, t: int, name: str, status: int, body: dict[str, Any]) -> list[dict[str, Any]]:
125 conn = self._conn(name)
126 if not conn["waiting"]:
127 return []
128 conn["waiting"] = False
129 if status != 200:
130 word = body.get("error") if isinstance(body, dict) else None
131 known = {"license_inactive": "license_inactive", "device_revoked": "device_revoked", "paused": "paused"}
132 return self._err(t, name, known.get(word, "control_unavailable" if status >= 500 or status == 429 else "auth"))
133 if c.SCHEMAS.validate("authorize-response.schema.json", body) or body["device_id"] != conn["claims"]["sub"]:
134 return self._err(t, name, "control_unavailable")
135 out: list[dict[str, Any]] = []
136 ok_body = {"v": 2, "node": self.node_id, "host": body["host"], "status": body["status"], "valid_until": body["valid_until"],
137 "server_time": body["server_time"], "policy": body["policy"]}
138 out.append(self._send(t, name, c.format_line("OK", body=ok_body)))
139 conn.update(phase="ready", host=body["host"], status=body["status"], valid_until=body["valid_until"], sub=conn["claims"]["sub"])
140 for other_name, other in self.conns.items():
141 if other_name != name and other["phase"] == "ready" and other["sub"] == conn["sub"]:
142 other["phase"] = "closed"
143 out.append(self._send(t, other_name, c.format_line("CLOSE", reason="replaced")))
144 return out
145
146 # ------------------------------------------------------------------------------------------------------------ clients and streams
147
148 def _client_hello(self, t: int, event: dict[str, Any]) -> list[dict[str, Any]]:
149 sni = event["sni"]
150 if not event["complete"]:
151 return [{"t": t, "ev": "tls_alert", "sni": sni}]
152 for name, conn in self.conns.items():
153 if conn["phase"] == "ready" and conn["host"] == sni and conn["status"] != "paused" and t < conn["valid_until"]:
154 stream_id = self.stream_ids.popleft() if self.stream_ids else os.urandom(c.pol("protocol.stream_id_bytes")).hex()
155 secret = self.secrets.popleft() if self.secrets else c.b64u_encode(os.urandom(c.pol("protocol.attach_secret_bytes")))
156 self.streams[stream_id] = {"secret": secret, "t": t, "consumed": False}
157 line = c.format_line("OPEN", stream_id=stream_id, secret=secret, peer_ip=event["peer_ip"], **{"class": event["class"]})
158 return [self._send(t, name, line)]
159 return [{"t": t, "ev": "tls_alert", "sni": sni}]
160
161 def _attach(self, t: int, event: dict[str, Any]) -> list[dict[str, Any]]:
162 out = {"t": t, "conn": event["conn"], "ev": "attach", "line": event["line"], "result": "err"}
163 try:
164 message = c.parse_line(event["line"], "d2n")
165 except c.ProtocolError:
166 return [out]
167 known = self.streams.get(message["stream_id"])
168 if (known is not None and known["secret"] == message["secret"] and not known["consumed"]
169 and t - known["t"] <= c.pol("protocol.attach_timeout_s")):
170 known["consumed"] = True
171 out["result"] = "ok"
172 return [out]
173
174 def _tick(self, t: int) -> list[dict[str, Any]]:
175 out: list[dict[str, Any]] = []
176 for name, conn in self.conns.items():
177 if conn["phase"] == "ready" and t >= conn["valid_until"]:
178 conn["phase"] = "closed"
179 out.append(self._send(t, name, c.format_line("CLOSE", reason="expired")))
180 return out
181
182
183INPUT_EVENTS = ("recv", "cp_reply", "client_hello", "attach", "tick", "end")
184
185
186def split_transcript(events: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
187 """(inputs, outputs) of a golden transcript: what was done to the node and what it did. An `attach` is an input without its result."""
188 inputs: list[dict[str, Any]] = []
189 outputs: list[dict[str, Any]] = []
190 for event in events:
191 if event["ev"] in INPUT_EVENTS:
192 inputs.append({k: v for k, v in event.items() if k != "result"} if event["ev"] == "attach" else event)
193 if event["ev"] in ("send", "cp", "tls_alert") or event["ev"] == "attach":
194 outputs.append(event)
195 return inputs, outputs
196
197
198def replay(ctx: dict[str, Any], events: list[dict[str, Any]]) -> list[dict[str, Any]]:
199 """Run the fake node over the inputs of a transcript and return the outputs it produced, in order."""
200 inputs, _ = split_transcript(events)
201 node = FakeNode(node_id=ctx["node_id"], node_url=ctx["node_url"], jwks=ctx["jwks"], skew=ctx["skew"], randoms=ctx["randoms"])
202 produced: list[dict[str, Any]] = []
203 for event in inputs:
204 produced.extend(node.handle(event))
205 return produced
206
207
208def rebuild(ctx: dict[str, Any], inputs: list[dict[str, Any]]) -> list[dict[str, Any]]:
209 """The full transcript (inputs and outputs interleaved as they happened) for scripted inputs: what the transcript generator stores."""
210 node = FakeNode(node_id=ctx["node_id"], node_url=ctx["node_url"], jwks=ctx["jwks"], skew=ctx["skew"], randoms=ctx["randoms"])
211 events: list[dict[str, Any]] = []
212 for event in inputs:
213 outs = node.handle(event)
214 if event["ev"] == "attach":
215 events.extend(outs) # the node's own record of the attach carries the result
216 else:
217 events.append(event)
218 events.extend(outs)
219 return events
220
221
222# ---------------------------------------------------------------------------------------------------------------------------------- sockets
223
224ControlPlane = Callable[[dict[str, Any]], "tuple[int, dict[str, Any]] | Awaitable[tuple[int, dict[str, Any]]]"]
225
226
227class OpenedStream:
228 """A client stream the fake node opened for a device: `stream_id` and `secret` are what the device's ATTACH must carry; `attached` resolves
229 to (reader, writer) of the node's end of the data connection once the device attached (bytes written go to the device, bytes read come from it)."""
230
231 def __init__(self, stream_id: str, secret: str, peer_ip: str, cls: str, attached: "asyncio.Future[tuple[asyncio.StreamReader, asyncio.StreamWriter]]"):
232 self.stream_id, self.secret, self.peer_ip, self.cls, self.attached = stream_id, secret, peer_ip, cls, attached
233
234
235class FakeNodeServer:
236 """The fake node over real sockets: a TCP (or TLS, with `ssl_context`) listener that speaks protocol v2 lines, the state machine above behind it.
237
238 `control_plane(body)` answers the node's authorize call (it is given the request body, returns (status, body); sync or async). `clock()` is the
239 node's time in whole seconds. `open_stream(sni, ...)` plays a client that connected to :443 with that SNI: when the node opens a stream for it,
240 the result says so and its `attached` future resolves when the device dials back with the right ATTACH. Nothing here is used by an installed ALL.
241 """
242
243 def __init__(self, *, node_id: str, node_url: str, jwks: dict[str, Any], control_plane: ControlPlane, skew: int | None = None,
244 clock: Callable[[], int] | None = None, tick_interval: float = 1.0):
245 self.node = FakeNode(node_id=node_id, node_url=node_url, jwks=jwks, skew=c.pol("authorization.node_skew_s") if skew is None else skew, randoms={})
246 self.control_plane = control_plane
247 self.clock = clock or (lambda: int(time.time()))
248 self.tick_interval = tick_interval
249 self.writers: dict[str, asyncio.StreamWriter] = {}
250 self.lines: list[tuple[str, str, str]] = [] # (direction, connection, line): what crossed the wire, for the tests
251 self._opened: dict[str, OpenedStream] = {}
252 self._open_writers: set[asyncio.StreamWriter] = set() # every accepted connection, control and data
253 self._count = 0
254 self._server: asyncio.AbstractServer | None = None
255 self._ticker: asyncio.Task[None] | None = None
256 self.port = 0
257
258 async def start(self, host: str = "127.0.0.1", port: int = 0, ssl_context: Any = None) -> int:
259 limit = c.pol("protocol.statement_line_max_bytes") + 16
260 self._server = await asyncio.start_server(self._serve, host, port, ssl=ssl_context, limit=limit)
261 self.port = self._server.sockets[0].getsockname()[1]
262 self._ticker = asyncio.create_task(self._tick_loop())
263 return self.port
264
265 async def stop(self) -> None:
266 if self._ticker:
267 self._ticker.cancel()
268 for writer in list(self._open_writers):
269 writer.close()
270 if self._server:
271 self._server.close()
272 try:
273 await asyncio.wait_for(self._server.wait_closed(), 2)
274 except asyncio.TimeoutError:
275 pass
276
277 async def _tick_loop(self) -> None:
278 while True:
279 await asyncio.sleep(self.tick_interval)
280 await self.tick()
281
282 async def tick(self) -> None:
283 await self._dispatch(self.node.handle({"t": self.clock(), "ev": "tick"}))
284
285 def drop_control(self, name: str) -> None:
286 """Cut one control connection (a node restart, a network drop)."""
287 writer = self.writers.pop(name, None)
288 if writer:
289 writer.close()
290
291 async def open_stream(self, sni: str, *, peer_ip: str = "203.0.113.9", cls: str = "anon", complete: bool = True) -> OpenedStream | None:
292 """A client with this SNI arrives. None when the node would send a TLS alert (unknown name, paused, expired, incomplete ClientHello)."""
293 outputs = self.node.handle({"t": self.clock(), "ev": "client_hello", "sni": sni, "complete": complete, "peer_ip": peer_ip, "class": cls})
294 opened = None
295 for event in outputs:
296 if event["ev"] == "send" and event["line"].startswith("OPEN "):
297 message = c.parse_line(event["line"], "n2d")
298 opened = OpenedStream(message["stream_id"], message["secret"], peer_ip, cls, asyncio.get_running_loop().create_future())
299 self._opened[opened.stream_id] = opened
300 await self._dispatch(outputs)
301 return opened
302
303 async def _dispatch(self, outputs: list[dict[str, Any]]) -> None:
304 pending = list(outputs)
305 while pending:
306 event = pending.pop(0)
307 if event["ev"] == "send":
308 writer = self.writers.get(event["conn"])
309 line = event["line"]
310 self.lines.append(("send", event["conn"], line))
311 if writer is None:
312 continue
313 writer.write(line.encode("ascii") + b"\n")
314 await writer.drain()
315 refused = line.startswith("ERR ") and self.node.conns[event["conn"]]["phase"] == "new"
316 if line.startswith("CLOSE ") or refused:
317 writer.close()
318 self.writers.pop(event["conn"], None)
319 elif event["ev"] == "cp":
320 reply = self.control_plane(event["body"])
321 if asyncio.iscoroutine(reply):
322 reply = await reply
323 status, body = reply
324 pending.extend(self.node.handle({"t": self.clock(), "ev": "cp_reply", "conn": event["conn"], "status": status, "body": body}))
325
326 async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
327 self._count += 1
328 name = f"c{self._count}"
329 self._open_writers.add(writer)
330 try:
331 first = await asyncio.wait_for(self._read(reader), c.pol("protocol.first_line_timeout_s"))
332 if first is None:
333 return
334 if first.startswith("ATTACH "):
335 await self._data(name, first, reader, writer)
336 return
337 self.writers[name] = writer
338 line: str | None = first
339 while line is not None:
340 self.lines.append(("recv", name, line))
341 await self._dispatch(self.node.handle({"t": self.clock(), "ev": "recv", "conn": name, "line": line}))
342 if name not in self.writers:
343 break
344 line = await self._read(reader)
345 except (asyncio.TimeoutError, ConnectionError, asyncio.IncompleteReadError, asyncio.LimitOverrunError):
346 pass
347 finally:
348 self._open_writers.discard(writer)
349 if self.writers.get(name) is writer:
350 self.writers.pop(name, None)
351 writer.close()
352
353 @staticmethod
354 async def _read(reader: asyncio.StreamReader) -> str | None:
355 try:
356 raw = await reader.readuntil(b"\n")
357 except asyncio.IncompleteReadError:
358 return None
359 return raw[:-1].decode("ascii", errors="replace")
360
361 async def _data(self, name: str, line: str, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
362 outputs = self.node.handle({"t": self.clock(), "ev": "attach", "conn": name, "line": line})
363 if outputs[0]["result"] != "ok":
364 writer.write(b"ERR auth\n")
365 await writer.drain()
366 writer.close()
367 return
368 message = c.parse_line(line, "d2n")
369 opened = self._opened.get(message["stream_id"])
370 if opened is not None and not opened.attached.done():
371 opened.attached.set_result((reader, writer))