move

main.py

@@ -43,7 +43,9 @@ def parse_utf16_hostname(data: bytes):
         hostname += chr(wchar)
     return hostname
 
-async def send_chat(writer, msg: str, player_type: int):
+
+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
@@ -55,7 +57,7 @@ async def send_chat(writer, msg: str, player_type: int):
         await writer.drain()
 
 
-async def close_connection(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
@@ -67,7 +69,6 @@ async def close_connection(writer):
         await writer.wait_closed()
 
 
-async def handle_client(reader, writer):
     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