From 52354b93d06e6085a79ccfe5ac9f57fded066c50 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sun, 17 Jun 2018 18:44:45 +0200
Subject: [PATCH] Add timeout when connecting

---
 pyrogram/connection/connection.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py
index a9aba7c1..4a25b72e 100644
--- a/pyrogram/connection/connection.py
+++ b/pyrogram/connection/connection.py
@@ -25,6 +25,7 @@ log = logging.getLogger(__name__)
 
 
 class Connection:
+    TIMEOUT = 10
     MAX_RETRIES = 3
 
     MODES = {
@@ -49,8 +50,8 @@ class Connection:
 
             try:
                 log.info("Connecting...")
-                await self.protocol.connect(self.address)
-            except OSError:
+                await asyncio.wait_for(self.protocol.connect(self.address), Connection.TIMEOUT)
+            except (OSError, asyncio.TimeoutError):
                 self.protocol.close()
                 await asyncio.sleep(1)
             else: