move

main.py

@@ -43,31 +43,32 @@ def parse_utf16_hostname(data: bytes):
         hostname += chr(wchar)
     return hostname
 
-async def send_chat(writer, msg: str, player_type: int):
-    # Truncate and encode msg to UTF-16-LE, with null terminator
-    encoded = msg.encode('utf-16le')[:510]  # max 255 UTF-16 chars
-    encoded += struct.pack('<H', 0)  # null terminator
 
-    payload = struct.pack('<H', player_type) + encoded
-    packet = struct.pack('<H', len(payload) + 1) + bytes([0x19]) + payload
-
-    writer.write(packet)
-    await writer.drain()
+async def handle_client(reader, writer):
+    async def send_chat(msg: str, player_type: int):
+        # Truncate and encode msg to UTF-16-LE, with null terminator
+        encoded = msg.encode('utf-16le')[:510]  # max 255 UTF-16 chars
+        encoded += struct.pack('<H', 0)  # null terminator
 
+        payload = struct.pack('<H', player_type) + encoded
+        packet = struct.pack('<H', len(payload) + 1) + bytes([0x19]) + payload
 
-async def close_connection(writer):
-    try:
-        payload = struct.pack('<BBBBI', 1, 0, 0, 0, 9)  # msg=1, code=9
-        packet = struct.pack('<H', len(payload) + 1) + bytes([0x02]) + payload  # STOC_ERROR_MSG = 0x02
         writer.write(packet)
         await writer.drain()
-    except Exception as e:
-        logger.warning(f"Failed to send error message before closing: {e}")
-    writer.close()
-    await writer.wait_closed()
 
 
-async def handle_client(reader, writer):
+    async def close_connection():
+        try:
+            payload = struct.pack('<BBBBI', 1, 0, 0, 0, 9)  # msg=1, code=9
+            packet = struct.pack('<H', len(payload) + 1) + bytes([0x02]) + payload  # STOC_ERROR_MSG = 0x02
+            writer.write(packet)
+            await writer.drain()
+        except Exception as e:
+            logger.warning(f"Failed to send error message before closing: {e}")
+        writer.close()
+        await writer.wait_closed()
+
+
     peer_ip = writer.get_extra_info("peername")[0]
     is_proxy = is_trusted(peer_ip)
 
@@ -82,18 +83,18 @@ async def handle_client(reader, writer):
 
             if packet_id != CTOS_EXTERNAL_ADDRESS:
                 logger.warning(f"First packet is not CTOS_EXTERNAL_ADDRESS from {peer_ip}, closing.")
-                await send_chat(writer, "400 Bad Request: CTOS_EXTERNAL_ADDRESS not found", player_type=11)
-                await close_connection(writer)
+                await send_chat( "400 Bad Request: CTOS_EXTERNAL_ADDRESS not found", player_type=11)
+                await close_connection()
                 return
 
             if length < 6 or length > 516:
                 logger.warning(f"Invalid packet length {length} from {peer_ip}, closing.")
-                await close_connection(writer)
+                await close_connection()
                 return
             payload = await asyncio.wait_for(reader.readexactly(length - 1), timeout=5.0)
         except asyncio.TimeoutError:
             logger.warning(f"Timeout while waiting for payload from {peer_ip}, closing.")
-            await close_connection(writer)
+            await close_connection()
             return
 
         real_ip = ipaddress.IPv4Address(payload[0:4])
@@ -112,8 +113,8 @@ async def handle_client(reader, writer):
 
         if not target_host:
             logger.warning(f"No route found for hostname: {hostname} from {client_ip}")
-            await send_chat(writer, f"404 Not Found: Host [{hostname}] not found", player_type=11)
-            await close_connection(writer)
+            await send_chat(f"404 Not Found: Host [{hostname}] not found", player_type=11)
+            await close_connection()
             return
 
         logger.info(f"{client_ip} requested {hostname} → forwarding to {target_host}:{target_port}")
@@ -124,8 +125,8 @@ async def handle_client(reader, writer):
                 asyncio.open_connection(target_host, target_port), timeout=5.0)
         except Exception as e:
             logger.warning(f"Failed to connect to {target_host}:{target_port} for client {client_ip}: {e}")
-            await send_chat(writer, f"502 Bad Gateway: Host [{hostname}] cannot be connected", player_type=11)
-            await close_connection(writer)
+            await send_chat(f"502 Bad Gateway: Host [{hostname}] cannot be connected", player_type=11)
+            await close_connection()
             return
 
         # Overwrite real_ip in payload with resolved client_ip