From 6de7329ce7d3288b8fa20c6b3eeeb4c0de77223d Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 21 Feb 2019 10:41:33 +0100
Subject: [PATCH] Fix Connection abstraction leak

---
 telethon/network/connection/connection.py      | 15 +++++++++++++++
 telethon/network/connection/tcpabridged.py     |  4 +---
 telethon/network/connection/tcpfull.py         |  3 +--
 telethon/network/connection/tcpintermediate.py |  4 +---
 telethon/network/connection/tcpobfuscated.py   |  6 +-----
 5 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/telethon/network/connection/connection.py b/telethon/network/connection/connection.py
index 029c19fe..747fa8de 100644
--- a/telethon/network/connection/connection.py
+++ b/telethon/network/connection/connection.py
@@ -76,6 +76,9 @@ class Connection(abc.ABC):
                 await asyncio.open_connection(sock=s, loop=self._loop)
 
         self._connected = True
+        self._init_conn()
+        await self._writer.drain()
+
         self._send_task = self._loop.create_task(self._send_loop())
         self._recv_task = self._loop.create_task(self._recv_loop())
 
@@ -170,6 +173,18 @@ class Connection(abc.ABC):
             except asyncio.CancelledError:
                 break
 
+    @abc.abstractmethod
+    def _init_conn(self):
+        """
+        This method will be called after `connect` is called.
+        After this method finishes, the writer will be drained.
+
+        Subclasses should make use of this if they need to send
+        data to Telegram to indicate which connection mode will
+        be used.
+        """
+        raise NotImplemented
+
     @abc.abstractmethod
     def _send(self, data):
         """
diff --git a/telethon/network/connection/tcpabridged.py b/telethon/network/connection/tcpabridged.py
index c9350da9..70cd610d 100644
--- a/telethon/network/connection/tcpabridged.py
+++ b/telethon/network/connection/tcpabridged.py
@@ -9,10 +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, timeout=None, ssl=None):
-        await super().connect(timeout=timeout, ssl=ssl)
+    def _init_conn(self):
         self._writer.write(b'\xef')
-        await self._writer.drain()
 
     def _write(self, data):
         """
diff --git a/telethon/network/connection/tcpfull.py b/telethon/network/connection/tcpfull.py
index 9eb0d934..05202edc 100644
--- a/telethon/network/connection/tcpfull.py
+++ b/telethon/network/connection/tcpfull.py
@@ -15,8 +15,7 @@ class ConnectionTcpFull(Connection):
             ip, port, dc_id, loop=loop, loggers=loggers, proxy=proxy)
         self._send_counter = 0
 
-    async def connect(self, timeout=None, ssl=None):
-        await super().connect(timeout=timeout, ssl=ssl)
+    def _init_conn(self):
         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 59467be7..273973c2 100644
--- a/telethon/network/connection/tcpintermediate.py
+++ b/telethon/network/connection/tcpintermediate.py
@@ -8,10 +8,8 @@ class ConnectionTcpIntermediate(Connection):
     Intermediate mode between `ConnectionTcpFull` and `ConnectionTcpAbridged`.
     Always sends 4 extra bytes for the packet length.
     """
-    async def connect(self, timeout=None, ssl=None):
-        await super().connect(timeout=timeout, ssl=ssl)
+    def _init_conn(self):
         self._writer.write(b'\xee\xee\xee\xee')
-        await self._writer.drain()
 
     def _send(self, data):
         self._writer.write(struct.pack('<i', len(data)) + data)
diff --git a/telethon/network/connection/tcpobfuscated.py b/telethon/network/connection/tcpobfuscated.py
index 57f9f878..3a0d3498 100644
--- a/telethon/network/connection/tcpobfuscated.py
+++ b/telethon/network/connection/tcpobfuscated.py
@@ -24,10 +24,7 @@ class ConnectionTcpObfuscated(ConnectionTcpAbridged):
     async def _read(self, n):
         return self._aes_decrypt.encrypt(await self._reader.readexactly(n))
 
-    async def connect(self, timeout=None, ssl=None):
-        # FIXME: that's an abstraction leak
-        await Connection.connect(self, timeout=timeout, ssl=ssl)
-
+    def _init_conn(self):
         # Obfuscated messages secrets cannot start with any of these
         keywords = (b'PVrG', b'GET ', b'POST', b'\xee\xee\xee\xee')
         while True:
@@ -53,7 +50,6 @@ class ConnectionTcpObfuscated(ConnectionTcpAbridged):
         random[56:64] = self._compose_tail(bytes(random))
 
         self._writer.write(random)
-        await self._writer.drain()
 
     # Next functions provide the variable parts of the connection handshake.
     # This is necessary to modify obfuscated2 the way that MTProxy requires.