ZKAC/python/zkac/udp.py

"""
Length-prefixed UDP datagram framing for ZKAC handshakes and encrypted sessions.

Wire format matches :mod:`zkac.tcp`: each datagram is ``uint32_le(length) || payload``
with *length* counting only *payload* bytes. **One datagram = one frame** (do not
split a frame across packets).

**Reliability:** UDP is unordered and lossy. This module does not add ACKs or
retransmits. Use TCP (``zkac.tcp``) if you need a reliable stream without
building your own reliability layer.

**Size:** Large BBS+ auth packets can exceed typical path MTUs (~1500 B). If
``send`` raises ``OSError`` (e.g. ``EMSGSIZE``), use TCP or reduce proof size /
raise MTU on controlled networks.
"""

from __future__ import annotations

import socket
import struct
from typing import TYPE_CHECKING, Optional, Tuple

from zkac import MAX_BBS_AUTH_PROOF_BYTES

if TYPE_CHECKING:
    from zkac import Credential, Node, PublicKey, RoleRegistry, Session

# Same logical cap as tcp framing; note UDP + large proofs may hit EMSGSIZE on send.
MAX_UDP_FRAME_BYTES: int = MAX_BBS_AUTH_PROOF_BYTES + 4096

# IPv4 max UDP payload (theoretical); recv buffer size hint.
MAX_UDP_DATAGRAM_BYTES: int = 65507

_HANDSHAKE_MSG_LEN = 32


def _build_framed_datagram(payload: bytes) -> bytes:
    if len(payload) > MAX_UDP_FRAME_BYTES:
        raise ValueError(
            f"payload length {len(payload)} exceeds maximum ({MAX_UDP_FRAME_BYTES})"
        )
    return struct.pack("<I", len(payload)) + payload


def _parse_framed_datagram(data: bytes) -> bytes:
    if len(data) < 4:
        raise ValueError("datagram too short for length prefix")
    (length,) = struct.unpack("<I", data[:4])
    if length > MAX_UDP_FRAME_BYTES:
        raise ValueError(f"frame length {length} exceeds maximum ({MAX_UDP_FRAME_BYTES})")
    if len(data) != 4 + length:
        raise ValueError(
            f"datagram size mismatch: expected {4 + length} bytes, got {len(data)}"
        )
    return data[4:] if length else b""


def write_datagram(sock: socket.socket, payload: bytes, addr: Optional[tuple] = None) -> None:
    """
    Send one framed datagram. If *addr* is ``None``, *sock* must be connected
    (e.g. after :meth:`socket.socket.connect`).
    """
    packet = _build_framed_datagram(payload)
    if len(packet) > MAX_UDP_DATAGRAM_BYTES:
        raise ValueError("framed datagram exceeds maximum UDP payload size")
    if addr is not None:
        sock.sendto(packet, addr)
    else:
        sock.send(packet)


def read_datagram(sock: socket.socket, bufsize: int = MAX_UDP_DATAGRAM_BYTES) -> bytes:
    """
    Receive one framed datagram on a **connected** UDP socket (``recv``).
    """
    data = sock.recv(bufsize)
    if not data:
        raise ConnectionError("received empty datagram (peer closed?)")
    return _parse_framed_datagram(data)


def read_datagram_from(
    sock: socket.socket, bufsize: int = MAX_UDP_DATAGRAM_BYTES
) -> Tuple[bytes, tuple]:
    """
    Receive one framed datagram on an **unconnected** UDP socket (``recvfrom``).
    Returns ``(payload, addr)``.
    """
    data, addr = sock.recvfrom(bufsize)
    if not data:
        raise ConnectionError("received empty datagram")
    return _parse_framed_datagram(data), addr


def client_handshake(
    sock: socket.socket,
    server_addr: tuple,
    node: Node,
    expected_server_pk: PublicKey,
    credential: Credential,
) -> Session:
    """
    Run the ZKAC client side over UDP. Connects *sock* to *server_addr* and
    exchanges three framed datagrams (init → server bundle → auth).

    *server_addr* is ``(host, port)`` for :meth:`socket.socket.connect`.
    """
    sock.connect(server_addr)

    pending, init_msg = node.connect()
    if len(init_msg) != _HANDSHAKE_MSG_LEN:
        raise ValueError("internal error: init_msg must be 32 bytes")
    write_datagram(sock, init_msg)

    bundle = read_datagram(sock)
    if len(bundle) < _HANDSHAKE_MSG_LEN:
        raise ValueError("server handshake bundle too short")
    response_msg = bundle[:_HANDSHAKE_MSG_LEN]
    identity_proof = bundle[_HANDSHAKE_MSG_LEN:]

    session, auth_packet = node.complete_connect(
        pending, response_msg, identity_proof, expected_server_pk, credential
    )
    write_datagram(sock, auth_packet)
    return session


def server_handshake(
    sock: socket.socket,
    node: Node,
    registry: RoleRegistry,
) -> Tuple[Session, bytes, tuple]:
    """
    Run the ZKAC server side over UDP. Waits for the first datagram, then
    :meth:`socket.socket.connect` to that peer so the rest of the handshake
    uses the same path.

    Returns ``(session, role_id, client_addr)``.
    """
    init_msg, client_addr = read_datagram_from(sock)
    if len(init_msg) != _HANDSHAKE_MSG_LEN:
        raise ValueError("init_msg must be 32 bytes")

    sock.connect(client_addr)

    session, response_msg = node.accept(init_msg)
    if len(response_msg) != _HANDSHAKE_MSG_LEN:
        raise ValueError("internal error: response_msg must be 32 bytes")

    identity_proof = node.prove_identity(session)
    bundle = response_msg + identity_proof
    write_datagram(sock, bundle)

    auth_packet = read_datagram(sock)
    role_id = node.verify_auth(session, auth_packet, registry)
    return session, role_id, client_addr


class FramedSession:
    """
    One ZKAC ciphertext per UDP datagram; *sock* must be connected.
    """

    def __init__(self, sock: socket.socket, session: Session) -> None:
        self._sock = sock
        self._session = session

    @property
    def session(self) -> Session:
        return self._session

    def send(self, plaintext: bytes) -> None:
        packet = self._session.encrypt(plaintext)
        write_datagram(self._sock, packet)

    def recv(self) -> bytes:
        return self._session.decrypt(read_datagram(self._sock))