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 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.groupbox import PURPOSE_ACK, PURPOSE_INDEX, unseal
from meshbay_common.protocol import MNP, file_chunk_plaintext
from meshbay_common.handshake import (
    MNP_MIN_SUPPORTED,
    NONCE_LEN,
    ROLE_CLIENT,
    ROLE_NODE,
    check_version,
    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
        # The handshake_ack's sealed payload, once connect() has opened it.
        self._node_config: dict = {}

    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,
            # The oldest node this build can talk to. Declared in the first
            # message so a mismatch is a refusal with a code, not a field that
            # turns up missing three messages later (L2).
            "v_min":    MNP_MIN_SUPPORTED,
            "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}")
        # The node's half of the range, checked before we speak to it further.
        check_version(reply.get("v", ""), reply.get("v_min", ""))

        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

        # Verify, then decrypt — in that order, and it is not incidental. The
        # proof and the signature above are what decide whether this peer is worth
        # trusting at all; opening the payload first would mean acting on data from
        # someone we have not authenticated. Empty on this transport today (D5),
        # but it must still open: a payload that does not is a peer we cannot talk
        # to, not a node with no configuration.
        self._node_config = unseal(
            self._gek, PURPOSE_ACK, MNP.HANDSHAKE_ACK, self._group_id, ack)

        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) -> dict:
        """
        Request the Mesh Group Index.

        Returns `{group_id, version, entries, dirs, roots}` — the sealed payload
        opened, with `group_id` from the envelope that carried it. It used to return
        the bytes of a `GroupIndex.serialize()` envelope for the caller to
        deserialize: the same message type carrying a different encoding on this
        transport alone.

        A payload that does not open raises. It is never an empty index — that is
        indistinguishable from a group with no files, which is why a fallback here
        would be worse than a stop (groupbox.py, §3.4).
        """
        sid = self._new_stream()
        self._proto._send(sid, {"type": MNP.INDEX_SYNC, "v": MNP_VERSION})
        msg = await self._proto._recv(sid)
        if msg.get("type") == "error":
            raise LookupError(msg.get("detail", "index_sync refused"))
        payload = unseal(
            self._gek, PURPOSE_INDEX, MNP.INDEX_SYNC, self._group_id, msg)
        return {"group_id": msg.get("group_id", self._group_id), **payload}

    async def fetch_chunk(self, file_id: str, chunk_index: int) -> bytes:
        """
        Fetch and decrypt one chunk over QUIC.

        The per-chunk Ed25519 signature this used to verify is gone (see
        `meshbay_common.protocol`): the AEAD tag authenticates the ciphertext under a
        GEK-derived key, and the node proved its identity in the handshake — which
        `connect()` verified and pinned — rather than once per megabyte.
        """
        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"))

        # From the id we asked for, not the one the answer claims: a peer that
        # substitutes it would otherwise choose which key we decrypt with.
        return file_chunk_plaintext(
            self._gek, msg, file_hash=bytes.fromhex(file_id))

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