"""Versioned, bounded JSON frames over a dedicated control socket.
stdout is deliberately NOT a control channel. Any framing or identity error is
fatal to the connection; never try to resynchronize a partially read frame.
"""
from __future__ import annotations
import asyncio
import json
import re
import struct
from dataclasses import dataclass
from 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 = 1
MAX_FRAME = 1_048_576
MAX_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