From 1bc599e26ca7d905c6b6d1fd57881f3a991da02a Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sun, 17 Jun 2018 19:20:22 +0200
Subject: [PATCH] Delegate timeout to TCP

---
 pyrogram/connection/connection.py        |  5 ++---
 pyrogram/connection/transport/tcp/tcp.py | 10 ++++++++--
 pyrogram/session/session.py              |  5 +----
 3 files changed, 11 insertions(+), 9 deletions(-)

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)