From 21ffa2f26b1d0ffe9533431aeee93b7a06b03320 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Mon, 1 Oct 2018 14:02:23 +0200
Subject: [PATCH] Fix DC migration and seqno

---
 telethon/client/telegrambaseclient.py |  2 +-
 telethon/network/mtprotolayer.py      | 12 ++++++++----
 telethon/network/mtprotostate.py      |  5 +++--
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py
index 5b82b625..89d59495 100644
--- a/telethon/client/telegrambaseclient.py
+++ b/telethon/client/telegrambaseclient.py
@@ -369,7 +369,7 @@ class TelegramBaseClient(abc.ABC):
         self.session.set_dc(dc.id, dc.ip_address, dc.port)
         # auth_key's are associated with a server, which has now changed
         # so it's not valid anymore. Set to None to force recreating it.
-        self.session.auth_key = self._sender.state.auth_key = None
+        self.session.auth_key = self._sender._connection._state.auth_key = None
         self.session.save()
         await self._disconnect()
         return await self.connect()
diff --git a/telethon/network/mtprotolayer.py b/telethon/network/mtprotolayer.py
index 8142a5e9..2960b28e 100644
--- a/telethon/network/mtprotolayer.py
+++ b/telethon/network/mtprotolayer.py
@@ -2,6 +2,7 @@ import io
 import struct
 
 from .mtprotostate import MTProtoState
+from ..tl import TLRequest
 from ..tl.core.messagecontainer import MessageContainer
 
 
@@ -78,14 +79,15 @@ class MTProtoLayer:
         for state in state_list:
             if not isinstance(state, list):
                 n += 1
-                state.msg_id = \
-                    self._state.write_data_as_message(buffer, state.data)
+                state.msg_id = self._state.write_data_as_message(
+                    buffer, state.data, isinstance(state.request, TLRequest))
             else:
                 last_id = None
                 for s in state:
                     n += 1
                     last_id = s.msg_id = self._state.write_data_as_message(
-                        buffer, s.data, after_id=last_id)
+                        buffer, s.data, isinstance(s.request, TLRequest),
+                        after_id=last_id)
 
         if n > 1:
             # Inlined code to pack several messages into a container
@@ -97,7 +99,9 @@ class MTProtoLayer:
                 '<Ii', MessageContainer.CONSTRUCTOR_ID, n
             ) + buffer.getvalue()
             buffer = io.BytesIO()
-            container_id = self._state.write_data_as_message(buffer, data)
+            container_id = self._state.write_data_as_message(
+                buffer, data, content_related=False
+            )
             for state in state_list:
                 if not isinstance(state, list):
                     state.container_id = container_id
diff --git a/telethon/network/mtprotostate.py b/telethon/network/mtprotostate.py
index 761b869f..e38dad60 100644
--- a/telethon/network/mtprotostate.py
+++ b/telethon/network/mtprotostate.py
@@ -69,14 +69,15 @@ class MTProtoState:
 
         return aes_key, aes_iv
 
-    def write_data_as_message(self, buffer, data, after_id=None):
+    def write_data_as_message(self, buffer, data, content_related,
+                              *, after_id=None):
         """
         Writes a message containing the given data into buffer.
 
         Returns the message id.
         """
         msg_id = self._get_new_msg_id()
-        seq_no = self._get_seq_no(True)  # TODO ack/ping are not content-related
+        seq_no = self._get_seq_no(content_related)
         if after_id is None:
             body = GzipPacked.gzip_if_smaller(data)
         else: