From 5ed357208027899e57307a5df83ae3c8d81ae47b Mon Sep 17 00:00:00 2001
From: Richard Neumann <mail@richard-neumann.de>
Date: Mon, 24 Jul 2023 14:34:11 +0200
Subject: [PATCH] Acknowledge server messages

---
 rcon/battleye/client.py | 18 +++++++++++++-----
 rcon/battleye/proto.py  | 17 ++++++++++++++++-
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/rcon/battleye/client.py b/rcon/battleye/client.py
index 2f710c9..face11c 100644
--- a/rcon/battleye/client.py
+++ b/rcon/battleye/client.py
@@ -1,5 +1,6 @@
 """BattlEye RCon client."""
 
+from io import BufferedWriter
 from logging import getLogger
 from socket import SOCK_DGRAM
 from typing import Callable
@@ -11,6 +12,7 @@ from rcon.battleye.proto import LoginRequest
 from rcon.battleye.proto import Request
 from rcon.battleye.proto import Response
 from rcon.battleye.proto import ServerMessage
+from rcon.battleye.proto import ServerMessageAck
 from rcon.client import BaseClient
 from rcon.exceptions import WrongPassword
 
@@ -31,11 +33,14 @@ class Client(BaseClient, socket_type=SOCK_DGRAM):
     """BattlEye RCon client."""
 
     def __init__(
-            self, *args,
+            self,
+            *args,
+            max_length: int = 4096,
             message_handler: MessageHandler = log_message,
             **kwargs
     ):
         super().__init__(*args, **kwargs)
+        self.max_length = max_length
         self._handle_server_message = message_handler
 
     def _receive(self, max_length: int) -> Response:
@@ -46,9 +51,13 @@ class Client(BaseClient, socket_type=SOCK_DGRAM):
             )).type
         ].from_bytes(header, data[8:])
 
-    def receive(self, max_length: int = 4096) -> Response:
+    def receive(self, file: BufferedWriter) -> Response:
         """Receive a message."""
-        while isinstance(response := self._receive(max_length), ServerMessage):
+        while isinstance(
+                response := self._receive(self.max_length),
+                ServerMessage
+        ):
+            file.write(bytes(ServerMessageAck(response.seq)))
             self._handle_server_message(response)
 
         return response
@@ -57,8 +66,7 @@ class Client(BaseClient, socket_type=SOCK_DGRAM):
         """Send a request and receive a response."""
         with self._socket.makefile('wb') as file:
             file.write(bytes(request))
-
-        return self.receive()
+            return self.receive(file)
 
     def login(self, passwd: str) -> bool:
         """Log-in the user."""
diff --git a/rcon/battleye/proto.py b/rcon/battleye/proto.py
index 5f1bff0..66d7eac 100644
--- a/rcon/battleye/proto.py
+++ b/rcon/battleye/proto.py
@@ -13,6 +13,7 @@ __all__ = [
     'CommandRequest',
     'CommandResponse',
     'ServerMessage',
+    'ServerMessageAck',
     'Request',
     'Response'
 ]
@@ -172,7 +173,21 @@ class ServerMessage(NamedTuple):
         return self.payload.decode('ascii')
 
 
-Request = LoginRequest | CommandRequest
+class ServerMessageAck(NamedTuple):
+    """An acknowledgement of a message from the server."""
+
+    seq: int
+
+    def __bytes__(self):
+        return (0x02).to_bytes(1, 'little') + self.payload
+
+    @property
+    def payload(self) -> bytes:
+        """Return the payload."""
+        return self.seq.to_bytes(1, 'little')
+
+
+Request = LoginRequest | CommandRequest | ServerMessageAck
 Response = LoginResponse | CommandResponse | ServerMessage
 
 RESPONSE_TYPES = {