From 1552a58e84821becd401af215ac470f3296f2340 Mon Sep 17 00:00:00 2001
From: Richard Neumann <mail@richard-neumann.de>
Date: Mon, 24 Jul 2023 18:15:29 +0200
Subject: [PATCH] Rewrite battleye client

---
 rcon/battleye/client.py | 54 +++++++++++++++++++++++------------------
 1 file changed, 30 insertions(+), 24 deletions(-)

diff --git a/rcon/battleye/client.py b/rcon/battleye/client.py
index 2309512..6568324 100644
--- a/rcon/battleye/client.py
+++ b/rcon/battleye/client.py
@@ -1,6 +1,6 @@
 """BattlEye RCon client."""
 
-from io import BufferedWriter
+from collections import defaultdict
 from logging import getLogger
 from socket import SOCK_DGRAM
 from typing import Callable
@@ -43,39 +43,45 @@ class Client(BaseClient, socket_type=SOCK_DGRAM):
         self.max_length = max_length
         self._handle_server_message = message_handler
 
-    def _receive(self, max_length: int) -> Response:
+    def receive(self) -> Response:
         """Receive a packet."""
         return RESPONSE_TYPES[
             (header := Header.from_bytes(
-                (data := self._socket.recv(max_length))[:8]
+                (data := self._socket.recv(self.max_length))[:8]
             )).type
         ].from_bytes(header, data[8:])
 
-    def receive(self, file: BufferedWriter) -> Response | str:
-        """Receive a message."""
-        server_messages = set()
+    def communicate(self, request: Request) -> Response | str:
+        """Send a request and receive a response."""
+        acknowledged = defaultdict(set)
+        messages = []
 
-        while isinstance(
-                response := self._receive(self.max_length),
-                ServerMessage
-        ):
-            server_messages.add(response)
-            file.write(bytes(ServerMessageAck(response.seq)))
-            self._handle_server_message(response)
+        while True:
+            with self._socket.makefile('wb') as file:
+                file.write(bytes(request))
 
-        if not server_messages:
-            return response
+            response = self.receive()
 
-        return ''.join(
-            msg.message for msg in
-            sorted(server_messages, key=lambda msg: msg.seq)
-        )
+            try:
+                seq = response.seq
+            except AttributeError:
+                return response
 
-    def communicate(self, request: Request) -> Response | str:
-        """Send a request and receive a response."""
-        with self._socket.makefile('wb') as file:
-            file.write(bytes(request))
-            return self.receive(file)
+            if seq in acknowledged[msg_type := type(response)]:
+                break
+            else:
+                acknowledged[msg_type].add(seq)
+
+            if isinstance(response, ServerMessage):
+                self._handle_server_message(response)
+
+                with self._socket.makefile('wb') as file:
+                    file.write(bytes(ServerMessageAck(response.seq)))
+
+            else:
+                messages.append(response)
+
+        return ''.join(msg.message for msg in sorted(messages, key=lambda msg: msg.seq))
 
     def login(self, passwd: str) -> bool:
         """Log-in the user."""