"""Golden-bytes agreement between the Python endpoint and the shared
envelope fixture, plus the wire-level target/namespace binding checks.
The same fixture bytes are asserted in Rust (klbr-runtime/tests/protocol.rs)
and TypeScript (klbr-web/src/lib/protocol-v2.test.ts)."""
from __future__ import annotations
import json
import os
import unittest
from pathlib import Path
from klbr_runtime.wire import (
BindingToken,
ProtocolError,
RELAY_PROTOCOL_VERSION,
VERSION,
Wire,
measured_runtime_identity,
validate_target_stamp,
validate_worker_identity,
)
FIXTURE = Path(__file__).resolve().parents[3] / "tests" / "runtime-fixtures" / "envelopes.json"
def load_fixture() -> list[dict]:
return json.loads(FIXTURE.read_text())
def remote_stamp() -> dict:
return {
"target_id": "linux-build",
"namespace_id": "main",
"namespace_generation": "ns-gen-remote-1",
}
class SharedFixtureAgreement(unittest.TestCase):
def test_fixture_parses_and_has_the_remote_handshake_entries(self):
entries = load_fixture()
self.assertEqual(len(entries), 11, "the shared fixture drives all three consumers")
hello_remote = [e for e in entries
if e["message"]["kind"] == "hello" and e["message"].get("identity")]
execute_remote = [e for e in entries if "target" in e]
self.assertEqual(len(hello_remote), 1)
self.assertEqual(len(execute_remote), 1)
identity = hello_remote[0]["message"]["identity"]
validate_worker_identity(identity)
validate_target_stamp(identity["target"]["stamp"])
validate_target_stamp(execute_remote[0]["target"])
def test_local_golden_hello_stays_valid_without_identity(self):
local_hello = load_fixture()[0]
self.assertEqual(local_hello["message"]["kind"], "hello")
self.assertNotIn("identity", local_hello["message"])
self.assertNotIn("target", local_hello)
class WorkerIdentityChecks(unittest.TestCase):
def base_identity(self) -> dict:
return {
"protocol_version": VERSION,
"relay_protocol_version": RELAY_PROTOCOL_VERSION,
"runtime": {
"platform": "Linux",
"arch": "aarch64",
"interpreter_path": "/usr/bin/python3",
"interpreter_version": "3.11.9",
"capabilities": [
{"capability": "shell", "shell": "/bin/zsh"},
{"capability": "closures", "dill_available": False},
],
},
}
def test_incompatible_worker_version_is_a_clean_refusal(self):
identity = self.base_identity()
identity["protocol_version"] = 2
with self.assertRaises(ProtocolError) as ctx:
validate_worker_identity(identity)
self.assertIn("unsupported worker protocol version 2", str(ctx.exception))
def test_incompatible_relay_version_is_a_clean_refusal(self):
identity = self.base_identity()
identity["relay_protocol_version"] = RELAY_PROTOCOL_VERSION + 1
with self.assertRaises(ProtocolError) as ctx:
validate_worker_identity(identity)
self.assertIn("unsupported relay protocol version", str(ctx.exception))
def test_incomplete_identity_is_refused_field_by_field(self):
identity = self.base_identity()
identity["runtime"]["platform"] = ""
with self.assertRaises(ProtocolError) as ctx:
validate_worker_identity(identity)
self.assertIn("runtime identity platform", str(ctx.exception))
def test_advertised_but_unusable_capability_is_expressible(self):
identity = self.base_identity()
validate_worker_identity(identity)
caps = {c["capability"]: c for c in identity["runtime"]["capabilities"]}
self.assertFalse(caps["closures"]["dill_available"], "advertised but unusable")
broken = self.base_identity()
broken["runtime"]["capabilities"] = [
{"capability": "relay_framing", "version": RELAY_PROTOCOL_VERSION + 3}
]
with self.assertRaises(ProtocolError) as ctx:
validate_worker_identity(broken)
self.assertIn("newer than this endpoint supports", str(ctx.exception))
def test_malformed_capability_entry_is_refused(self):
identity = self.base_identity()
identity["runtime"]["capabilities"] = [{"capability": 7}]
with self.assertRaises(ProtocolError):
validate_worker_identity(identity)
class EnvelopeTargetBindingChecks(unittest.TestCase):
def _stubs(self):
class FakeWriter:
def __init__(self):
self.sent = b""
def write(self, data):
self.sent += data
async def drain(self):
return None
class FakeReader:
def __init__(self, data: bytes):
self.data = data
self.pos = 0
async def readexactly(self, n: int) -> bytes:
chunk = self.data[self.pos:self.pos + n]
self.pos += n
if len(chunk) < n:
raise asyncio.IncompleteReadError(chunk, n)
return chunk
return FakeReader, FakeWriter
def test_local_envelope_claiming_a_stamp_is_forged(self):
import asyncio
import struct as _struct
body = json.dumps({
"version": VERSION,
"session_id": "test-session",
"generation": "test-generation",
"message": {"kind": "shutdown"},
"target": remote_stamp(),
}).encode("utf-8")
FakeReader, _ = self._stubs()
reader = FakeReader(_struct.pack(">I", len(body)) + body)
wire = Wire(reader, None, "test-session", "test-generation") # no target
with self.assertRaises(ProtocolError) as ctx:
asyncio.run(wire.receive())
self.assertIn("must not claim a remote target stamp", str(ctx.exception))
def test_remote_envelope_missing_or_foreign_stamp_is_fatal(self):
import asyncio
import struct as _struct
def frame(packet: dict) -> bytes:
body = json.dumps(packet).encode("utf-8")
return _struct.pack(">I", len(body)) + body
FakeReader, _ = self._stubs()
missing = frame({
"version": VERSION,
"session_id": "test-session",
"generation": "test-generation",
"message": {"kind": "shutdown"},
})
foreign = frame({
"version": VERSION,
"session_id": "test-session",
"generation": "test-generation",
"message": {"kind": "shutdown"},
"target": {**remote_stamp(), "namespace_generation": "ns-gen-other"},
})
matching = frame({
"version": VERSION,
"session_id": "test-session",
"generation": "test-generation",
"message": {"kind": "shutdown"},
"target": remote_stamp(),
})
wire = Wire(FakeReader(missing), None, "test-session", "test-generation",
target=remote_stamp())
with self.assertRaises(ProtocolError) as ctx:
asyncio.run(wire.receive())
self.assertIn("missing its target stamp", str(ctx.exception))
wire = Wire(FakeReader(foreign), None, "test-session", "test-generation",
target=remote_stamp())
with self.assertRaises(ProtocolError) as ctx:
asyncio.run(wire.receive())
self.assertIn("foreign target stamp", str(ctx.exception))
wire = Wire(FakeReader(matching), None, "test-session", "test-generation",
target=remote_stamp())
message = asyncio.run(wire.receive())
self.assertEqual(message["kind"], "shutdown")
def test_remote_wire_stamps_its_own_sends(self):
import asyncio
FakeReader, FakeWriter = self._stubs()
writer = FakeWriter()
wire = Wire(FakeReader(b""), writer, "test-session", "test-generation",
target=remote_stamp())
asyncio.run(wire.send({"kind": "shutdown"}))
packet = json.loads(writer.sent[4:].decode("utf-8"))
self.assertEqual(packet["target"], remote_stamp())
local_writer = FakeWriter()
local = Wire(FakeReader(b""), local_writer, "test-session", "test-generation")
asyncio.run(local.send({"kind": "shutdown"}))
local_packet = json.loads(local_writer.sent[4:].decode("utf-8"))
self.assertNotIn("target", local_packet, "local workers omit the stamp")
if __name__ == "__main__":
unittest.main()
class LaunchBindingVsMeasuredIdentity(unittest.TestCase):
"""The echoed --target binding is NOT measured identity (parent's trap)."""
def test_binding_token_carries_no_identity_fields(self):
token = BindingToken.parse(remote_stamp())
stamp = token.as_stamp()
self.assertEqual(stamp, remote_stamp())
# Structurally: no account/machine/platform/cwd fields exist on the token.
self.assertEqual(sorted(token.__dict__), ["namespace_generation", "namespace_id", "target_id"])
for forbidden in ("account", "machine", "platform", "cwd", "interpreter_path"):
self.assertNotIn(forbidden, stamp)
def test_hello_identity_runtime_is_measured_not_echoed(self):
runtime = measured_runtime_identity()
this_platform = __import__("platform").system()
this_arch = __import__("platform").machine()
self.assertEqual(runtime["platform"], this_platform, "the endpoint measures its own platform")
self.assertEqual(runtime["arch"], this_arch)
import sys
self.assertEqual(runtime["interpreter_path"], sys.executable)
# The binding token is absent from the runtime identity.
self.assertNotIn("namespace_generation", runtime)
def test_parse_rejects_a_malformed_binding(self):
with self.assertRaises(ProtocolError):
BindingToken.parse({**remote_stamp(), "namespace_generation": "bad id"})
with self.assertRaises(ProtocolError):
BindingToken.parse({"target_id": "t", "namespace_id": "n"})