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('