ive harnessed the harness
Something went wrong. Try again.
9.7 kB · 232 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233"""Versioned, bounded JSON frames over a dedicated control socket.
stdout is deliberately NOT a control channel. Any framing or identity error isfatal to the connection; never try to resynchronize a partially read frame."""from __future__ import annotations
import asyncioimport jsonimport reimport structfrom dataclasses import dataclassfrom typing import Any
VERSION = 1# The relay protocol is a DISTINCT version plane from the worker protocol:# a remote launcher multiplexes framed control over a different transport# and the two must never be conflated.RELAY_PROTOCOL_VERSION = 1MAX_FRAME = 1_048_576MAX_CODE = 262_144_ID = re.compile(r"[A-Za-z0-9_.-]{1,128}\Z")
class ProtocolError(Exception): pass
def identifier(value: Any) -> str: if not isinstance(value, str) or not _ID.fullmatch(value): raise ProtocolError("invalid protocol identifier") return value
def fields(value: Any, required: set[str], optional: set[str] = frozenset()) -> dict: if not isinstance(value, dict) or not required <= value.keys() or value.keys() - required - optional: raise ProtocolError(f"expected fields {sorted(required)}, optional {sorted(optional)}") return value
def text(value: Any, maximum: int = MAX_CODE) -> str: if not isinstance(value, str) or len(value.encode("utf-8")) > maximum: raise ProtocolError("expected bounded UTF-8 text") return value
def _unique(pairs: list[tuple[str, Any]]) -> dict: result = {} for key, value in pairs: if key in result: raise ProtocolError(f"duplicate JSON key: {key}") result[key] = value return result
def _nonfinite(value: str) -> Any: raise ProtocolError(f"non-finite JSON number: {value}")
def encode(value: dict) -> bytes: body = json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode("utf-8") if not 0 < len(body) <= MAX_FRAME: raise ProtocolError("frame exceeds byte limit") return struct.pack(">I", len(body)) + body
async def read_frame(reader: asyncio.StreamReader) -> dict: size = struct.unpack(">I", await reader.readexactly(4))[0] if not 0 < size <= MAX_FRAME: raise ProtocolError("invalid frame length") raw = await reader.readexactly(size) try: return json.loads(raw.decode("utf-8"), object_pairs_hook=_unique, parse_constant=_nonfinite) except (ValueError, UnicodeError, RecursionError) as error: raise ProtocolError("invalid JSON frame") from error
@dataclass(frozen=True)class BindingToken: """The host-issued launch binding a worker echoes in Hello.
This is a binding token, NOT measured identity: it proves this process is the one the host launched (the same evidence class as --token), and nothing about the target itself. It deliberately carries no account, machine, platform or cwd fields, so it cannot be used where a measured identity is required. Verified identity comes from an actual probe on the target; if the two disagree, that is a refusal, not a merge. """
target_id: str namespace_id: str namespace_generation: str
@classmethod def parse(cls, value: Any) -> "BindingToken": stamp = validate_target_stamp(value) return cls(**stamp)
def as_stamp(self) -> dict: return { "target_id": self.target_id, "namespace_id": self.namespace_id, "namespace_generation": self.namespace_generation, }
def measured_runtime_identity() -> dict: """Measure THIS process's runtime identity on the machine it runs on.
This is the endpoint's own probe: platform, arch, interpreter path and version. It is the only allowed source for the Hello runtime identity; the echoed launch binding is never a substitute for it. """ import platform as _platform import sys as _sys
return { "platform": _platform.system() or "unknown", "arch": _platform.machine() or "unknown", "interpreter_path": _sys.executable, "interpreter_version": f"{_sys.version_info.major}.{_sys.version_info.minor}.{_sys.version_info.micro}", "capabilities": [], }
def validate_target_stamp(stamp: Any) -> dict: """Validate and normalize a target/namespace stamp on the wire.""" fields(stamp, {"target_id", "namespace_id", "namespace_generation"}) return { "target_id": identifier(stamp["target_id"]), "namespace_id": identifier(stamp["namespace_id"]), "namespace_generation": identifier(stamp["namespace_generation"]), }
def validate_worker_identity(identity: Any) -> dict: """Validate the runtime/target identity a worker reports in its Hello.
An incompatible or incomplete identity is a clean typed refusal, never a crash and never silently mapped to something acceptable. """ fields(identity, {"protocol_version", "relay_protocol_version", "runtime"}, {"target"}) if type(identity["protocol_version"]) is not int: raise ProtocolError("identity protocol_version must be an integer") if type(identity["relay_protocol_version"]) is not int: raise ProtocolError("identity relay_protocol_version must be an integer") runtime = fields(identity["runtime"], {"platform", "arch", "interpreter_path", "interpreter_version"}, {"environment_fingerprint", "capabilities"}) for field in ("platform", "arch", "interpreter_path", "interpreter_version"): value = runtime[field] if not isinstance(value, str) or not value.strip(): raise ProtocolError(f"runtime identity {field} must be non-empty") caps = runtime.get("capabilities", []) if not isinstance(caps, list): raise ProtocolError("capabilities must be a list") for cap in caps: if not isinstance(cap, dict) or not isinstance(cap.get("capability"), str): raise ProtocolError("malformed capability entry") if cap["capability"] == "closures": if type(cap.get("dill_available")) is not bool: raise ProtocolError("closures capability requires a boolean dill_available") elif cap["capability"] == "skill_manifest": revision = cap.get("revision") if not isinstance(revision, str) or not revision.strip(): raise ProtocolError("skill_manifest capability requires a revision") elif cap["capability"] == "relay_framing": version = cap.get("version") if type(version) is not int: raise ProtocolError("relay_framing capability requires an integer version") if version > RELAY_PROTOCOL_VERSION: raise ProtocolError("relay framing version is newer than this endpoint supports") if identity["protocol_version"] != VERSION: raise ProtocolError( f"unsupported worker protocol version {identity['protocol_version']}; " f"this endpoint speaks version {VERSION}" ) if identity["relay_protocol_version"] != RELAY_PROTOCOL_VERSION: raise ProtocolError( f"unsupported relay protocol version {identity['relay_protocol_version']}; " f"this endpoint speaks relay version {RELAY_PROTOCOL_VERSION}" ) return identity
class Wire: def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, session_id: str, generation: str, target: dict | None = None): self.reader, self.writer = reader, writer self.session_id = identifier(session_id) self.generation = identifier(generation) self.target = validate_target_stamp(target) if target is not None else None self._write_lock = asyncio.Lock()
async def send(self, message: dict) -> None: packet = {"version": VERSION, "session_id": self.session_id, "generation": self.generation, "message": message} if self.target is not None: packet["target"] = dict(self.target) packet = encode(packet) async with self._write_lock: self.writer.write(packet) await self.writer.drain()
async def receive(self) -> dict: packet = fields(await read_frame(self.reader), {"version", "session_id", "generation", "message"}, {"target"}) if type(packet["version"]) is not int or packet["version"] != VERSION: raise ProtocolError("unsupported protocol version") if packet["session_id"] != self.session_id or packet["generation"] != self.generation: raise ProtocolError("stale or foreign worker identity") # Target/namespace mismatch joins session/generation mismatch in the # fatal class: a local frame claiming a stamp is a forged identity, # a remote frame missing its stamp cannot be routed. if self.target is None: if "target" in packet: raise ProtocolError("local worker envelope must not claim a remote target stamp") else: if "target" not in packet: raise ProtocolError("remote worker envelope is missing its target stamp") actual = validate_target_stamp(packet["target"]) if actual != self.target: raise ProtocolError( f"foreign target stamp: expected {self.target}, got {actual}" ) message = packet["message"] if not isinstance(message, dict) or not isinstance(message.get("kind"), str): raise ProtocolError("missing message kind") return message