hms-mediaengine/packages/protocol/hms_protocol/framing.py

"""Length-prefixed MessagePack-Framing (ADR-0003).

4-Byte-Big-Endian-Länge, danach MessagePack-Payload. Maximale Payloadgröße
schützt vor unkontrollierten Queues/Resourcenerschöpfung (§6.2, §33).
"""

from __future__ import annotations

import asyncio
import socket
import struct

import msgpack

MAX_PAYLOAD_SIZE = 16 * 1024 * 1024  # 16 MiB Obergrenze je Nachricht
_LENGTH = struct.Struct(">I")


def encode_frame(payload: dict) -> bytes:
    """Serialisiert ein dict zu length-prefixed MessagePack."""
    body = msgpack.packb(payload, use_bin_type=True)
    if len(body) > MAX_PAYLOAD_SIZE:
        raise ValueError(f"payload too large: {len(body)} > {MAX_PAYLOAD_SIZE}")
    return _LENGTH.pack(len(body)) + body


def decode_frame(frame: bytes) -> dict:
    """Dekodiert einen vollständigen Frame (Länge + Body)."""
    if len(frame) < _LENGTH.size:
        raise ValueError("frame too short")
    (length,) = _LENGTH.unpack_from(frame, 0)
    if length > MAX_PAYLOAD_SIZE:
        raise ValueError(f"declared length {length} exceeds limit")
    body = frame[_LENGTH.size : _LENGTH.size + length]
    if len(body) != length:
        raise ValueError(f"truncated frame: expected {length}, got {len(body)}")
    return msgpack.unpackb(body, raw=False)


def read_frame(sock: socket.socket) -> dict:
    """Liest einen Frame von einem verbundenen Socket."""
    header = _recv_exact(sock, _LENGTH.size)
    (length,) = _LENGTH.unpack(header)
    if length > MAX_PAYLOAD_SIZE:
        raise ValueError(f"declared length {length} exceeds limit")
    body = _recv_exact(sock, length)
    return msgpack.unpackb(body, raw=False)


async def read_frame_async(reader: asyncio.StreamReader) -> dict:
    """Liest einen Frame von einem asyncio-StreamReader (IPC-Server/Client)."""
    header = await reader.readexactly(_LENGTH.size)
    (length,) = _LENGTH.unpack(header)
    if length > MAX_PAYLOAD_SIZE:
        raise ValueError(f"declared length {length} exceeds limit")
    body = await reader.readexactly(length)
    return msgpack.unpackb(body, raw=False)


def write_frame(sock: socket.socket, payload: dict) -> None:
    """Schreibt einen Frame auf einen verbundenen Socket."""
    sock.sendall(encode_frame(payload))


def _recv_exact(sock: socket.socket, count: int) -> bytes:
    buf = bytearray()
    while len(buf) < count:
        chunk = sock.recv(count - len(buf))
        if not chunk:
            raise ConnectionError("socket closed mid-frame")
        buf.extend(chunk)
    return bytes(buf)