56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import socket
|
|
|
|
from packet_codec import MAX_MESSAGE_BYTES
|
|
|
|
|
|
class TcpFramingError(ValueError):
|
|
pass
|
|
|
|
|
|
def frame_message(payload: bytes | bytearray | memoryview) -> bytes:
|
|
raw = bytes(payload)
|
|
if len(raw) > MAX_MESSAGE_BYTES:
|
|
raise TcpFramingError("TCP message exceeds maximum message size")
|
|
if b"\n" in raw or b"\r" in raw:
|
|
raise TcpFramingError("transport-neutral payload must not contain raw line delimiters")
|
|
return raw + b"\n"
|
|
|
|
|
|
def send_message(connection: socket.socket, payload: bytes | bytearray | memoryview) -> None:
|
|
connection.sendall(frame_message(payload))
|
|
|
|
|
|
class LineMessageBuffer:
|
|
def __init__(self) -> None:
|
|
self._buffer = bytearray()
|
|
|
|
@property
|
|
def buffered_bytes(self) -> int:
|
|
return len(self._buffer)
|
|
|
|
def clear(self) -> None:
|
|
self._buffer.clear()
|
|
|
|
def feed(self, chunk: bytes | bytearray | memoryview) -> list[bytes]:
|
|
raw = bytes(chunk)
|
|
if not raw:
|
|
return []
|
|
self._buffer.extend(raw)
|
|
messages: list[bytes] = []
|
|
while True:
|
|
newline = self._buffer.find(b"\n")
|
|
if newline < 0:
|
|
break
|
|
if newline > MAX_MESSAGE_BYTES:
|
|
raise TcpFramingError("TCP message exceeds maximum message size")
|
|
message = bytes(self._buffer[:newline])
|
|
del self._buffer[: newline + 1]
|
|
if message.endswith(b"\r"):
|
|
message = message[:-1]
|
|
messages.append(message)
|
|
if len(self._buffer) > MAX_MESSAGE_BYTES:
|
|
raise TcpFramingError("unterminated TCP message exceeds maximum message size")
|
|
return messages
|