codexPySnake/server/webtransport_server.py

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import Awaitable, Callable, Dict, Optional

from .transport import DatagramServerTransport, OnDatagram, TransportPeer
import logging


try:
    from aioquic.asyncio import QuicConnectionProtocol, serve
    from aioquic.quic.configuration import QuicConfiguration
    from aioquic.h3.connection import H3Connection
    from aioquic.h3.events import HeadersReceived
    # Datagram event names vary by aioquic version; try both
    try:
        from aioquic.h3.events import DatagramReceived as H3DatagramReceived  # type: ignore
    except Exception:  # pragma: no cover
        H3DatagramReceived = object  # type: ignore
except Exception:  # pragma: no cover - optional dependency not installed
    QuicConnectionProtocol = object  # type: ignore
    QuicConfiguration = object  # type: ignore
    H3Connection = object  # type: ignore
    HeadersReceived = object  # type: ignore
    H3DatagramReceived = object  # type: ignore
    serve = None  # type: ignore


@dataclass
class WTSession:
    flow_id: int  # HTTP/3 datagram flow id (usually CONNECT stream id)
    proto: "GameWTProtocol"


class GameWTProtocol(QuicConnectionProtocol):  # type: ignore[misc]
    def __init__(self, *args, on_datagram: OnDatagram, sessions: Dict[int, WTSession], **kwargs):
        super().__init__(*args, **kwargs)
        self._on_datagram = on_datagram
        self._sessions = sessions
        self._http: Optional[H3Connection] = None

    def http_event_received(self, event) -> None:  # type: ignore[override]
        # Headers for CONNECT :protocol = webtransport open a session
        if isinstance(event, HeadersReceived):
            headers = {k.decode().lower(): v.decode() for k, v in event.headers}
            method = headers.get(":method")
            protocol = headers.get(":protocol") or headers.get("sec-webtransport-protocol")
            if method == "CONNECT" and (protocol == "webtransport"):
                # In WebTransport over H3, datagrams use the CONNECT stream id as flow_id
                flow_id = event.stream_id  # type: ignore[attr-defined]
                self._sessions[flow_id] = WTSession(flow_id=flow_id, proto=self)
                logging.info("WT CONNECT accepted: flow_id=%s", flow_id)
                # Send 2xx to accept the session
                if self._http is not None:
                    self._http.send_headers(event.stream_id, [(b":status", b"200")])
        elif isinstance(event, H3DatagramReceived):  # type: ignore[misc]
            # Route datagram to session by flow_id
            flow_id = getattr(event, "flow_id", None)
            data = getattr(event, "data", None)
            if flow_id is None or data is None:
                return
            sess = self._sessions.get(flow_id)
            if not sess:
                return
            peer = TransportPeer(addr=(self, flow_id))
            logging.debug("WT datagram received: flow_id=%s, %d bytes", flow_id, len(data))
            asyncio.ensure_future(self._on_datagram(bytes(data), peer))

    def quic_event_received(self, event) -> None:  # type: ignore[override]
        # Lazily create H3 connection wrapper
        if self._http is None and hasattr(self, "_quic"):
            try:
                self._http = H3Connection(self._quic, enable_webtransport=True)  # type: ignore[attr-defined]
            except Exception:
                self._http = None
        if self._http is not None:
            for http_event in self._http.handle_event(event):
                self.http_event_received(http_event)

    async def send_h3_datagram(self, flow_id: int, data: bytes) -> None:
        if self._http is None:
            return
        try:
            logging.debug("WT send datagram: flow_id=%s, %d bytes", flow_id, len(data))
            self._http.send_datagram(flow_id, data)
            await self._loop.run_in_executor(None, self.transmit)  # type: ignore[attr-defined]
        except Exception:
            pass


class WebTransportServer(DatagramServerTransport):
    """HTTP/3 WebTransport datagram server using aioquic.

    Accepts CONNECT requests with :protocol = webtransport and exchanges
    RFC 9297 HTTP/3 datagrams bound to the CONNECT stream id.
    """

    def __init__(self, host: str, port: int, certfile: str, keyfile: str, on_datagram: OnDatagram):
        if serve is None:
            raise RuntimeError("aioquic is not installed. Please `pip install aioquic`.")
        self.host = host
        self.port = port
        self.certfile = certfile
        self.keyfile = keyfile
        self._on_datagram = on_datagram
        self._server = None
        self._sessions: Dict[int, WTSession] = {}

    async def send(self, data: bytes, peer: TransportPeer) -> None:
        # peer.addr is (protocol, flow_id)
        if isinstance(peer.addr, tuple) and len(peer.addr) == 2:
            proto, flow_id = peer.addr
            if isinstance(proto, GameWTProtocol) and isinstance(flow_id, int):
                await proto.send_h3_datagram(flow_id, data)

    async def run(self) -> None:
        configuration = QuicConfiguration(is_client=False, alpn_protocols=["h3"])
        configuration.load_cert_chain(self.certfile, self.keyfile)

        async def _create_protocol(*args, **kwargs):
            return GameWTProtocol(*args, on_datagram=self._on_datagram, sessions=self._sessions, **kwargs)

        logging.info("WebTransport (H3) server listening on %s:%d", self.host, self.port)
        self._server = await serve(self.host, self.port, configuration=configuration, create_protocol=_create_protocol)
        try:
            await self._server.wait_closed()
        finally:
            self._server.close()