diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py index c7210ad6..3e27638b 100644 --- a/pyrogram/connection/connection.py +++ b/pyrogram/connection/connection.py @@ -25,7 +25,6 @@ log = logging.getLogger(__name__) class Connection: - TIMEOUT = 10 MAX_RETRIES = 3 MODES = { @@ -50,8 +49,8 @@ class Connection: try: log.info("Connecting...") - await asyncio.wait_for(self.protocol.connect(self.address), Connection.TIMEOUT) - except (OSError, asyncio.TimeoutError): + await self.protocol.connect(self.address) + except OSError: self.protocol.close() await asyncio.sleep(1) else: diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py index f006cadd..f541153e 100644 --- a/pyrogram/connection/transport/tcp/tcp.py +++ b/pyrogram/connection/transport/tcp/tcp.py @@ -34,6 +34,8 @@ log = logging.getLogger(__name__) class TCP: + TIMEOUT = 10 + def __init__(self, proxy: dict): self.proxy = proxy @@ -41,6 +43,7 @@ class TCP: self.reader = None # type: asyncio.StreamReader self.writer = None # type: asyncio.StreamWriter + self.socket.settimeout(TCP.TIMEOUT) self.proxy_enabled = proxy.get("enabled", False) if proxy and self.proxy_enabled: @@ -81,8 +84,11 @@ class TCP: while len(data) < length: try: - chunk = await self.reader.read(length - len(data)) - except OSError: + chunk = await asyncio.wait_for( + self.reader.read(length - len(data)), + TCP.TIMEOUT + ) + except (OSError, asyncio.TimeoutError): return None else: if chunk: diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index d58eae37..6479dfdd 100644 --- a/pyrogram/session/session.py +++ b/pyrogram/session/session.py @@ -332,10 +332,7 @@ class Session: log.info("RecvTask started") while True: - try: - packet = await asyncio.wait_for(self.connection.recv(), self.connection.TIMEOUT) - except asyncio.TimeoutError: - packet = None + packet = await self.connection.recv() if packet is None or len(packet) == 4: self.recv_queue.put_nowait(None)