protocol.py

# NetFS Protocol Module
# ABOUTME: Message framing, encoding/decoding for TCP communication with length-prefix framing

import json
import socket
import struct
from typing import Any

from netfs.common.errors import ErrorCode, NetFSError


class MessageTooLargeError(NetFSError):
    """Raised when a message exceeds the maximum allowed size."""

    def __init__(self, size: int, max_size: int):
        super().__init__(
            f"Message size {size} bytes exceeds maximum {max_size} bytes",
            ErrorCode.INTERNAL_ERROR,
        )
        self.size = size
        self.max_size = max_size


def encode_message(data: dict[str, Any], encoding: str = "utf-8") -> bytes:
    """
    Encode a dictionary as JSON with a 4-byte big-endian length prefix.

    Frame format: [4-byte length (big-endian)][JSON payload]

    Args:
        data: Dictionary to encode
        encoding: Character encoding (default: utf-8)

    Returns:
        Encoded bytes with length prefix
    """
    json_str = json.dumps(data)
    json_bytes = json_str.encode(encoding)
    length = len(json_bytes)

    # Pack length as 4-byte big-endian unsigned int
    length_prefix = struct.pack(">I", length)

    return length_prefix + json_bytes


def decode_message(data: bytes, encoding: str = "utf-8", max_size_mb: int = 16) -> dict[str, Any]:
    """
    Decode a message with 4-byte length prefix.

    Args:
        data: Encoded bytes with length prefix
        encoding: Character encoding (default: utf-8)
        max_size_mb: Maximum message size in MB

    Returns:
        Decoded dictionary

    Raises:
        MessageTooLargeError: If message exceeds max_size_mb
    """
    # Extract length from first 4 bytes
    length = struct.unpack(">I", data[:4])[0]

    # Check max size
    max_size_bytes = max_size_mb * 1024 * 1024
    if length > max_size_bytes:
        raise MessageTooLargeError(length, max_size_bytes)

    # Extract and decode JSON
    json_bytes = data[4 : 4 + length]
    json_str = json_bytes.decode(encoding)
    return json.loads(json_str)


def recv_exactly(sock: socket.socket, num_bytes: int) -> bytes:
    """
    Receive exactly num_bytes from socket, handling partial receives.

    This is critical for TCP because recv() may return fewer bytes than requested.

    Args:
        sock: Socket to receive from
        num_bytes: Exact number of bytes to receive

    Returns:
        Exactly num_bytes of data

    Raises:
        ConnectionError: If connection closes before receiving all bytes
    """
    buffer = b""

    while len(buffer) < num_bytes:
        chunk = sock.recv(num_bytes - len(buffer))
        if not chunk:
            raise ConnectionError("Connection closed before receiving complete message")
        buffer += chunk

    return buffer


def send_message(sock: socket.socket, data: dict[str, Any], encoding: str = "utf-8") -> None:
    """
    Send a message over a socket with length-prefix framing.

    Args:
        sock: Socket to send over
        data: Dictionary to send
        encoding: Character encoding (default: utf-8)
    """
    encoded = encode_message(data, encoding)
    sock.sendall(encoded)


def recv_message(
    sock: socket.socket, encoding: str = "utf-8", max_size_mb: int = 16
) -> dict[str, Any]:
    """
    Receive a message from a socket with length-prefix framing.

    Args:
        sock: Socket to receive from
        encoding: Character encoding (default: utf-8)
        max_size_mb: Maximum message size in MB

    Returns:
        Decoded dictionary

    Raises:
        MessageTooLargeError: If message exceeds max_size_mb
        ConnectionError: If connection closes unexpectedly
    """
    # First, receive the 4-byte length prefix
    length_bytes = recv_exactly(sock, 4)
    length = struct.unpack(">I", length_bytes)[0]

    # Check max size before receiving full message
    max_size_bytes = max_size_mb * 1024 * 1024
    if length > max_size_bytes:
        raise MessageTooLargeError(length, max_size_bytes)

    # Receive the JSON payload
    json_bytes = recv_exactly(sock, length)

    # Decode and parse
    json_str = json_bytes.decode(encoding)
    return json.loads(json_str)