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

"""
MeshBay Node — QUIC chunk server (MNP v2).

Replaces the TCP+TLS ChunkServer with QUIC transport.
Advantages over TCP:
  - UDP-based → works with hole punching (Spike 4 confirmed Cone NAT on SFR)
  - Multiplexed streams — each request is an independent QUIC stream
  - 0-RTT reconnection (connection resumption)
  - Built-in TLS 1.3

Wire protocol:
  - Each bidirectional QUIC stream carries one request/response exchange
  - Messages: length-prefixed msgpack (4-byte big-endian, same as TCP+TLS)
  - MNP handshake on stream 0 (control stream); subsequent streams = requests

Application protocol (MNP) is identical to TCP+TLS version.
The transport is the only change — all crypto, auth, and message types stay the same.
"""

import asyncio
import base64
import logging
import os
import struct
import subprocess
from pathlib import Path
from typing import Any, Callable

import blake3
import jwt
import msgpack
from aioquic.asyncio import QuicConnectionProtocol, serve
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.events import QuicEvent, StreamDataReceived, StreamReset
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

from meshbay_common import MNP_VERSION
from meshbay_node.roots import RootSet, entry_abs_path
from meshbay_common.handshake import (
    NONCE_LEN,
    ROLE_CLIENT,
    ROLE_NODE,
    HandshakeError,
    authorize_token,
    handshake_transcript,
    make_proof,
    quic_binding,
    verify_proof,
)
from meshbay_common.crypto import (
    sign_chunk,
    pk_to_b64,
)
from meshbay_common.webcrypto import chunk_key_aes as derive_chunk_key, encrypt_chunk_aes as encrypt_chunk
from meshbay_common.protocol import MNP
from meshbay_node.indexer import GroupIndex

log = logging.getLogger(__name__)

CHUNK_SIZE = 1024 * 1024
MAX_MSG    = 64 * 1024 * 1024
ALPN       = ["meshbay-mnp"]


class Denylist:
    """
    Denylist for revoked users, groups and invalidated JWTs.

    Finding H4: revocations used to live only in memory, so a node restart silently
    un-revoked everyone, and group revocations were dropped entirely — the hub
    signed and broadcast them but the node's handler only understood "user" and
    "jti". Now persisted to disk and group targets are honoured.
    """

    def __init__(self, path: Path | None = None):
        self.user_ids: set[str] = set()
        self.group_ids: set[str] = set()
        self.jtis: set[str] = set()
        self._path = path
        self._load()

    def is_denied(self, user_id: str, jti: str, group_id: str = "") -> bool:
        return (user_id in self.user_ids
                or jti in self.jtis
                or (bool(group_id) and group_id in self.group_ids))

    def deny_user(self, user_id: str) -> None:
        self.user_ids.add(user_id)
        log.info("Denied user: %s", user_id[:8])
        self._save()

    def deny_group(self, group_id: str) -> None:
        self.group_ids.add(group_id)
        log.info("Denied group: %s", group_id[:8])
        self._save()

    def deny_jti(self, jti: str) -> None:
        self.jtis.add(jti)
        log.info("Denied jti: %s", jti[:8])
        self._save()

    def entries(self) -> dict[str, list[str]]:
        """What is currently refused, for the operator to inspect (14.10)."""
        return {
            "users": sorted(self.user_ids),
            "groups": sorted(self.group_ids),
            "jtis": sorted(self.jtis),
        }

    def clear(self, subject: str = "") -> int:
        """
        Drop everything, or one identifier. Returns how many entries went.

        Not silent by design: clearing re-admits whoever it was keeping out, and
        the count is what tells the operator whether they undid one revocation
        or all of them.
        """
        before = len(self.user_ids) + len(self.group_ids) + len(self.jtis)
        if subject:
            self.user_ids.discard(subject)
            self.group_ids.discard(subject)
            self.jtis.discard(subject)
        else:
            self.user_ids.clear()
            self.group_ids.clear()
            self.jtis.clear()
        after = len(self.user_ids) + len(self.group_ids) + len(self.jtis)
        removed = before - after
        if removed:
            self._save()
        return removed

    def _load(self) -> None:
        if not self._path or not self._path.exists():
            return
        try:
            import json
            data = json.loads(self._path.read_text())
            self.user_ids = set(data.get("users", []))
            self.group_ids = set(data.get("groups", []))
            self.jtis = set(data.get("jtis", []))
            log.info("Denylist loaded: %d users, %d groups, %d jtis",
                     len(self.user_ids), len(self.group_ids), len(self.jtis))
        except Exception as e:
            log.warning("Could not load denylist from %s: %s", self._path, e)

    def _save(self) -> None:
        if not self._path:
            return
        try:
            import json
            self._path.parent.mkdir(parents=True, exist_ok=True)
            self._path.write_text(json.dumps({
                "users": sorted(self.user_ids),
                "groups": sorted(self.group_ids),
                "jtis": sorted(self.jtis),
            }))
        except Exception as e:
            log.warning("Could not persist denylist to %s: %s", self._path, e)


# ── Wire helpers ──────────────────────────────────────────────────────────────

def _pack(obj: dict) -> bytes:
    data = msgpack.packb(obj, use_bin_type=True)
    return struct.pack(">I", len(data)) + data


class _StreamBuffer:
    """Accumulate incoming QUIC stream data and extract length-prefixed messages."""

    def __init__(self):
        self._buf = bytearray()

    def feed(self, data: bytes):
        self._buf.extend(data)

    def messages(self):
        while len(self._buf) >= 4:
            length = struct.unpack(">I", self._buf[:4])[0]
            if length > MAX_MSG:
                raise ValueError(f"Message too large: {length}")
            if len(self._buf) < 4 + length:
                break
            msg_bytes = bytes(self._buf[4:4 + length])
            del self._buf[:4 + length]
            yield msgpack.unpackb(msg_bytes, raw=False)


# ── Per-connection server protocol ────────────────────────────────────────────

class _MNPServerProtocol(QuicConnectionProtocol):
    """
    One instance per QUIC connection.
    Handles the MNP handshake and all subsequent streams.
    """

    def __init__(self, *args, node_ctx: dict, **kwargs):
        super().__init__(*args, **kwargs)
        self._ctx        = node_ctx    # shared server context (keys, index, etc.)
        self._user_id: str | None = None
        self._group_id: str | None = None
        self._buffers: dict[int, _StreamBuffer] = {}
        self._nonce_client: bytes = b""
        self._gek_challenge: bytes | None = None
        self._pending = None

    def quic_event_received(self, event: QuicEvent) -> None:
        if isinstance(event, StreamDataReceived):
            sid = event.stream_id
            if sid not in self._buffers:
                self._buffers[sid] = _StreamBuffer()
            self._buffers[sid].feed(event.data)
            for msg in self._buffers[sid].messages():
                self._handle_message_sync(sid, msg)

        elif isinstance(event, StreamReset):
            self._buffers.pop(event.stream_id, None)

    def _handle_message_sync(self, stream_id: int, msg: dict) -> None:
        """Handle an MNP message synchronously (called from quic_event_received)."""
        mtype = msg.get("type")
        try:
            if mtype == MNP.HANDSHAKE:
                self._do_handshake_sync(stream_id, msg)
            elif mtype == MNP.HANDSHAKE_RESPONSE:
                self._do_handshake_response_sync(stream_id, msg)
            elif self._user_id is None:
                self._send(stream_id, {"type": "error", "detail": "Handshake required"})
            elif mtype == MNP.INDEX_SYNC:
                self._do_index_sync_sync(stream_id)
            elif mtype == MNP.FILE_REQUEST:
                self._do_file_request_sync(stream_id, msg)
            elif mtype == MNP.STREAM_SEGMENT:
                self._do_stream_segment_sync(stream_id, msg)
            elif mtype == MNP.CHAT_MESSAGE:
                self._do_chat_message_sync(stream_id, msg)
            elif mtype == MNP.PING:
                # Liveness is transport-agnostic, and a native client over QUIC
                # has the same half-open problem a DataChannel does.
                self._send(stream_id, {"type": MNP.PONG, "v": MNP_VERSION,
                                       "token": msg.get("token")})
            else:
                log.warning("Unknown MNP message type: %s", mtype)
        except Exception as e:
            log.error("Error handling %s: %s", mtype, e)
            self._send(stream_id, {"type": "error", "detail": str(e)})

    def _do_handshake_sync(self, stream_id: int, msg: dict) -> None:
        """
        Authorization half of the unified handshake (11.5.4).

        This used to be a second, weaker copy of the WebRTC logic: group_id was
        optional (so omitting it skipped the membership check entirely — M1),
        node-scoped daemon tokens were accepted as client tokens (M9), and the
        checks could drift from the WebRTC path independently. All of that now
        comes from meshbay_common.handshake, shared with WebRTC.

        NOT YET DONE — finding C6 remains open on this transport: there is still no
        GEK proof here, so a forged or stolen token reaches the node and can inject
        chat without holding the group key. The challenge/response and mutual node
        proof (quic_binding() is written and unit-tested for exactly this) are the
        remaining work in 11.5.4/5/6.
        """
        try:
            peer = authorize_token(
                msg.get("token", ""),
                self._ctx["hub_pk_pem"],
                group_id=msg.get("group_id", ""),
                hosted_groups=self._ctx.get("groups"),
                denylist=self._ctx.get("denylist"),
            )
        except HandshakeError as refusal:
            self._send(stream_id, {"type": "error", "detail": str(refusal)})
            self._quic.close()
            return

        try:
            self._nonce_client = base64.b64decode(msg.get("nonce", ""))
        except Exception:
            self._nonce_client = b""
        if len(self._nonce_client) < NONCE_LEN:
            self._send(stream_id, {"type": "error", "detail": "Client nonce required"})
            self._quic.close()
            return

        gctx = self._ctx["groups"][peer.group_id] if "groups" in self._ctx else self._ctx
        if not gctx.get("gek"):
            self._send(stream_id, {
                "type": "error",
                "detail": "Group encryption not initialized — contact node operator",
            })
            self._quic.close()
            return

        # Decoded but NOT authenticated: authentication is the GEK proof below.
        self._pending = peer
        self._gek_challenge = os.urandom(NONCE_LEN)
        self._send(stream_id, {
            "type":  MNP.HANDSHAKE_CHALLENGE,
            "v":     MNP_VERSION,
            "nonce": base64.b64encode(self._gek_challenge).decode(),
        })

    def _do_handshake_response_sync(self, stream_id: int, msg: dict) -> None:
        """Verify the client's GEK proof, then prove the node in return (C6, C3)."""
        if not self._gek_challenge or self._pending is None:
            self._send(stream_id, {"type": "error", "detail": "No pending handshake challenge"})
            return

        peer = self._pending
        gctx = self._ctx["groups"][peer.group_id] if "groups" in self._ctx else self._ctx
        gek = gctx.get("gek")
        if not gek:
            self._send(stream_id, {"type": "error", "detail": "Group encryption not initialized"})
            self._quic.close()
            return

        binding = self._ctx.get("server_cert_der")
        if not binding:
            # Refuse rather than fall back to an unbound proof (L4).
            self._send(stream_id, {"type": "error", "detail": "Channel binding unavailable"})
            self._quic.close()
            return
        binding = quic_binding(binding)

        try:
            proof = base64.b64decode(msg.get("proof", ""))
        except Exception:
            self._send(stream_id, {"type": "error", "detail": "Invalid proof encoding"})
            return

        if not verify_proof(gek, proof, ROLE_CLIENT, peer.group_id,
                            self._nonce_client, self._gek_challenge, binding):
            self._send(stream_id, {"type": "error", "detail": "GEK proof failed"})
            self._quic.close()
            return

        self._user_id = peer.user_id
        self._group_id = peer.group_id

        peers = self._ctx.get("_peers")
        if peers is not None:
            peers[self._user_id] = self

        transcript = handshake_transcript(
            ROLE_NODE, peer.group_id, self._nonce_client, self._gek_challenge, binding)
        node_proof = make_proof(
            gek, ROLE_NODE, peer.group_id, self._nonce_client, self._gek_challenge, binding)

        log.info("QUIC handshake OK — user=%s group=%s",
                 self._user_id[:8], self._group_id[:8])
        self._send(stream_id, {
            "type":    MNP.HANDSHAKE_ACK,
            "v":       MNP_VERSION,
            "node_pk": pk_to_b64(self._ctx["sk_node"].public_key()),
            "proof":   base64.b64encode(node_proof).decode(),
            "sig":     base64.b64encode(self._ctx["sk_node"].sign(transcript)).decode(),
        })
        self._gek_challenge = None

    def _group_ctx(self) -> dict:
        """Resolve the active group context (multi-group or legacy single-group)."""
        if "groups" in self._ctx and self._group_id:
            return self._ctx["groups"][self._group_id]
        return self._ctx

    def _do_index_sync_sync(self, stream_id: int) -> None:
        ctx = self._group_ctx()
        wire = ctx["index"].serialize()
        self._send(stream_id, {
            "type":      MNP.INDEX_SYNC,
            "v":         MNP_VERSION,
            "index_b64": base64.b64encode(wire).decode(),
        })

    def _do_file_request_sync(self, stream_id: int, msg: dict) -> None:
        """Serve file chunk synchronously (blocking I/O — acceptable for test sizes)."""
        ctx = self._group_ctx()
        file_id     = msg["file_id"]
        chunk_index = msg["chunk_index"]
        entry = ctx["index"].get_entry(file_id)
        if not entry:
            self._send(stream_id, {"type": "error", "detail": "File not found"})
            return

        file_path = entry_abs_path(ctx["roots"], entry)
        if not file_path.exists():
            self._send(stream_id, {"type": "error", "detail": "File not on disk"})
            return

        file_hash = bytes.fromhex(entry.id)
        chunk_data = _read_and_encrypt(
            self._ctx["sk_node"],
            ctx["gek"],
            file_path,
            chunk_index,
            file_hash,
        )
        self._send(stream_id, chunk_data)

    def _do_stream_segment_sync(self, stream_id: int, msg: dict) -> None:
        """Extract and serve one HLS segment via ffmpeg."""
        ctx = self._group_ctx()
        file_id = msg["file_id"]
        segment_index = msg["segment_index"]
        segment_duration = msg.get("segment_duration", 4)

        entry = ctx["index"].get_entry(file_id)
        if not entry:
            self._send(stream_id, {"type": "error", "detail": "File not found"})
            return

        file_path = entry_abs_path(ctx["roots"], entry)
        if not file_path.exists():
            self._send(stream_id, {"type": "error", "detail": "File not on disk"})
            return

        start_time = segment_index * segment_duration
        segment_data = _extract_segment(file_path, start_time, segment_duration)
        if segment_data is None:
            self._send(stream_id, {"type": "error", "detail": "Segment extraction failed"})
            return

        self._send(stream_id, {
            "type": MNP.STREAM_SEGMENT,
            "v": MNP_VERSION,
            "file_id": file_id,
            "segment_index": segment_index,
            "data_b64": base64.b64encode(segment_data).decode(),
            "size": len(segment_data),
        })

    def _do_chat_message_sync(self, stream_id: int, msg: dict) -> None:
        """Receive a chat message, store it, and broadcast to other connected peers."""
        chat_store = self._ctx.get("chat_store")
        if chat_store:
            import asyncio
            asyncio.ensure_future(chat_store.save_message(
                sender_id=msg.get("sender_id", self._user_id),
                iteration=msg.get("iteration", 0),
                payload=msg.get("payload", b"").encode() if isinstance(msg.get("payload"), str) else msg.get("payload", b""),
                thread_id=msg.get("thread_id"),
            ))

        peers = self._ctx.get("_peers", {})
        broadcast = {
            "type": MNP.CHAT_MESSAGE,
            "v": MNP_VERSION,
            "sender_id": msg.get("sender_id", self._user_id),
            "iteration": msg.get("iteration", 0),
            "payload": msg.get("payload", ""),
            "thread_id": msg.get("thread_id"),
            "group_id": self._group_id or "",
        }
        for uid, proto in peers.items():
            if uid != self._user_id and proto is not self:
                try:
                    proto._send(0, broadcast)
                except Exception:
                    pass

        self._send(stream_id, {"type": "ack", "v": MNP_VERSION})

    def connection_lost(self, exc) -> None:
        peers = self._ctx.get("_peers")
        if peers and self._user_id:
            peers.pop(self._user_id, None)
        super().connection_lost(exc)

    def _send(self, stream_id: int, obj: dict) -> None:
        self._quic.send_stream_data(stream_id, _pack(obj))
        self.transmit()


def _read_and_encrypt(
    sk_node: Ed25519PrivateKey,
    gek: bytes,
    file_path: Path,
    chunk_index: int,
    file_hash: bytes,
) -> dict:
    """Read and encrypt one chunk (blocking — runs in executor)."""
    with open(file_path, "rb") as f:
        f.seek(chunk_index * CHUNK_SIZE)
        plaintext = f.read(CHUNK_SIZE)

    pt_hash   = blake3.blake3(plaintext).digest()
    ckey      = derive_chunk_key(gek, file_hash, chunk_index)
    nonce, ct = encrypt_chunk(ckey, plaintext)
    ct_hash   = blake3.blake3(ct).digest()
    sig       = sign_chunk(sk_node, chunk_index, nonce, ct_hash)

    return {
        "type":           MNP.FILE_CHUNK,
        "v":              MNP_VERSION,
        "chunk_index":    chunk_index,
        "plaintext_size": len(plaintext),
        "nonce_b64":      base64.b64encode(nonce).decode(),
        "ct_b64":         base64.b64encode(ct).decode(),
        "ct_hash_b64":    base64.b64encode(ct_hash).decode(),
        "pt_hash_b64":    base64.b64encode(pt_hash).decode(),
        "sig_b64":        base64.b64encode(sig).decode(),
        "pk_node_b64":    pk_to_b64(sk_node.public_key()),
        "file_hash_b64":  base64.b64encode(file_hash).decode(),
    }


def _extract_segment(file_path: Path, start_time: float, duration: float) -> bytes | None:
    """Extract one HLS segment via ffmpeg. Returns MPEG-TS bytes or None on failure."""
    try:
        result = subprocess.run(
            ["ffmpeg", "-hide_banner", "-loglevel", "error",
             "-ss", str(start_time),
             "-i", str(file_path),
             "-t", str(duration),
             "-c:v", "copy", "-c:a", "copy",
             "-f", "mpegts", "pipe:1"],
            capture_output=True, timeout=30,
        )
        if result.returncode == 0 and result.stdout:
            return result.stdout
        return None
    except Exception:
        return None


# ── QuicChunkServer ────────────────────────────────────────────────────────────

class QuicChunkServer:
    """
    QUIC-based MNP chunk server (MNP v2).
    Drop-in replacement for ChunkServer with UDP transport.
    """

    def __init__(
        self,
        sk_node: Ed25519PrivateKey,
        hub_pk_pem: bytes,
        gek: bytes,
        roots: RootSet,
        index: GroupIndex,
        host: str = "::",   # listen IPv4 + IPv6 (dual-stack Linux)
        port: int = 19000,
        cert_path: Path | None = None,
        key_path: Path | None = None,
        groups: dict[str, dict] | None = None,
        denylist: Denylist | None = None,
    ):
        self._ctx = {
            "sk_node":     sk_node,
            "hub_pk_pem":  hub_pk_pem,
            "gek":         gek,
            "roots": roots,
            "index":       index,
        }
        if groups:
            self._ctx["groups"] = groups
        self._denylist = denylist or Denylist()
        self._ctx["denylist"] = self._denylist
        self._ctx["_peers"] = {}
        self._host      = host
        self._port      = port
        self._cert_path = cert_path or Path.home() / ".config/meshbay/node_tls.crt"
        self._key_path  = key_path  or Path.home() / ".config/meshbay/node_tls.key"
        self._server    = None
        self._task      = None
        self._session_tickets: dict[bytes, Any] = {}

    @property
    def port(self) -> int:
        return self._port

    @property
    def denylist(self) -> Denylist:
        return self._denylist

    def _make_config(self) -> QuicConfiguration:
        from meshbay_node.transport.tls_cert import generate_self_signed_cert
        if not self._cert_path.exists():
            generate_self_signed_cert(self._cert_path, self._key_path)
        config = QuicConfiguration(is_client=False, alpn_protocols=ALPN)
        config.load_cert_chain(str(self._cert_path), str(self._key_path))

        # Channel-binding anchor for the handshake proof (11.5.6). Read from our own
        # cert file — no aioquic internals needed on this side.
        from cryptography import x509
        from cryptography.hazmat.primitives import serialization as _ser
        self._ctx["server_cert_der"] = x509.load_pem_x509_certificate(
            self._cert_path.read_bytes()).public_bytes(_ser.Encoding.DER)
        return config

    def _store_ticket(self, ticket: Any) -> None:
        self._session_tickets[ticket.ticket] = ticket

    def _fetch_ticket(self, label: bytes) -> Any | None:
        return self._session_tickets.pop(label, None)

    async def start(self) -> None:
        config = self._make_config()
        ctx = self._ctx

        def protocol_factory(*args, **kwargs):
            return _MNPServerProtocol(*args, node_ctx=ctx, **kwargs)

        self._server = await serve(
            self._host, self._port,
            configuration=config,
            create_protocol=protocol_factory,
            session_ticket_handler=self._store_ticket,
            session_ticket_fetcher=self._fetch_ticket,
        )
        log.info("QuicChunkServer listening on %s:%d (QUIC/UDP)", self._host, self._port)

    def punch_nat(self, peer_ip: str, peer_port: int) -> None:
        """
        Send a probe UDP packet FROM the QUIC server's own socket.

        Critical for Port-Restricted Cone NAT (SFR résidentiel) :
        the NAT only allows inbound from (peer_ip, peer_port) if the server
        previously sent a packet TO (peer_ip, peer_port) from this same socket.

        The QUIC client must connect FROM peer_port for the NAT entry to match.
        """
        if self._server and hasattr(self._server, '_transport') and self._server._transport:
            self._server._transport.sendto(b'MESHBAY:NAT:PUNCH', (peer_ip, peer_port))
            log.info("NAT probe sent → %s:%d", peer_ip, peer_port)
        else:
            log.warning("punch_nat: server transport not available")

    async def stop(self) -> None:
        if self._server:
            self._server.close()
            self._server = None
            log.info("QuicChunkServer stopped")