Delegate timeout to TCP

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)
+                await self.protocol.connect(self.address)
-            except (OSError, asyncio.TimeoutError):
+            except OSError:
                 self.protocol.close()
                 await asyncio.sleep(1)
             else:

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))
+                chunk = await asyncio.wait_for(
-            except OSError:
+                    self.reader.read(length - len(data)),
+                    TCP.TIMEOUT
+                )
+            except (OSError, asyncio.TimeoutError):
                 return None
             else:
                 if chunk:

pyrogram/session/session.py

@@ -332,10 +332,7 @@ class Session:
         log.info("RecvTask started")
 
         while True:
-            try:
+            packet = await self.connection.recv()
-                packet = await asyncio.wait_for(self.connection.recv(), self.connection.TIMEOUT)
-            except asyncio.TimeoutError:
-                packet = None
 
             if packet is None or len(packet) == 4:
                 self.recv_queue.put_nowait(None)