Fix Connection abstraction leak

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):
         """

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):
+    def _init_conn(self):
-        await super().connect(timeout=timeout, ssl=ssl)
         self._writer.write(b'\xef')
-        await self._writer.drain()
 
     def _write(self, data):
         """

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):
+    def _init_conn(self):
-        await super().connect(timeout=timeout, ssl=ssl)
         self._send_counter = 0  # Important or Telegram won't reply
 
     def _send(self, data):

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):
+    def _init_conn(self):
-        await super().connect(timeout=timeout, ssl=ssl)
         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)

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):
+    def _init_conn(self):
-        # FIXME: that's an abstraction leak
-        await Connection.connect(self, timeout=timeout, ssl=ssl)
-
         # 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.