[!] Fix closing connection

This commit is contained in:
Maxim Khomutov 2024-07-27 12:21:29 +03:00
parent a406956080
commit b0303f3e6d
3 changed files with 22 additions and 6 deletions

View File

@ -754,11 +754,13 @@ class Client:
try: try:
if not self.__writer.is_closing(): if not self.__writer.is_closing():
self.__writer.close() self.__writer.close()
await self.__writer.wait_closed()
except Exception as e: except Exception as e:
self.log.debug(f"Error while closing writer: {e}") self.log.debug(f"Error while closing writer: {e}")
try: try:
_, down_w = self._down_sock _, down_w = self._down_sock
if down_w and not down_w.is_closing(): if down_w and not down_w.is_closing():
down_w.close() down_w.close()
await down_w.wait_closed()
except Exception as e: except Exception as e:
self.log.debug(f"Error while closing download writer: {e}") self.log.debug(f"Error while closing download writer: {e}")

View File

@ -22,6 +22,8 @@ class TCPServer:
self.host = host self.host = host
self.port = port self.port = port
self.run = False self.run = False
self._connections = set()
self.server = None
self.rl = RateLimiter(50, 10, 300) self.rl = RateLimiter(50, 10, 300)
console.add_command("rl", self.rl.parse_console, None, "RateLimiter menu", console.add_command("rl", self.rl.parse_console, None, "RateLimiter menu",
{"rl": {"info": None, "unban": None, "ban": None, "help": None}}) {"rl": {"info": None, "unban": None, "ban": None, "help": None}})
@ -135,7 +137,8 @@ class TCPServer:
return False, None return False, None
async def handle_client(self, reader, writer): async def handle_client(self, reader, writer):
while True: while self.run:
self._connections.add(writer)
try: try:
ip = writer.get_extra_info('peername')[0] ip = writer.get_extra_info('peername')[0]
if self.rl.is_banned(ip): if self.rl.is_banned(ip):
@ -163,12 +166,12 @@ class TCPServer:
self.log.debug("Starting TCP server.") self.log.debug("Starting TCP server.")
self.run = True self.run = True
try: try:
server = await asyncio.start_server(self.handle_client, self.host, self.port, self.server = await asyncio.start_server(self.handle_client, self.host, self.port,
backlog=int(config.Game["players"] * 2.3)) backlog=int(config.Game["players"] * 2.3))
self.log.debug(f"TCP server started on {server.sockets[0].getsockname()!r}") self.log.debug(f"TCP server started on {self.server.sockets[0].getsockname()!r}")
while True: while True:
async with server: async with self.server:
await server.serve_forever() await self.server.serve_forever()
except OSError as e: except OSError as e:
self.log.error(i18n.core_bind_failed.format(e)) self.log.error(i18n.core_bind_failed.format(e))
raise e raise e
@ -181,5 +184,14 @@ class TCPServer:
self.run = False self.run = False
self.Core.run = False self.Core.run = False
def stop(self): async def stop(self):
self.log.debug("Stopping TCP server") self.log.debug("Stopping TCP server")
try:
self.server.close()
for conn in self._connections:
conn.close()
await conn.wait_closed()
await self.server.wait_closed()
except Exception as e:
self.log.exception(e)
self.log.debug("Stopped")

View File

@ -15,11 +15,13 @@ from modules import RateLimiter
class TCPServer: class TCPServer:
def __init__(self, core: Core, host, port): def __init__(self, core: Core, host, port):
self.server = await asyncio.start_server(self.handle_client, "", 0, backlog=int(config.Game["players"] * 2.3))
self.log = utils.get_logger("TCPServer") self.log = utils.get_logger("TCPServer")
self.loop = asyncio.get_event_loop() self.loop = asyncio.get_event_loop()
self.Core = core self.Core = core
self.host = host self.host = host
self.port = port self.port = port
self._connections = set()
self.run = False self.run = False
self.rl = RateLimiter(50, 10, 15) self.rl = RateLimiter(50, 10, 15)