From 89e3d4b1c71a8a4b1e66deef4dbcc4dfcd614464 Mon Sep 17 00:00:00 2001 From: Richard Neumann Date: Sun, 14 Aug 2022 23:28:07 +0200 Subject: [PATCH] Attempt to read following packets on large responses --- rcon/source/client.py | 2 +- rcon/source/proto.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/rcon/source/client.py b/rcon/source/client.py index d03ff46..7e992ea 100644 --- a/rcon/source/client.py +++ b/rcon/source/client.py @@ -23,7 +23,7 @@ class Client(BaseClient, socket_type=SOCK_STREAM): def read(self) -> Packet: """Read a packet.""" with self._socket.makefile('rb') as file: - return Packet.read(file) + return Packet.read(file, max_pkg_size=self.max_pkg_size) def login(self, passwd: str, *, encoding: str = 'utf-8') -> bool: """Perform a login.""" diff --git a/rcon/source/proto.py b/rcon/source/proto.py index fe5c845..4c00af9 100644 --- a/rcon/source/proto.py +++ b/rcon/source/proto.py @@ -103,17 +103,21 @@ class Packet(NamedTuple): return cls(id_, type_, payload, terminator) @classmethod - def read(cls, file: IO) -> Packet: + def read(cls, file: IO, *, max_pkg_size: int | None = None) -> Packet: """Read a packet from a file-like object.""" size = LittleEndianSignedInt32.read(file) id_ = LittleEndianSignedInt32.read(file) type_ = Type.read(file) - payload = file.read(size - 10) + payload = file.read(size := size - 10) terminator = file.read(2) if terminator != TERMINATOR: LOGGER.warning('Unexpected terminator: %s', terminator) + # Attempt to read following packets on large responses. + if size >= max_pkg_size: + payload += cls.read(file, max_pkg_size=max_pkg_size).payload + return cls(id_, type_, payload, terminator) @classmethod