diff options
Diffstat (limited to 'packages/meshbay-node/src/meshbay_node/transport/quic_client.py')
| -rw-r--r-- | packages/meshbay-node/src/meshbay_node/transport/quic_client.py | 197 |
1 files changed, 197 insertions, 0 deletions
diff --git a/packages/meshbay-node/src/meshbay_node/transport/quic_client.py b/packages/meshbay-node/src/meshbay_node/transport/quic_client.py new file mode 100644 index 0000000..6041fef --- /dev/null +++ b/packages/meshbay-node/src/meshbay_node/transport/quic_client.py @@ -0,0 +1,197 @@ +""" +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, + ): + self._host = host + self._port = port + self._jwt_token = jwt_token + self._gek = gek + self._pk_node = Ed25519PublicKey.from_public_bytes( + base64.b64decode(pk_node_b64)) + self._proto: _MNPClientProtocol | None = None + self._cm = None + self._ctrl_stream = 0 # stream 0 = control/handshake + + async def __aenter__(self): + await self.connect() + return self + + async def __aexit__(self, *_): + await self.close() + + 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 + ) + self._cm = connect( + self._host, self._port, + configuration=config, + create_protocol=_MNPClientProtocol, + ) + self._proto = await self._cm.__aenter__() + + # MNP handshake on stream 0 + self._proto._send(self._ctrl_stream, { + "type": MNP.HANDSHAKE, + "v": MNP_VERSION, + "token": self._jwt_token, + }) + 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 |