Support connection timeout

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')