Isolate TCP message framing
This commit is contained in:
@@ -0,0 +1,55 @@
|
||||
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
|
||||
Reference in New Issue
Block a user