"""TLS gRPC candidate carrying the existing raw-first binary wire unchanged. One stream per resident profile. The trusted local controller issues one-shot, short-lived capabilities AFTER choosing the input epoch. Network code cannot activate models, renew leases, reset temporal stores or authorize a reconnect. No insecure listener/channel, route upload, background source queue or retry. The gRPC method uses identity byte serialization (not generated protobufs): each request is <=64 KiB of the existing MCI2 byte stream, each response one bounded MCR3 envelope. Source/result semantics remain in the existing adapters. """ from __future__ import annotations import asyncio import hmac import secrets import socket import threading import time from collections.abc import Callable, Iterable from contextlib import suppress from dataclasses import dataclass, field from hashlib import sha256 from typing import Any import grpc # type: ignore[import-untyped] # Upstream wheel does not ship type stubs. from k1link.compute.live_perception import LiveIngressEvent from . import streaming_wire as wire from .realtime_contract import StreamStart from .streaming_ingress import StreamingIngress from .streaming_lifecycle import StreamingLifecycle from .streaming_network import ( BRIDGE_RESERVATION, INPUT_CHUNK, MAX_REPLY, LatestReplies, decode_reply, encode_reply, ) SERVICE = "missioncore.perception.v1.BinaryStream" METHOD = f"/{SERVICE}/Exchange" OPTIONS = ( ("grpc.enable_retries", 0), ("grpc.http2.bdp_probe", 0), # A fixed 1 MiB HTTP/2 receive window, not the default 64 KiB which # serializes large observations into many WAN/SSH round trips. Keep BDP # auto-growth disabled; native buffering remains part of the RSS gate. ("grpc.http2.lookahead_bytes", wire.MAX_FRAGMENT), ("grpc.http2.write_buffer_size", INPUT_CHUNK), ("grpc.max_metadata_size", 4096), ("grpc.max_concurrent_streams", 2), ) @dataclass(frozen=True) class StreamAccess: epoch: StreamStart session_id: str session_generation: int token: str = field(repr=False) def metadata(self) -> tuple[tuple[str, str], ...]: return ( ("authorization", "Bearer " + self.token), ("x-stream-binding", wire.binding(self.epoch)), ) class GrpcStreamEndpoint: def __init__( self, runtime: StreamingLifecycle, consume: Callable[[LiveIngressEvent], None], notice: Callable[[dict[str, Any]], None], *, idle_timeout: float = 2.0, io_timeout: float = 0.25, ) -> None: if runtime.continuity is None: raise ValueError("external stream requires a resumable resident runtime") if not 0 < idle_timeout <= 30 or not 0 < io_timeout <= 2: raise ValueError("invalid bounded network timeouts") self.runtime, self.consume, self.notice = runtime, consume, notice self.idle_timeout, self.io_timeout = idle_timeout, io_timeout self.lock = threading.Lock() self.grant: tuple[StreamAccess, bytes, float] | None = None self.active: tuple[StreamStart, LatestReplies] | None = None self.last: dict[str, Any] | None = None self._quarantine_thread: threading.Thread | None = None def issue(self, epoch: StreamStart, session_id: str, session_generation: int) -> StreamAccess: """Local control plane only. Delivery of this secret is outside data-plane.""" wire.identifier(session_id) wire.decimal(session_generation) if session_generation < 1: raise ValueError("invalid acquisition generation") self.runtime.check_input(epoch, synchronizing=True) with self.lock: if self.active is not None: raise wire.StreamWireError("previous network stream has not drained") access = StreamAccess(epoch, session_id, session_generation, secrets.token_hex(32)) # Do not retain the plaintext secret on the endpoint. self.grant = ( StreamAccess(epoch, session_id, session_generation, ""), sha256(access.token.encode()).digest(), time.monotonic() + 30, ) return access def _claim(self, metadata: Iterable[tuple[str, str | bytes]]) -> StreamAccess: values: dict[str, str | bytes] = {} for key, value in metadata: if key in ("authorization", "x-stream-binding"): if key in values: raise wire.StreamWireError("stream authorization rejected") values[key] = value token = values.get("authorization") if not isinstance(token, str) or len(token) != 71 or not token.startswith("Bearer "): raise wire.StreamWireError("stream authorization rejected") with self.lock: grant = self.grant if ( grant is None or self.active is not None or time.monotonic() >= grant[2] or values.get("x-stream-binding") != wire.binding(grant[0].epoch) or not hmac.compare_digest(sha256(token[7:].encode()).digest(), grant[1]) ): raise wire.StreamWireError("stream authorization rejected") self.runtime.check_input(grant[0].epoch, synchronizing=True) self.grant = None self.active = (grant[0].epoch, LatestReplies()) return grant[0] def publish(self, epoch: StreamStart, sequence: int, payload: bytes) -> bool: self.runtime.check_input(epoch) with self.lock: if self.active is None or self.active[0] != epoch: return False return self.active[1].put(sequence, payload) def finish_results(self, epoch: StreamStart) -> None: """Graph controller signals output EOF AFTER all admitted work is accounted.""" with self.lock: if self.active is not None and self.active[0] == epoch: self.active[1].finish() async def serve( self, address: str, *, certificate: bytes, private_key: bytes, control_handlers: tuple[Any, ...] = (), ) -> Any: if not certificate or not private_key: raise ValueError("TLS identity is required") server = grpc.aio.server( options=OPTIONS + ( ("grpc.max_receive_message_length", INPUT_CHUNK), ("grpc.max_send_message_length", MAX_REPLY), ), # aio reserves its next accept before awaiting a request. Leave # one bounded accept slot beside Exchange + a completing Poll; # otherwise the next Poll can inherit a stale "limit exceeded". # Data ownership still admits exactly one Exchange, never 2 GPUs. maximum_concurrent_rpcs=3 if control_handlers else 2, compression=grpc.Compression.NoCompression, ) server.add_generic_rpc_handlers( ( grpc.method_handlers_generic_handler( SERVICE, {"Exchange": grpc.stream_stream_rpc_method_handler(self.exchange)} ), *control_handlers, ) ) port = server.add_secure_port( address, grpc.ssl_server_credentials(((private_key, certificate),)) ) if port == 0: raise RuntimeError("TLS stream listener failed to bind") await server.start() return server, port async def exchange(self, unused_iterator: Any, context: Any) -> None: if b"ssl" not in context.auth_context().get("transport_security_type", ()): await context.abort(grpc.StatusCode.UNAUTHENTICATED, "TLS required") try: access = self._claim(context.invocation_metadata()) except (ValueError, RuntimeError): await context.abort(grpc.StatusCode.UNAUTHENTICATED, "stream authorization rejected") return assert self.active is not None outbox = self.active[1] left = right = receiver = reservation = None tasks: list[asyncio.Task[None]] = [] try: reservation = self.runtime.mailbox.reserve_ingress(BRIDGE_RESERVATION) left, right = socket.socketpair() for sock in (left, right): sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 32 * 1024) sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 32 * 1024) sock.setblocking(False) receiver = StreamingIngress( right, self.runtime, access.session_id, access.session_generation, self.consume, self.notice, input_epoch=access.epoch, fragment_timeout=self.io_timeout, idle_timeout=self.idle_timeout, ) receiver.start() await context.send_initial_metadata((("x-stream-binding", wire.binding(access.epoch)),)) async def read_input() -> None: while True: async with asyncio.timeout(self.idle_timeout): raw = await context.read() if raw is grpc.aio.EOF: # Explicit wire End can already have closed the peer. with suppress(OSError): left.shutdown(socket.SHUT_WR) return if not isinstance(raw, bytes) or not 0 < len(raw) <= INPUT_CHUNK: raise wire.StreamWireError("invalid network chunk") async with asyncio.timeout(self.io_timeout): await asyncio.get_running_loop().sock_sendall(left, raw) async def write_output() -> None: while True: if not receiver.thread.is_alive() and receiver.terminal != "end": raise ConnectionError("input interrupted or rejected") self.runtime.check_input(access.epoch, synchronizing=True) item = outbox.take() if item is not None: self.runtime.check_input(access.epoch) raw = encode_reply(access.epoch, *item) async with asyncio.timeout(self.io_timeout): await context.write(raw) del raw item = None elif receiver.terminal == "end" and outbox.drained(): return else: await asyncio.sleep(0.01) tasks = [asyncio.create_task(read_input()), asyncio.create_task(write_output())] await asyncio.gather(*tasks) except (TimeoutError, OSError, ValueError, RuntimeError): try: # Even error-status delivery can block behind an unread result. # It must never delay closing ingress/releasing local buffers. async with asyncio.timeout(self.io_timeout): await context.abort( grpc.StatusCode.UNAVAILABLE, "stream interrupted; fresh controller grant required", ) except TimeoutError: context.set_code(grpc.StatusCode.UNAVAILABLE) except Exception as exc: # A blocked peer can cancel the HTTP/2 call while the server is # sending its final status. Cleanup still owns the real outcome. if not isinstance(exc, grpc.RpcError): raise finally: for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) if left is not None: left.close() if receiver is not None and receiver.thread.ident is not None: deadline = time.monotonic() + 1 while receiver.thread.is_alive() and time.monotonic() < deadline: await asyncio.sleep(0.01) elif right is not None: right.close() outbox.close() drained = receiver is None or not receiver.thread.is_alive() if drained: self._complete(receiver, reservation, outbox) else: self._quarantine_thread = threading.Thread( target=self._finish_quarantine, args=(receiver, reservation, outbox), daemon=True, name="perception-network-quarantine", ) self._quarantine_thread.start() # A hung trusted consumer is quarantined: no replacement grant or # premature reservation release while it still owns raw bytes. def _finish_quarantine( self, receiver: StreamingIngress, reservation: Any, outbox: LatestReplies ) -> None: receiver.thread.join() # Bounded to this one already-quarantined stream. self._complete(receiver, reservation, outbox) def _complete( self, receiver: StreamingIngress | None, reservation: Any, outbox: LatestReplies ) -> None: if reservation is not None: reservation.release() with self.lock: self.last = { "ingress": receiver.snapshot() if receiver else None, "reply_drops": outbox.dropped, "network_threads_drained": True, "transport": "grpc-tls-binary/v1", "actuation_allowed": False, } self.active = None class GrpcStreamClient: """Explicit async source/result API. Exactly one writer and one reader.""" def __init__( self, target: str, roots: bytes, access: StreamAccess, *, timeout: float = 0.25 ) -> None: if not roots or not 0 < timeout <= 2: raise ValueError("trusted TLS roots and bounded timeout required") self.access, self.timeout = access, timeout self.channel = grpc.aio.secure_channel( target, grpc.ssl_channel_credentials(root_certificates=roots), options=OPTIONS + ( ("grpc.max_send_message_length", INPUT_CHUNK), ("grpc.max_receive_message_length", MAX_REPLY), ), compression=grpc.Compression.NoCompression, ) self.call = self.channel.stream_stream(METHOD)( metadata=access.metadata(), wait_for_ready=False ) self.last_reply = -1 self._sending = self._reading = False async def open(self) -> None: try: async with asyncio.timeout(2): metadata = await self.call.initial_metadata() if dict(metadata).get("x-stream-binding") != wire.binding(self.access.epoch): raise wire.StreamWireError("server did not admit current input") await self._send( ( wire.open_packet( self.access.epoch, self.access.session_id, self.access.session_generation ), ) ) except BaseException: self.call.cancel() raise async def _send(self, pieces: Iterable[bytes | memoryview]) -> None: if self._sending: self.call.cancel() raise wire.StreamWireError("one source writer per stream") self._sending = True try: async with asyncio.timeout(self.timeout): for piece in pieces: for offset in range(0, len(piece), INPUT_CHUNK): await self.call.write(bytes(piece[offset : offset + INPUT_CHUNK])) except BaseException: self.call.cancel() raise finally: self._sending = False async def send(self, event: LiveIngressEvent) -> None: if (event.session_id, event.session_generation) != ( self.access.session_id, self.access.session_generation, ): self.call.cancel() raise wire.StreamWireError("source acquisition binding changed") await self._send( piece for pair in wire.event_packets(self.access.epoch, event) for piece in pair ) async def end(self) -> None: await self._send((wire.terminal_packet(self.access.epoch),)) await self.call.done_writing() async def receive(self, *, timeout: float = 2) -> tuple[int, bytes] | None: if not 0 < timeout <= 30: raise ValueError("invalid bounded result deadline") if self._reading: self.call.cancel() raise wire.StreamWireError("one result reader per stream") self._reading = True try: async with asyncio.timeout(timeout): raw = await self.call.read() if raw is grpc.aio.EOF: return None sequence, payload = decode_reply(self.access.epoch, raw) if sequence <= self.last_reply: raise wire.StreamWireError("network result sequence regressed") self.last_reply = sequence return sequence, payload except BaseException: self.call.cancel() raise finally: self._reading = False async def close(self) -> None: self.call.cancel() await self.channel.close()