feat(perception): add bounded authenticated gRPC stream transport
This commit is contained in:
@@ -0,0 +1,463 @@
|
||||
"""Small TLS/socket integration checks, not a model benchmark or load test."""
|
||||
|
||||
# ruff: noqa: E402 -- optional transport extra must be checked before its import.
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import subprocess
|
||||
import threading
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import replace
|
||||
from hashlib import sha256
|
||||
|
||||
import pytest
|
||||
|
||||
grpc = pytest.importorskip("grpc", reason="install the perception-stream extra")
|
||||
|
||||
from k1link.compute.live_perception import LiveIngressEvent
|
||||
from k1link.perception import streaming_wire as wire
|
||||
from k1link.perception.graph_contracts import GraphState
|
||||
from k1link.perception.realtime_contract import StreamStart
|
||||
from k1link.perception.streaming_continuity import ResumeEvidence, StreamSuspended
|
||||
from k1link.perception.streaming_grpc import GrpcStreamClient, GrpcStreamEndpoint
|
||||
from k1link.perception.streaming_lifecycle import StreamingLifecycle
|
||||
from k1link.perception.streaming_network import (
|
||||
MAX_RESULT,
|
||||
LatestReplies,
|
||||
decode_reply,
|
||||
encode_reply,
|
||||
)
|
||||
from k1link.perception.streaming_queue import StreamMailbox
|
||||
|
||||
|
||||
def identity():
|
||||
return StreamStart("run", "source", "worker", "epoch", 1, *(["a" * 64] * 4), "clock", "live")
|
||||
|
||||
|
||||
def event(sequence=1, payload=b"synthetic-raw-points", stamp=1_000_000_000):
|
||||
return LiveIngressEvent(
|
||||
sequence,
|
||||
"capture",
|
||||
1,
|
||||
"lidar",
|
||||
"lidar",
|
||||
sequence,
|
||||
1_799_999_999_123_456_789,
|
||||
stamp,
|
||||
payload,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def tls(tmp_path_factory):
|
||||
root = tmp_path_factory.mktemp("stream-tls")
|
||||
key, cert = root / "key.pem", root / "cert.pem"
|
||||
subprocess.run(
|
||||
[
|
||||
"openssl",
|
||||
"req",
|
||||
"-x509",
|
||||
"-newkey",
|
||||
"rsa:2048",
|
||||
"-nodes",
|
||||
"-days",
|
||||
"1",
|
||||
"-subj",
|
||||
"/CN=localhost",
|
||||
"-keyout",
|
||||
str(key),
|
||||
"-out",
|
||||
str(cert),
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
timeout=10,
|
||||
)
|
||||
return cert.read_bytes(), key.read_bytes()
|
||||
|
||||
|
||||
async def eventually(predicate):
|
||||
async with asyncio.timeout(2):
|
||||
while not predicate():
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def harness(tmp_path, tls, *, consume=None, **options):
|
||||
clock = [1_000_000_000]
|
||||
run = StreamingLifecycle(
|
||||
identity(),
|
||||
tmp_path,
|
||||
StreamMailbox(),
|
||||
threading.Event(),
|
||||
clock_ns=lambda: clock[0],
|
||||
recover_input=True,
|
||||
source_clock_ns=lambda: clock[0],
|
||||
)
|
||||
run.ready()
|
||||
seen = []
|
||||
|
||||
def accept(value):
|
||||
seen.append(value)
|
||||
if consume:
|
||||
consume(value)
|
||||
else:
|
||||
endpoint.publish(
|
||||
run.continuity.epoch, value.ingress_sequence, sha256(value.payload).digest()
|
||||
)
|
||||
|
||||
endpoint = GrpcStreamEndpoint(run, accept, lambda _: None, **options)
|
||||
server, port = await endpoint.serve("localhost:0", certificate=tls[0], private_key=tls[1])
|
||||
clients = []
|
||||
|
||||
def client(access=None, **kwargs):
|
||||
access = access or endpoint.issue(run.continuity.epoch, "capture", 1)
|
||||
result = GrpcStreamClient(f"localhost:{port}", tls[0], access, **kwargs)
|
||||
clients.append(result)
|
||||
return result
|
||||
|
||||
try:
|
||||
yield run, endpoint, client, seen, clock
|
||||
finally:
|
||||
for item in clients:
|
||||
await item.close()
|
||||
await server.stop(0)
|
||||
await eventually(lambda: endpoint.active is None)
|
||||
assert run.close()
|
||||
assert run.mailbox.bytes == 0
|
||||
|
||||
|
||||
def test_delivers_before_eof_both_directions_and_explicit_result_drain(tmp_path, tls):
|
||||
async def check():
|
||||
async with harness(tmp_path, tls) as (run, endpoint, client, seen, _):
|
||||
stream = client()
|
||||
await stream.open()
|
||||
original = event()
|
||||
await stream.send(original)
|
||||
assert await stream.receive() == (1, sha256(original.payload).digest())
|
||||
assert seen == [original]
|
||||
assert not run.mailbox.done
|
||||
await stream.end()
|
||||
# Input EOF alone must not discard delayed graph results.
|
||||
await eventually(lambda: run.mailbox.done)
|
||||
assert endpoint.active is not None
|
||||
endpoint.finish_results(identity())
|
||||
assert await stream.receive() is None
|
||||
await eventually(lambda: endpoint.active is None)
|
||||
assert endpoint.last["ingress"]["terminal"] == "end"
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fault", ["token", "binding", "expired", "duplicate"])
|
||||
def test_untrusted_grant_cannot_reserve_memory_or_stop_owner(tmp_path, tls, fault):
|
||||
async def check():
|
||||
async with harness(tmp_path, tls) as (run, endpoint, client, _, _):
|
||||
access = endpoint.issue(identity(), "capture", 1)
|
||||
if fault == "token":
|
||||
bad = replace(access, token="0" * 64)
|
||||
elif fault == "binding":
|
||||
bad = replace(access, epoch=replace(identity(), calibration_sha256="f" * 64))
|
||||
elif fault == "expired":
|
||||
a, b, _ = endpoint.grant
|
||||
endpoint.grant = (a, b, 0)
|
||||
bad = access
|
||||
else:
|
||||
with pytest.raises(ValueError):
|
||||
endpoint._claim((*access.metadata(), *access.metadata()))
|
||||
bad = replace(access, token="0" * 64)
|
||||
stream = client(bad)
|
||||
with pytest.raises((grpc.aio.AioRpcError, wire.StreamWireError)):
|
||||
await stream.open()
|
||||
assert endpoint.active is None and run.mailbox.bytes == 0
|
||||
assert run.state == GraphState.RUNNING and endpoint.grant is not None
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_disconnect_new_controller_epoch_stale_secret_and_local_lease(tmp_path, tls):
|
||||
async def check():
|
||||
async with harness(tmp_path, tls) as (run, endpoint, client, _, clock):
|
||||
stream = client()
|
||||
original_access = stream.access
|
||||
await stream.open()
|
||||
await stream.send(event())
|
||||
await stream.receive()
|
||||
with pytest.raises(ValueError):
|
||||
endpoint.issue(identity(), "capture", 1)
|
||||
await stream.close()
|
||||
await eventually(lambda: endpoint.active is None)
|
||||
assert run.state == GraphState.RUNNING
|
||||
assert run.continuity.phase == "waiting" and run.lease.start == identity()
|
||||
clock[0] += 100_000_000
|
||||
run.renew(identity()) # Local controller, unrelated to network traffic.
|
||||
epoch = run.begin_input(identity())
|
||||
access = endpoint.issue(epoch, "capture", 1)
|
||||
stale = client(original_access)
|
||||
with pytest.raises((grpc.aio.AioRpcError, wire.StreamWireError)):
|
||||
await stale.open()
|
||||
assert run.continuity.epoch == epoch and run.state == GraphState.RUNNING
|
||||
resumed = client(access)
|
||||
await resumed.open()
|
||||
# Transport does NOT invent a decoded keyframe/causal sensor proof.
|
||||
assert run.continuity.phase == "synchronizing"
|
||||
with pytest.raises(StreamSuspended):
|
||||
endpoint.publish(epoch, 1, b"premature")
|
||||
run.resume_input(epoch, ResumeEvidence(*([clock[0]] * 4), True), lambda: None)
|
||||
await resumed.send(event(2, stamp=clock[0]))
|
||||
assert (await resumed.receive())[0] == 2
|
||||
with pytest.raises(StreamSuspended):
|
||||
endpoint.publish(identity(), 3, b"obsolete")
|
||||
await resumed.end()
|
||||
endpoint.finish_results(epoch)
|
||||
assert await resumed.receive() is None
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_input_idle_deadline_pauses_without_source_eof(tmp_path, tls):
|
||||
async def check():
|
||||
async with harness(tmp_path, tls, idle_timeout=0.15) as (run, endpoint, client, _, _):
|
||||
stream = client()
|
||||
await stream.open()
|
||||
await stream.send(event())
|
||||
await stream.receive()
|
||||
await eventually(lambda: endpoint.active is None)
|
||||
assert run.state == GraphState.RUNNING
|
||||
assert run.continuity.phase == "waiting" and not run.mailbox.done
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_rpc_half_close_without_wire_end_is_not_success(tmp_path, tls):
|
||||
async def check():
|
||||
async with harness(tmp_path, tls) as (run, endpoint, client, _, _):
|
||||
stream = client()
|
||||
await stream.open()
|
||||
await stream.send(event())
|
||||
await stream.receive()
|
||||
await stream.call.done_writing()
|
||||
with pytest.raises(grpc.aio.AioRpcError):
|
||||
await stream.receive()
|
||||
await eventually(lambda: endpoint.active is None)
|
||||
assert run.continuity.phase == "waiting" and not run.mailbox.done
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_authorized_malformed_wire_remains_fatal(tmp_path, tls):
|
||||
async def check():
|
||||
async with harness(tmp_path, tls) as (run, endpoint, client, _, _):
|
||||
stream = client()
|
||||
await stream.open()
|
||||
await stream.send(event())
|
||||
await stream.receive()
|
||||
await stream.call.write(wire.PREFIX.pack(b"BAD!", wire.OPEN, 1) + b"x")
|
||||
with pytest.raises(grpc.aio.AioRpcError):
|
||||
await stream.receive()
|
||||
await eventually(lambda: endpoint.active is None)
|
||||
assert run.state == GraphState.STOPPING
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_oversized_rpc_rejected_before_domain_consumer(tmp_path, tls):
|
||||
async def check():
|
||||
async with harness(tmp_path, tls) as (run, endpoint, client, seen, _):
|
||||
stream = client()
|
||||
await stream.open()
|
||||
with pytest.raises(grpc.aio.AioRpcError):
|
||||
await stream.call.write(b"x" * (64 * 1024 + 1))
|
||||
await stream.receive()
|
||||
assert seen == []
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_latest_reply_bounds_no_hidden_queue():
|
||||
queue = LatestReplies()
|
||||
for i in range(8):
|
||||
assert queue.put(i, b"result")
|
||||
assert len(queue.pending) == 2 and queue.dropped == 6
|
||||
assert queue.take() == (6, b"result")
|
||||
assert queue.take() == (7, b"result")
|
||||
with pytest.raises(ValueError):
|
||||
queue.put(7, b"regressed")
|
||||
queue.finish()
|
||||
assert not queue.put(8, b"late") and queue.drained()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [b"", b"x" * (MAX_RESULT + 1), bytearray(b"x")])
|
||||
def test_result_size_type_rejected_before_queue(payload):
|
||||
with pytest.raises(ValueError):
|
||||
LatestReplies().put(0, payload)
|
||||
|
||||
|
||||
def test_reply_digest_epoch_and_exact_uint64():
|
||||
raw = encode_reply(identity(), (1 << 53) + 1, b"existing-domain-wire")
|
||||
assert decode_reply(identity(), raw) == ((1 << 53) + 1, b"existing-domain-wire")
|
||||
for altered in (raw[:-1] + b"?", raw[:20]):
|
||||
with pytest.raises(ValueError):
|
||||
decode_reply(identity(), altered)
|
||||
with pytest.raises(ValueError):
|
||||
decode_reply(replace(identity(), epoch_id="another"), raw)
|
||||
|
||||
|
||||
def test_tls_name_verification_is_not_disabled(tmp_path, tls):
|
||||
async def check():
|
||||
async with harness(tmp_path, tls) as (run, endpoint, _, seen, _):
|
||||
access = endpoint.issue(identity(), "capture", 1)
|
||||
server, port = await endpoint.serve(
|
||||
"127.0.0.1:0", certificate=tls[0], private_key=tls[1]
|
||||
)
|
||||
stream = GrpcStreamClient(f"127.0.0.1:{port}", tls[0], access)
|
||||
try:
|
||||
# The certificate is for localhost, not this numeric address.
|
||||
with pytest.raises((grpc.aio.AioRpcError, wire.StreamWireError)):
|
||||
await stream.open()
|
||||
assert endpoint.active is None and not seen and run.mailbox.bytes == 0
|
||||
finally:
|
||||
await stream.close()
|
||||
await server.stop(0)
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_slow_result_sink_times_out_without_killing_resident_runtime(tmp_path):
|
||||
class Aborted(Exception):
|
||||
pass
|
||||
|
||||
async def check():
|
||||
run = StreamingLifecycle(
|
||||
identity(),
|
||||
tmp_path,
|
||||
StreamMailbox(),
|
||||
threading.Event(),
|
||||
recover_input=True,
|
||||
source_clock_ns=lambda: 1_000_000_000,
|
||||
)
|
||||
run.ready()
|
||||
endpoint = GrpcStreamEndpoint(
|
||||
run,
|
||||
lambda e: endpoint.publish(identity(), e.ingress_sequence, b"result"),
|
||||
lambda _: None,
|
||||
io_timeout=0.05,
|
||||
)
|
||||
access = endpoint.issue(identity(), "capture", 1)
|
||||
packets = [wire.open_packet(identity(), "capture", 1)]
|
||||
packets += [
|
||||
bytes(piece) for pair in wire.event_packets(identity(), event()) for piece in pair
|
||||
]
|
||||
|
||||
class Context:
|
||||
def auth_context(self):
|
||||
return {"transport_security_type": (b"ssl",)}
|
||||
|
||||
def invocation_metadata(self):
|
||||
return access.metadata()
|
||||
|
||||
async def send_initial_metadata(self, _):
|
||||
pass
|
||||
|
||||
async def read(self):
|
||||
if packets:
|
||||
return packets.pop(0)
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def write(self, _):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def abort(self, *_):
|
||||
raise Aborted()
|
||||
|
||||
try:
|
||||
with pytest.raises(Aborted):
|
||||
async with asyncio.timeout(1):
|
||||
await endpoint.exchange(None, Context())
|
||||
assert run.state == GraphState.RUNNING and run.continuity.phase == "waiting"
|
||||
assert endpoint.active is None and run.mailbox.bytes == 0
|
||||
finally:
|
||||
assert run.close()
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_slow_trusted_callback_quarantines_then_releases_without_replacing_owner(tmp_path, tls):
|
||||
entered, leave = threading.Event(), threading.Event()
|
||||
|
||||
def consume(_):
|
||||
entered.set()
|
||||
leave.wait(3)
|
||||
|
||||
async def check():
|
||||
async with harness(tmp_path, tls, consume=consume) as (run, endpoint, client, _, _):
|
||||
stream = client()
|
||||
try:
|
||||
await stream.open()
|
||||
await stream.send(event())
|
||||
await eventually(entered.is_set)
|
||||
await stream.close()
|
||||
await eventually(lambda: endpoint._quarantine_thread is not None)
|
||||
assert endpoint.active is not None and run.mailbox.bytes > 0
|
||||
with pytest.raises((ValueError, StreamSuspended)):
|
||||
endpoint.issue(identity(), "capture", 1)
|
||||
finally:
|
||||
leave.set()
|
||||
await eventually(lambda: endpoint.active is None)
|
||||
assert run.state == GraphState.RUNNING and run.mailbox.bytes == 0
|
||||
await eventually(lambda: not endpoint._quarantine_thread.is_alive())
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("NDC_STREAM_WORKER_PROBE") != "1", reason="Worker-only bounded probe"
|
||||
)
|
||||
def test_worker_maximum_result_crosses_tls_intact(tmp_path, tls):
|
||||
payload = b"r" * MAX_RESULT
|
||||
|
||||
async def check():
|
||||
async with harness(
|
||||
tmp_path,
|
||||
tls,
|
||||
consume=lambda e: endpoint.publish(identity(), e.ingress_sequence, payload),
|
||||
) as (_, endpoint, client, _, _):
|
||||
stream = client()
|
||||
await stream.open()
|
||||
await stream.send(event())
|
||||
assert await stream.receive() == (1, payload)
|
||||
await stream.end()
|
||||
endpoint.finish_results(identity())
|
||||
assert await stream.receive() is None
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("NDC_STREAM_WORKER_PROBE") != "1", reason="Worker-only bounded probe"
|
||||
)
|
||||
def test_worker_actual_grpc_slow_reader_releases_stream_not_runtime(tmp_path, tls):
|
||||
payload = b"r" * MAX_RESULT
|
||||
|
||||
async def check():
|
||||
async with harness(
|
||||
tmp_path,
|
||||
tls,
|
||||
consume=lambda e: endpoint.publish(identity(), e.ingress_sequence, payload),
|
||||
io_timeout=0.1,
|
||||
) as (run, endpoint, client, _, _):
|
||||
stream = client()
|
||||
await stream.open()
|
||||
# At most 24 tiny inputs; server outputs at most 24 x 1 MiB on Worker.
|
||||
# Never read results: real gRPC/TLS flow control must reach its bound.
|
||||
for sequence in range(1, 25):
|
||||
try:
|
||||
await stream.send(event(sequence))
|
||||
except grpc.aio.AioRpcError:
|
||||
break
|
||||
await asyncio.sleep(0.02)
|
||||
await eventually(lambda: endpoint.active is None)
|
||||
assert run.state == GraphState.RUNNING and run.continuity.phase == "waiting"
|
||||
assert run.mailbox.bytes == 0 and endpoint.last["reply_drops"] > 0
|
||||
|
||||
asyncio.run(check())
|
||||
Reference in New Issue
Block a user