Fix socket not close issue when error

Author: BattlefieldDuck

opengsq/protocols/gamespy1.py

@@ -115,14 +115,12 @@ async def __get_packets_response(self, sock):
 
     async def __connect_and_send(self, data) -> BinaryReader:
         # Connect to remote host
-        sock = SocketAsync()
-        sock.settimeout(self._timeout)
-        await sock.connect((self._address, self._query_port))
+        with SocketAsync() as sock:
+            sock.settimeout(self._timeout)
+            await sock.connect((self._address, self._query_port))
 
-        sock.send(data)
-        br = BinaryReader(await self.__get_packets_response(sock))
-
-        sock.close()
+            sock.send(data)
+            br = BinaryReader(await self.__get_packets_response(sock))
 
         return br
 

opengsq/protocols/gamespy2.py

@@ -21,17 +21,15 @@ def __init__(self, address: str, query_port: int, timeout: float = 5.0):
     async def get_status(self, request: Request = Request.INFO | Request.PLAYERS | Request.TEAMS) -> dict:
         """Retrieves information about the server including, Info, Players, and Teams."""
         # Connect to remote host
-        sock = SocketAsync()
-        sock.settimeout(self._timeout)
-        await sock.connect((self._address, self._query_port))
+        with SocketAsync() as sock:
+            sock.settimeout(self._timeout)
+            await sock.connect((self._address, self._query_port))
 
-        # Send Request
-        sock.send(b'\xFE\xFD\x00\x04\x05\x06\x07' + self.__get_request_bytes(request))
+            # Send Request
+            sock.send(b'\xFE\xFD\x00\x04\x05\x06\x07' + self.__get_request_bytes(request))
 
-        # Server response
-        response = await sock.recv()
-
-        sock.close()
+            # Server response
+            response = await sock.recv()
 
         # Remove the first 5 bytes { 0x00, 0x04, 0x05, 0x06, 0x07 }
         br = BinaryReader(response[5:])

opengsq/protocols/gamespy3.py

@@ -17,39 +17,37 @@ def __init__(self, address: str, query_port: int, timeout: float = 5.0):
     async def get_status(self):
         """Retrieves information about the server including, Info, Players, and Teams."""
         # Connect to remote host
-        sock = SocketAsync()
-        sock.settimeout(self._timeout)
-        await sock.connect((self._address, self._query_port))
+        with SocketAsync() as sock:
+            sock.settimeout(self._timeout)
+            await sock.connect((self._address, self._query_port))
 
-        request_h = b'\xFE\xFD'
-        timestamp = b'\x04\x05\x06\x07'
-        challenge = b''
+            request_h = b'\xFE\xFD'
+            timestamp = b'\x04\x05\x06\x07'
+            challenge = b''
 
-        if self.challenge:
-            # Packet 1: Initial request - (https://wiki.unrealadmin.org/UT3_query_protocol#Packet_1:_Initial_request)
-            sock.send(request_h + b'\x09' + timestamp)
+            if self.challenge:
+                # Packet 1: Initial request - (https://wiki.unrealadmin.org/UT3_query_protocol#Packet_1:_Initial_request)
+                sock.send(request_h + b'\x09' + timestamp)
 
-            # Packet 2: First response - (https://wiki.unrealadmin.org/UT3_query_protocol#Packet_2:_First_response)
-            response = await sock.recv()
-
-            if response[0] != 9:
-                raise InvalidPacketException(
-                    'Packet header mismatch. Received: {}. Expected: {}.'
-                    .format(chr(response[0]), chr(9))
-                )
+                # Packet 2: First response - (https://wiki.unrealadmin.org/UT3_query_protocol#Packet_2:_First_response)
+                response = await sock.recv()
 
-            # Packet 3: Second request - (http://wiki.unrealadmin.org/UT3_query_protocol#Packet_3:_Second_request)
-            challenge = int(response[5:].decode('ascii').strip('\x00'))
-            challenge = b'' if challenge == 0 else challenge.to_bytes(4, 'big', signed=True)
+                if response[0] != 9:
+                    raise InvalidPacketException(
+                        'Packet header mismatch. Received: {}. Expected: {}.'
+                        .format(chr(response[0]), chr(9))
+                    )
 
-        request_data = request_h + b'\x00' + timestamp + challenge
-        sock.send(request_data + b'\xFF\xFF\xFF\x01')
+                # Packet 3: Second request - (http://wiki.unrealadmin.org/UT3_query_protocol#Packet_3:_Second_request)
+                challenge = int(response[5:].decode('ascii').strip('\x00'))
+                challenge = b'' if challenge == 0 else challenge.to_bytes(4, 'big', signed=True)
 
-        # Packet 4: Server information response
-        # (http://wiki.unrealadmin.org/UT3_query_protocol#Packet_4:_Server_information_response)
-        response = await self.__read(sock)
+            request_data = request_h + b'\x00' + timestamp + challenge
+            sock.send(request_data + b'\xFF\xFF\xFF\x01')
 
-        sock.close()
+            # Packet 4: Server information response
+            # (http://wiki.unrealadmin.org/UT3_query_protocol#Packet_4:_Server_information_response)
+            response = await self.__read(sock)
 
         br = BinaryReader(response)
 

opengsq/protocols/quake1.py

@@ -89,19 +89,17 @@ def _get_player_match_collections(self, br: BinaryReader):
 
     async def _connect_and_send(self, data):
         # Connect to remote host
-        sock = SocketAsync()
-        sock.settimeout(self._timeout)
-        await sock.connect((self._address, self._query_port))
+        with SocketAsync() as sock:
+            sock.settimeout(self._timeout)
+            await sock.connect((self._address, self._query_port))
 
-        header = b'\xFF\xFF\xFF\xFF'
+            header = b'\xFF\xFF\xFF\xFF'
 
-        # Send Request
-        sock.send(header + data + b'\x00')
+            # Send Request
+            sock.send(header + data + b'\x00')
 
-        # Server response
-        response_data = await sock.recv()
-
-        sock.close()
+            # Server response
+            response_data = await sock.recv()
 
         # Remove the last 0x00 if exists (Only if Quake1)
         if response_data[-1] == 0:

opengsq/protocols/source.py

@@ -189,33 +189,31 @@ async def get_rules(self) -> dict:
 
     async def __connect_and_send_challenge(self, header: __RequestHeader) -> bytes:
         # Connect to remote host
-        sock = SocketAsync()
-        sock.settimeout(self._timeout)
-        await sock.connect((self._address, self._query_port))
-        
-        # Send and receive
-        request_base = b'\xFF\xFF\xFF\xFF' + header
-        request_data = request_base
-
-        if len(self._challenge) > 0:
-            request_data += self._challenge
-        elif header != self.__RequestHeader.A2S_INFO:
-            request_data += b'\xFF\xFF\xFF\xFF'
+        with SocketAsync() as sock:
+            sock.settimeout(self._timeout)
+            await sock.connect((self._address, self._query_port))
+            
+            # Send and receive
+            request_base = b'\xFF\xFF\xFF\xFF' + header
+            request_data = request_base
 
-        sock.send(request_data)
-        response_data = await self.__receive(sock)
-        br = BinaryReader(response_data)
-        header = br.read_byte()
+            if len(self._challenge) > 0:
+                request_data += self._challenge
+            elif header != self.__RequestHeader.A2S_INFO:
+                request_data += b'\xFF\xFF\xFF\xFF'
 
-        # The server may reply with a challenge
-        if header == self.__ResponseHeader.S2C_CHALLENGE:
-            self._challenge = br.read()
-            
-            # Send the challenge and receive
-            sock.send(request_base + self._challenge)
+            sock.send(request_data)
             response_data = await self.__receive(sock)
-            
-        sock.close()
+            br = BinaryReader(response_data)
+            header = br.read_byte()
+
+            # The server may reply with a challenge
+            if header == self.__ResponseHeader.S2C_CHALLENGE:
+                self._challenge = br.read()
+                
+                # Send the challenge and receive
+                sock.send(request_base + self._challenge)
+                response_data = await self.__receive(sock)
 
         return response_data
 

opengsq/socket_async.py

@@ -7,7 +7,7 @@ class SocketAsync():
     class SocketKind(Enum):
         SOCK_STREAM = auto() 
         SOCK_DGRAM = auto()
-    
+
     @staticmethod
     def gethostbyname(hostname: str) -> str:
         return socket.gethostbyname(hostname)
@@ -18,6 +18,12 @@ def __init__(self, kind: SocketKind = SocketKind.SOCK_DGRAM):
         self.__protocol = None
         self.__kind = kind
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+
     def settimeout(self, value: float):
         self.__timeout = value
 
@@ -87,12 +93,11 @@ def error_received(self, exc):
 
 if __name__ == '__main__':
     async def test_socket_async():
-        socket_async = SocketAsync()
-        socket_async.settimeout(5)
-        await socket_async.connect(('', 27015))
-        socket_async.send(b'\xFF\xFF\xFF\xFFTSource Engine Query\x00\xFF\xFF\xFF\xFF')
-        print(await socket_async.recv())
-        socket_async.close()
+        with SocketAsync() as socket_async:
+            socket_async.settimeout(5)
+            await socket_async.connect(('122.128.109.245', 27015))
+            socket_async.send(b'\xFF\xFF\xFF\xFFTSource Engine Query\x00\xFF\xFF\xFF\xFF')
+            print(await socket_async.recv())
 
     loop = asyncio.get_event_loop()
     loop.run_until_complete(test_socket_async())

opengsq/version.py

@@ -1 +1 @@
-__version__ = '1.1.1'
+__version__ = '1.2.0'