Beta 13 - fix bugs (#11)

* Account for partial reads from socket

* Remove unused code

Author: VirxEC

rlbot/interface.py

@@ -10,7 +10,8 @@
 from rlbot import flat
 from rlbot.utils.logging import get_logger
 
-# We can connect to RLBotServer on this port.
+MAX_SIZE_2_BYTES = 2**16 - 1
+# The default port we can expect RLBotServer to be listening on.
 RLBOT_SERVER_PORT = 23234
 
 
@@ -38,30 +39,12 @@ class SocketDataType(IntEnum):
     CONTROLLABLE_TEAM_INFO = 15
 
 
-MAX_SIZE_2_BYTES = 2**16 - 1
-
-
-def int_to_bytes(val: int) -> bytes:
-    return val.to_bytes(2, byteorder="big")
-
-
-def int_from_bytes(bytes: bytes) -> int:
-    return int.from_bytes(bytes, "big")
-
-
 class SocketMessage:
     def __init__(self, type: int, data: bytes):
         self.type = SocketDataType(type)
         self.data = data
 
 
-def read_message_from_socket(s: socket) -> SocketMessage:
-    type_int = int_from_bytes(s.recv(2))
-    size = int_from_bytes(s.recv(2))
-    data = s.recv(size)
-    return SocketMessage(type_int, data)
-
-
 class SocketRelay:
     """
     The SocketRelay provides an abstraction over the direct communication with
@@ -97,11 +80,36 @@ def __init__(
         self.logger = get_logger("interface") if logger is None else logger
 
         self.socket = socket()
+
+        # Allow sending packets before getting a response from core
         self.socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)
 
     def __del__(self):
         self.socket.close()
 
+    @staticmethod
+    def _int_to_bytes(val: int) -> bytes:
+        return val.to_bytes(2, byteorder="big")
+
+    def _read_int(self) -> int:
+        return int.from_bytes(self._read_exact(2), "big")
+
+    def _read_exact(self, n: int) -> bytes:
+        buff = bytearray(n)
+        pos = 0
+        while pos < n:
+            cr = self.socket.recv_into(memoryview(buff)[pos:])
+            if cr == 0:
+                raise EOFError
+            pos += cr
+        return bytes(buff)
+
+    def read_message(self) -> SocketMessage:
+        type_int = self._read_int()
+        size = self._read_int()
+        data = self._read_exact(size)
+        return SocketMessage(type_int, data)
+
     def send_bytes(self, data: bytes, data_type: SocketDataType):
         assert self.is_connected, "Connection has not been established"
 
@@ -112,7 +120,7 @@ def send_bytes(self, data: bytes, data_type: SocketDataType):
             )
             return
 
-        message = int_to_bytes(data_type) + int_to_bytes(size) + data
+        message = self._int_to_bytes(data_type) + self._int_to_bytes(size) + data
         self.socket.sendall(message)
 
     def send_init_complete(self):
@@ -255,7 +263,7 @@ def handle_incoming_messages(self, blocking: bool = False) -> bool:
         assert self.is_connected, "Connection has not been established"
         try:
             self.socket.setblocking(blocking)
-            incoming_message = read_message_from_socket(self.socket)
+            incoming_message = self.read_message()
             try:
                 return self.handle_incoming_message(incoming_message)
             except flat.InvalidFlatbuffer as e:

rlbot/version.py

@@ -1,4 +1,4 @@
-__version__ = "5.0.0-beta.12"
+__version__ = "5.0.0-beta.13"
 
 
 RESET_SEQ = "\033[0m"

tests/run_only.py

@@ -4,7 +4,7 @@
 
 CURRENT_FILE = Path(__file__).parent
 
-MATCH_CONFIG_PATH = CURRENT_FILE / "human_vs_atba.toml"
+MATCH_CONFIG_PATH = CURRENT_FILE / "human_vs_necto.toml"
 RLBOT_SERVER_FOLDER = CURRENT_FILE / "../../core/RLBotCS/bin/Release/"
 
 if __name__ == "__main__":