Support connection timeout

telethon/client/telegrambaseclient.py

@@ -202,6 +202,7 @@ class TelegramBaseClient(abc.ABC):
 
         self._request_retries = request_retries or sys.maxsize
         self._connection_retries = connection_retries or sys.maxsize
+        self._timeout = timeout
         self._auto_reconnect = auto_reconnect
 
         assert isinstance(connection, type)
@@ -228,6 +229,7 @@ class TelegramBaseClient(abc.ABC):
             self._loop,
             retries=self._connection_retries,
             auto_reconnect=self._auto_reconnect,
+            connect_timeout=self._timeout,
             update_callback=self._handle_update,
             auth_key_callback=self._auth_key_callback,
             auto_reconnect_callback=self._handle_auto_reconnect

telethon/network/connection/connection.py

@@ -16,7 +16,7 @@ class Connection(abc.ABC):
     under any conditions for as long as the user doesn't disconnect or
     the input parameters to auto-reconnect dictate otherwise.
     """
-    # TODO Support proxy. Support timeout?
+    # TODO Support proxy
     def __init__(self, ip, port, *, loop):
         self._ip = ip
         self._port = port
@@ -31,12 +31,14 @@ class Connection(abc.ABC):
         self._send_queue = asyncio.Queue(1)
         self._recv_queue = asyncio.Queue(1)
 
-    async def connect(self):
+    async def connect(self, timeout=None):
         """
         Establishes a connection with the server.
         """
-        self._reader, self._writer = await asyncio.open_connection(
-            self._ip, self._port, loop=self._loop)
+        self._reader, self._writer = await asyncio.wait_for(
+            asyncio.open_connection(self._ip, self._port, loop=self._loop),
+            loop=self._loop, timeout=timeout
+        )
 
         self._disconnected.clear()
         self._disconnected_future = None

telethon/network/connection/http.py

@@ -4,13 +4,17 @@ from .connection import Connection
 
 
 class ConnectionHttp(Connection):
-    async def connect(self):
+    async def connect(self, timeout=None):
         # TODO Test if the ssl part works or it needs to be as before:
         # dict(ssl_version=ssl.PROTOCOL_SSLv23, ciphers='ADH-AES256-SHA')
-        self._reader, self._writer = await asyncio.open_connection(
-            self._ip, self._port, loop=self._loop, ssl=True)
+        self._reader, self._writer = await asyncio.wait_for(
+            asyncio.open_connection(
+                self._ip, self._port, loop=self._loop, ssl=True),
+            loop=self._loop, timeout=timeout
+        )
 
         self._disconnected.clear()
+        self._disconnected_future = None
         self._send_task = self._loop.create_task(self._send_loop())
         self._recv_task = self._loop.create_task(self._send_loop())
 

telethon/network/connection/tcpabridged.py

@@ -9,8 +9,8 @@ class ConnectionTcpAbridged(Connection):
     only require 1 byte if the packet length is less than
     508 bytes (127 << 2, which is very common).
     """
-    async def connect(self):
-        await super().connect()
+    async def connect(self, timeout=None):
+        await super().connect(timeout=timeout)
         await self.send(b'\xef')
 
     def _write(self, data):

telethon/network/connection/tcpfull.py

@@ -14,8 +14,8 @@ class ConnectionTcpFull(Connection):
         super().__init__(ip, port, loop=loop)
         self._send_counter = 0
 
-    async def connect(self):
-        await super().connect()
+    async def connect(self, timeout=None):
+        await super().connect(timeout=timeout)
         self._send_counter = 0  # Important or Telegram won't reply
 
     def _send(self, data):

telethon/network/connection/tcpintermediate.py

@@ -8,8 +8,8 @@ class ConnectionTcpIntermediate(Connection):
     Intermediate mode between `ConnectionTcpFull` and `ConnectionTcpAbridged`.
     Always sends 4 extra bytes for the packet length.
     """
-    async def connect(self):
-        await super().connect()
+    async def connect(self, timeout=None):
+        await super().connect(timeout=timeout)
         await self.send(b'\xee\xee\xee\xee')
 
     def _send(self, data):

telethon/network/connection/tcpobfuscated.py

@@ -22,8 +22,8 @@ class ConnectionTcpObfuscated(ConnectionTcpAbridged):
     async def _read(self, n):
         return self._aes_decrypt.encrypt(await self._reader.readexactly(n))
 
-    async def connect(self):
-        await Connection.connect(self)
+    async def connect(self, timeout=None):
+        await Connection.connect(self, timeout=timeout)
 
         # Obfuscated messages secrets cannot start with any of these
         keywords = (b'PVrG', b'GET ', b'POST', b'\xee\xee\xee\xee')

telethon/network/mtprotolayer.py

@@ -22,11 +22,11 @@ class MTProtoLayer:
         self._connection = connection
         self._state = MTProtoState(auth_key)
 
-    def connect(self):
+    def connect(self, timeout=None):
         """
         Wrapper for ``self._connection.connect()``.
         """
-        return self._connection.connect()
+        return self._connection.connect(timeout=timeout)
 
     def disconnect(self):
         """

telethon/network/mtprotosender.py

@@ -40,12 +40,14 @@ class MTProtoSender:
     key exists yet.
     """
     def __init__(self, loop, *,
-                 retries=5, auto_reconnect=True, update_callback=None,
+                 retries=5, auto_reconnect=True, connect_timeout=None,
+                 update_callback=None,
                  auth_key_callback=None, auto_reconnect_callback=None):
         self._connection = None  # MTProtoLayer, a.k.a. encrypted connection
         self._loop = loop
         self._retries = retries
         self._auto_reconnect = auto_reconnect
+        self._connect_timeout = connect_timeout
         self._update_callback = update_callback
         self._auth_key_callback = auth_key_callback
         self._auto_reconnect_callback = auto_reconnect_callback
@@ -189,14 +191,12 @@ class MTProtoSender:
         authorization key if necessary, and starting the send and
         receive loops.
         """
-        # TODO With ``asyncio.open_connection``, no timeout occurs
-        # However, these are probably desirable in some circumstances.
         __log__.info('Connecting to %s...', self._connection)
         for retry in range(1, self._retries + 1):
             try:
                 __log__.debug('Connection attempt {}...'.format(retry))
-                await self._connection.connect()
-            except OSError as e:
+                await self._connection.connect(timeout=self._connect_timeout)
+            except (OSError, asyncio.TimeoutError) as e:
                 __log__.warning('Attempt {} at connecting failed: {}: {}'
                                 .format(retry, type(e).__name__, e))
             else: