From db83709c6b48c1d25ac4da1b3dbf3a2839386d8c Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Thu, 4 Oct 2018 16:39:57 +0200 Subject: [PATCH] Support connection timeout --- telethon/client/telegrambaseclient.py | 2 ++ telethon/network/connection/connection.py | 10 ++++++---- telethon/network/connection/http.py | 10 +++++++--- telethon/network/connection/tcpabridged.py | 4 ++-- telethon/network/connection/tcpfull.py | 4 ++-- telethon/network/connection/tcpintermediate.py | 4 ++-- telethon/network/connection/tcpobfuscated.py | 4 ++-- telethon/network/mtprotolayer.py | 4 ++-- telethon/network/mtprotosender.py | 10 +++++----- 9 files changed, 30 insertions(+), 22 deletions(-) diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py index 8ef6103a..b0face7a 100644 --- a/telethon/client/telegrambaseclient.py +++ b/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 diff --git a/telethon/network/connection/connection.py b/telethon/network/connection/connection.py index 5d650a38..8e3acb44 100644 --- a/telethon/network/connection/connection.py +++ b/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 diff --git a/telethon/network/connection/http.py b/telethon/network/connection/http.py index c346d2f8..09096454 100644 --- a/telethon/network/connection/http.py +++ b/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()) diff --git a/telethon/network/connection/tcpabridged.py b/telethon/network/connection/tcpabridged.py index 6352413e..dced879f 100644 --- a/telethon/network/connection/tcpabridged.py +++ b/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): diff --git a/telethon/network/connection/tcpfull.py b/telethon/network/connection/tcpfull.py index e07e3558..1d278c30 100644 --- a/telethon/network/connection/tcpfull.py +++ b/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): diff --git a/telethon/network/connection/tcpintermediate.py b/telethon/network/connection/tcpintermediate.py index 82a99250..7b8e3c79 100644 --- a/telethon/network/connection/tcpintermediate.py +++ b/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): diff --git a/telethon/network/connection/tcpobfuscated.py b/telethon/network/connection/tcpobfuscated.py index 393977b2..eb4fd7ff 100644 --- a/telethon/network/connection/tcpobfuscated.py +++ b/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') diff --git a/telethon/network/mtprotolayer.py b/telethon/network/mtprotolayer.py index d3cd7ae4..731cabb8 100644 --- a/telethon/network/mtprotolayer.py +++ b/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): """ diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py index d398eccc..4d96cb2c 100644 --- a/telethon/network/mtprotosender.py +++ b/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: