diff --git a/telethon/network/connection/connection.py b/telethon/network/connection/connection.py index 5263c5ff..2987bc11 100644 --- a/telethon/network/connection/connection.py +++ b/telethon/network/connection/connection.py @@ -39,10 +39,7 @@ class Connection(abc.ABC): self._send_queue = asyncio.Queue(1) self._recv_queue = asyncio.Queue(1) - async def connect(self, timeout=None, ssl=None): - """ - Establishes a connection with the server. - """ + async def _connect(self, timeout=None, ssl=None): if not self._proxy: self._reader, self._writer = await asyncio.wait_for( asyncio.open_connection( @@ -81,11 +78,17 @@ class Connection(abc.ABC): self._reader, self._writer = \ await asyncio.open_connection(sock=s, loop=self._loop) - self._connected = True self._codec = self.packet_codec(self) self._init_conn() await self._writer.drain() + async def connect(self, timeout=None, ssl=None): + """ + Establishes a connection with the server. + """ + await self._connect(timeout=timeout, ssl=ssl) + self._connected = True + self._send_task = self._loop.create_task(self._send_loop()) self._recv_task = self._loop.create_task(self._recv_loop()) diff --git a/telethon/network/connection/tcpmtproxy.py b/telethon/network/connection/tcpmtproxy.py index fc8947db..ac3d2ba5 100644 --- a/telethon/network/connection/tcpmtproxy.py +++ b/telethon/network/connection/tcpmtproxy.py @@ -1,3 +1,4 @@ +import asyncio import hashlib import os @@ -92,6 +93,28 @@ class TcpMTProxy(ObfuscatedConnection): packet_codec = None obfuscated_io = MTProxyIO + async def _connect(self, timeout=None, ssl=None): + await super()._connect(timeout=timeout, ssl=ssl) + + # Wait for EOF for 2 seconds (or if _wait_for_data's definition + # is missing or different, just sleep for 2 seconds). This way + # we give the proxy a chance to close the connection if the current + # codec (which the proxy detects with the data we sent) cannot + # be used for this proxy. This is a work around for #1134. + # TODO Sleeping for N seconds may not be the best solution + # TODO This fix could be welcome for HTTP proxies as well + try: + await asyncio.wait_for(self._reader._wait_for_data('proxy'), 2) + except asyncio.TimeoutError: + pass + except Exception: + await asyncio.sleep(2) + + if self._reader.at_eof(): + self.disconnect() + raise ConnectionError( + 'Proxy closed the connection after sending initial payload') + @staticmethod def address_info(proxy_info): if proxy_info is None: