path: root/packages/meshbay-node/src/meshbay_node/transport/quic_client.py

"""
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 os
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 verify_chunk_signature
from meshbay_common.webcrypto import chunk_key_aes as derive_chunk_key, decrypt_chunk_aes as decrypt_chunk
from meshbay_common.protocol import MNP
from meshbay_common.handshake import (
    NONCE_LEN,
    ROLE_CLIENT,
    ROLE_NODE,
    handshake_transcript,
    make_proof,
    quic_binding,
    verify_proof,
)


def _peer_cert_der(proto) -> bytes | None:
    """
    The server certificate as seen by the client — the channel-binding anchor.

    Spike 11.5.6: aioquic 1.3.0 exposes no RFC 5705 exporter, and the peer
    certificate only through a private attribute. Returns None when it is absent;
    callers decide, because absence is not always an error — see below.
    """
    from cryptography.hazmat.primitives import serialization

    tls = getattr(getattr(proto, "_quic", None), "tls", None)
    cert = getattr(tls, "_peer_certificate", None) if tls is not None else None
    if cert is None:
        return None
    return cert.public_bytes(serialization.Encoding.DER)

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 = "",
        peer_cert_der: bytes | None = None,
        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
        # 11.5.6 binding anchor. Travels with the session ticket: on a resumed
        # TLS session the server does not re-send its certificate.
        self._peer_cert_der: bytes | None = peer_cert_der
        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__()

        nonce_c = os.urandom(NONCE_LEN)
        self._proto._send(self._ctrl_stream, {
            "type":     MNP.HANDSHAKE,
            "v":        MNP_VERSION,
            "token":    self._jwt_token,
            "group_id": self._group_id,
            "nonce":    base64.b64encode(nonce_c).decode(),
        })

        reply = await self._proto._recv(self._ctrl_stream)
        if reply.get("type") != MNP.HANDSHAKE_CHALLENGE:
            raise ConnectionError(f"QUIC handshake rejected: {reply}")

        nonce_s = base64.b64decode(reply["nonce"])

        # On a RESUMED TLS session the server does not re-send its certificate, so
        # there is nothing live to bind to. The session ticket is cryptographically
        # derived from the original handshake, so binding to the certificate seen
        # then is sound — but only if we actually saw one. We never fall back to an
        # unbound proof: that would silently drop MitM detection (L4).
        cert_der = _peer_cert_der(self._proto)
        if cert_der is not None:
            self._peer_cert_der = cert_der
        elif getattr(self, "_peer_cert_der", None) is None:
            raise ConnectionError(
                "QUIC peer certificate unavailable and none cached from a prior "
                "session — refusing to handshake without channel binding")
        binding = quic_binding(self._peer_cert_der)

        self._proto._send(self._ctrl_stream, {
            "type":  MNP.HANDSHAKE_RESPONSE,
            "v":     MNP_VERSION,
            "proof": base64.b64encode(make_proof(
                self._gek, ROLE_CLIENT, self._group_id,
                nonce_c, nonce_s, binding)).decode(),
        })

        ack = await self._proto._recv(self._ctrl_stream)
        if ack.get("type") != MNP.HANDSHAKE_ACK:
            raise ConnectionError(f"QUIC handshake rejected: {ack}")

        # Authenticate the node before trusting anything it serves (C3).
        if not verify_proof(
            self._gek, base64.b64decode(ack.get("proof", "")), ROLE_NODE,
            self._group_id, nonce_c, nonce_s, binding,
        ):
            raise ConnectionError("Node failed to prove GEK possession")

        transcript = handshake_transcript(
            ROLE_NODE, self._group_id, nonce_c, nonce_s, binding)
        try:
            Ed25519PublicKey.from_public_bytes(
                base64.b64decode(ack["node_pk"])
            ).verify(base64.b64decode(ack["sig"]), transcript)
        except Exception as exc:
            raise ConnectionError(f"Node signature invalid: {exc}") from exc

        log.debug("QUIC connected to %s:%d", self._host, self._port)

    @property
    def peer_cert_der(self) -> bytes | None:
        """Binding anchor to carry alongside a saved session ticket (11.5.6)."""
        return self._peer_cert_der

    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"])