Source code for edri.protocol.frame

"""Protocol v2 frame: one 48-byte header, identical in both directions.

::

    offset  size  field
    0       1     version
    1       1     type
    2       2     flags
    4       8     seq
    12      4     length
    16      16    target
    32      16    mac
    48      ...   payload

The MAC covers ``header[0:32] ‖ payload`` — every header byte except the MAC
field itself — so type, sequence and target are all authenticated, not just the
body.
"""
from dataclasses import dataclass
from enum import IntEnum, unique
from hmac import compare_digest, new as hmac_new
from socket import socket
from struct import error as StructError, pack, unpack

from edri.protocol.errors import ErrorReason, ProtocolError

__all__ = [
    "HEADER_LENGTH", "MAC_LENGTH", "MAC_OFFSET", "NONCE_LENGTH",
    "MAX_LENGTH_FIELD", "PROTOCOL_VERSION", "ZERO_MAC", "ZERO_TARGET",
    "Frame", "FrameType", "compute_mac", "decode_header", "encode_frame",
    "verify_frame", "read_exactly", "read_frame",
]

PROTOCOL_VERSION = 2
HEADER_FORMAT = "!BBHQI16s16s"
HEADER_LENGTH = 48
MAC_OFFSET = 32
MAC_LENGTH = 16
NONCE_LENGTH = 32
MAX_LENGTH_FIELD = 0xFFFFFFFF

ZERO_TARGET = bytes(16)
ZERO_MAC = bytes(MAC_LENGTH)




@unique
class FrameType(IntEnum):
    """Frame types occupy their own space, no longer shared with router UUIDs."""

    HELLO_SERVER = 0x01
    HELLO_CLIENT = 0x02
    HELLO_ACCEPT = 0x03
    REKEY = 0x04  # reserved for a forced re-handshake; not implemented
    EVENT = 0x10
    DEMANDS = 0x11
    LAST_EVENTS = 0x12
    ACK = 0x13
    ERROR = 0x20





@dataclass(frozen=True)
class Frame:
    type: FrameType
    seq: int
    target: bytes
    payload: bytes
    flags: int = 0





def compute_mac(key: bytes, header_without_mac: bytes, payload: bytes) -> bytes:
    return hmac_new(key, header_without_mac + payload, "sha256").digest()[:MAC_LENGTH]





def encode_frame(frame: Frame, key: bytes | None) -> bytes:
    """Serialize ``frame``; ``key`` of ``None`` zeroes the MAC (HELLO_SERVER only)."""
    length = len(frame.payload)
    if length > MAX_LENGTH_FIELD:
        raise ProtocolError(f"payload of {length} bytes exceeds the length field", ErrorReason.TOO_LARGE)
    header_without_mac = pack(
        "!BBHQI16s", PROTOCOL_VERSION, frame.type, frame.flags, frame.seq, length, frame.target
    )
    mac = ZERO_MAC if key is None else compute_mac(key, header_without_mac, frame.payload)
    return header_without_mac + mac + frame.payload





def decode_header(header: bytes) -> tuple[int, FrameType, int, int, int, bytes, bytes]:
    """Split a 48-byte header. Does *not* verify the MAC — see :func:`verify_frame`."""
    try:
        version, raw_type, flags, seq, length, target, mac = unpack(HEADER_FORMAT, header)
    except StructError as e:
        raise ProtocolError(f"malformed header: {e}", ErrorReason.BAD_VERSION) from e
    if version != PROTOCOL_VERSION:
        raise ProtocolError(f"unsupported protocol version {version}", ErrorReason.BAD_VERSION)
    try:
        frame_type = FrameType(raw_type)
    except ValueError as e:
        raise ProtocolError(f"unknown frame type {raw_type:#x}", ErrorReason.BAD_VERSION) from e
    return version, frame_type, flags, seq, length, target, mac





def verify_frame(header: bytes, payload: bytes, key: bytes) -> None:
    """Raise :class:`ProtocolError` unless ``header`` + ``payload`` authenticate."""
    expected = compute_mac(key, header[:MAC_OFFSET], payload)
    if not compare_digest(expected, header[MAC_OFFSET:HEADER_LENGTH]):
        raise ProtocolError("frame MAC does not match", ErrorReason.BAD_MAC)





def read_exactly(sock: socket, amount: int) -> bytes:
    """Read exactly ``amount`` bytes or raise.

    ``recv`` may return fewer bytes than asked for at any time; the pre-v2 code
    assumed otherwise in three places, which is how a frame could be parsed from
    a partial read.
    """
    chunks: list[bytes] = []
    remaining = amount
    while remaining > 0:
        chunk = sock.recv(remaining)
        if not chunk:
            raise ConnectionError(f"connection closed with {remaining} of {amount} bytes outstanding")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)





def read_frame(sock: socket, max_message: int) -> tuple[bytes, bytes]:
    """Read one frame. The length is validated *before* the payload is read."""
    header = read_exactly(sock, HEADER_LENGTH)
    _, _, _, _, length, _, _ = decode_header(header)
    if length > max_message:
        raise ProtocolError(
            f"declared payload of {length} bytes exceeds the {max_message} byte ceiling",
            ErrorReason.TOO_LARGE,
        )
    payload = read_exactly(sock, length) if length else b""
    return header, payload