""" MeshBay Node — QUIC chunk client (MNP v2). Drop-in replacement for ChunkClient with QUIC/UDP transport. Same application-layer protocol: MNP handshake → stream requests. Node identity verified via Ed25519 PK from hub (not TLS cert chain). TLS cert is self-signed; we use CERT_NONE equivalent in QUIC config. """ import asyncio import base64 import logging import struct from pathlib import Path import blake3 import jwt import msgpack from aioquic.asyncio import connect, QuicConnectionProtocol from aioquic.quic.configuration import QuicConfiguration from aioquic.quic.events import QuicEvent, StreamDataReceived from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey from meshbay_common import MNP_VERSION from meshbay_common.crypto import ( chunk_key as derive_chunk_key, decrypt_chunk, verify_chunk_signature, ) from meshbay_common.protocol import MNP log = logging.getLogger(__name__) ALPN = ["meshbay-mnp"] MAX_MSG = 64 * 1024 * 1024 # ── Client-side protocol ────────────────────────────────────────────────────── class _MNPClientProtocol(QuicConnectionProtocol): """ Client-side QUIC protocol. Uses per-stream asyncio.Queue for received messages — no race conditions. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._bufs: dict[int, bytearray] = {} self._queues: dict[int, asyncio.Queue] = {} def _queue_for(self, stream_id: int) -> asyncio.Queue: if stream_id not in self._queues: self._queues[stream_id] = asyncio.Queue() self._bufs[stream_id] = bytearray() return self._queues[stream_id] def quic_event_received(self, event: QuicEvent) -> None: if isinstance(event, StreamDataReceived): sid = event.stream_id q = self._queue_for(sid) self._bufs[sid].extend(event.data) buf = self._bufs[sid] while len(buf) >= 4: length = struct.unpack(">I", buf[:4])[0] if length > MAX_MSG: raise ValueError(f"Message too large: {length}") if len(buf) < 4 + length: break msg_bytes = bytes(buf[4:4 + length]) del buf[:4 + length] q.put_nowait(msgpack.unpackb(msg_bytes, raw=False)) async def _recv(self, stream_id: int, timeout: float = 10.0) -> dict: q = self._queue_for(stream_id) return await asyncio.wait_for(q.get(), timeout=timeout) def _send(self, stream_id: int, obj: dict) -> None: data = msgpack.packb(obj, use_bin_type=True) payload = struct.pack(">I", len(data)) + data self._quic.send_stream_data(stream_id, payload) self.transmit() # ── QuicChunkClient ─────────────────────────────────────────────────────────── class QuicChunkClient: """ QUIC-based MNP client. Drop-in replacement for ChunkClient. Usage: async with QuicChunkClient(host, port, token, gek, pk_node_b64) as client: data = await client.fetch_chunk(file_id, chunk_index=0) """ def __init__( self, host: str, port: int, jwt_token: str, gek: bytes, pk_node_b64: str, local_port: int = 0, # 0 = OS picks; set for hole punching (Port-Restricted) group_id: str = "", session_ticket: object | None = None, ): self._host = host self._port = port self._jwt_token = jwt_token self._gek = gek self._local_port = local_port self._group_id = group_id self._pk_node = Ed25519PublicKey.from_public_bytes( base64.b64decode(pk_node_b64)) self._proto: _MNPClientProtocol | None = None self._cm = None self._ctrl_stream = 0 self._session_ticket = session_ticket async def __aenter__(self): await self.connect() return self async def __aexit__(self, *_): await self.close() @property def session_ticket(self) -> object | None: return self._session_ticket def _save_ticket(self, ticket: object) -> None: self._session_ticket = ticket async def connect(self) -> None: import ssl config = QuicConfiguration( is_client=True, alpn_protocols=ALPN, verify_mode=ssl.CERT_NONE, # identity verified via Ed25519 at MNP layer ) if self._session_ticket: config.session_ticket = self._session_ticket self._cm = connect( self._host, self._port, configuration=config, create_protocol=_MNPClientProtocol, local_port=self._local_port, session_ticket_handler=self._save_ticket, ) self._proto = await self._cm.__aenter__() handshake_msg = { "type": MNP.HANDSHAKE, "v": MNP_VERSION, "token": self._jwt_token, } if self._group_id: handshake_msg["group_id"] = self._group_id self._proto._send(self._ctrl_stream, handshake_msg) ack = await self._proto._recv(self._ctrl_stream) if ack.get("type") != MNP.HANDSHAKE_ACK: raise ConnectionError(f"QUIC handshake rejected: {ack}") log.debug("QUIC connected to %s:%d", self._host, self._port) async def close(self) -> None: if self._cm: await self._cm.__aexit__(None, None, None) self._cm = None self._proto = None def _new_stream(self) -> int: """Open a new bidirectional QUIC stream for a request.""" stream_id = self._proto._quic.get_next_available_stream_id(is_unidirectional=False) return stream_id async def fetch_index(self) -> bytes: """Request the Mesh Group Index.""" sid = self._new_stream() self._proto._send(sid, {"type": MNP.INDEX_SYNC, "v": MNP_VERSION}) msg = await self._proto._recv(sid) return base64.b64decode(msg["index_b64"]) async def fetch_chunk(self, file_id: str, chunk_index: int) -> bytes: """Fetch, verify, and decrypt one chunk over QUIC.""" sid = self._new_stream() self._proto._send(sid, { "type": MNP.FILE_REQUEST, "v": MNP_VERSION, "file_id": file_id, "chunk_index": chunk_index, }) msg = await self._proto._recv(sid, timeout=30.0) if msg.get("type") == "error": raise LookupError(msg.get("detail", "Unknown error")) ct = base64.b64decode(msg["ct_b64"]) nonce = base64.b64decode(msg["nonce_b64"]) ct_hash = base64.b64decode(msg["ct_hash_b64"]) pt_hash = base64.b64decode(msg["pt_hash_b64"]) sig = base64.b64decode(msg["sig_b64"]) file_hash = base64.b64decode(msg["file_hash_b64"]) ci = msg["chunk_index"] verify_chunk_signature(self._pk_node, ci, nonce, ct_hash, sig) if blake3.blake3(ct).digest() != ct_hash: raise ValueError("Ciphertext hash mismatch") ckey = derive_chunk_key(self._gek, file_hash, ci) plaintext = decrypt_chunk(ckey, nonce, ct) if blake3.blake3(plaintext).digest() != pt_hash: raise ValueError("Plaintext hash mismatch after decryption") return plaintext async def fetch_stream_segment( self, file_id: str, segment_index: int, segment_duration: int = 4, ) -> bytes: """Fetch one HLS segment (MPEG-TS bytes) over QUIC.""" sid = self._new_stream() self._proto._send(sid, { "type": MNP.STREAM_SEGMENT, "v": MNP_VERSION, "file_id": file_id, "segment_index": segment_index, "segment_duration": segment_duration, }) msg = await self._proto._recv(sid, timeout=30.0) if msg.get("type") == "error": raise LookupError(msg.get("detail", "Unknown error")) return base64.b64decode(msg["data_b64"])