From 2e4476a7546ce5909761ffc37049747ed270709b Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 21 Mar 2019 11:22:09 +0100
Subject: [PATCH] Workaround #1134 by early checking if proxy closes connection

---
 telethon/network/connection/connection.py | 13 ++++++++-----
 telethon/network/connection/tcpmtproxy.py | 23 +++++++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

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: