From 1f40372235e0ffbddff7832e48e46252b2efb1d6 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Sun, 23 Jan 2022 12:43:41 +0100
Subject: [PATCH] Fix update handling for channels

---
 telethon/_client/updates.py     |  5 ++++-
 telethon/_updates/messagebox.py | 20 ++++++++++----------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/telethon/_client/updates.py b/telethon/_client/updates.py
index b0d632b5..e3e8fd78 100644
--- a/telethon/_client/updates.py
+++ b/telethon/_client/updates.py
@@ -96,6 +96,7 @@ async def _update_loop(self: 'TelegramClient'):
                 self._log[__name__].info('Getting difference for account updates')
                 diff = await self(get_diff)
                 updates, users, chats = self._message_box.apply_difference(diff, self._entity_cache)
+                self._entity_cache.extend(users, chats)
                 updates_to_dispatch.extend(updates)
                 continue
 
@@ -103,7 +104,8 @@ async def _update_loop(self: 'TelegramClient'):
             if get_diff:
                 self._log[__name__].info('Getting difference for channel updates')
                 diff = await self(get_diff)
-                updates, users, chats = self._message_box.apply_channel_difference(diff, self._entity_cache)
+                updates, users, chats = self._message_box.apply_channel_difference(get_diff, diff, self._entity_cache)
+                self._entity_cache.extend(users, chats)
                 updates_to_dispatch.extend(updates)
                 continue
 
@@ -119,6 +121,7 @@ async def _update_loop(self: 'TelegramClient'):
 
             processed = []
             users, chats = self._message_box.process_updates(updates, self._entity_cache, processed)
+            self._entity_cache.extend(users, chats)
             updates_to_dispatch.extend(processed)
     except Exception:
         self._log[__name__].exception('Fatal error handling updates (this is a bug in Telethon, please report it)')
diff --git a/telethon/_updates/messagebox.py b/telethon/_updates/messagebox.py
index 542fc629..235247f3 100644
--- a/telethon/_updates/messagebox.py
+++ b/telethon/_updates/messagebox.py
@@ -232,7 +232,7 @@ class MessageBox:
 
     # Convenience to reset a channel's deadline, with optional timeout.
     def reset_channel_deadline(self, channel_id, timeout):
-        self.reset_deadlines(channel_id, asyncio.get_running_loop().time() + (timeout or NO_UPDATES_TIMEOUT))
+        self.reset_deadline(channel_id, asyncio.get_running_loop().time() + (timeout or NO_UPDATES_TIMEOUT))
 
     # Reset all the deadlines in `reset_deadlines_for` and then empty the set.
     def apply_deadlines_reset(self):
@@ -524,10 +524,10 @@ class MessageBox:
 
         return _tl.fn.updates.GetChannelDifference(
             force=False,
-            channel=channel,
+            channel=packed.try_to_input_channel(),
             filter=_tl.ChannelMessagesFilterEmpty(),
             pts=state.pts,
-            limit=BOT_CHANNEL_DIFF_LIMIT if chat_hashes.is_self_bot() else USER_CHANNEL_DIFF_LIMIT
+            limit=BOT_CHANNEL_DIFF_LIMIT if chat_hashes.self_bot else USER_CHANNEL_DIFF_LIMIT
         )
 
     # Similar to [`MessageBox::process_updates`], but using the result from getting difference.
@@ -538,7 +538,7 @@ class MessageBox:
         chat_hashes,
     ):
         entry = request.channel.channel_id
-        self.possible_gaps.remove(entry)
+        self.possible_gaps.pop(entry, None)
 
         if isinstance(diff, _tl.updates.ChannelDifferenceEmpty):
             assert diff.final
@@ -549,7 +549,7 @@ class MessageBox:
             assert diff.final
             self.map[entry].pts = diff.dialog.pts
             chat_hashes.extend(diff.users, diff.chats)
-            self.reset_channel_deadline(channel_id, diff.timeout)
+            self.reset_channel_deadline(entry, diff.timeout)
             # This `diff` has the "latest messages and corresponding chats", but it would
             # be strange to give the user only partial changes of these when they would
             # expect all updates to be fetched. Instead, nothing is returned.
@@ -558,15 +558,15 @@ class MessageBox:
             if diff.final:
                 self.end_get_diff(entry)
 
-            self.map[entry].pts = pts
-            updates.extend(_tl.UpdateNewMessage(
+            self.map[entry].pts = diff.pts
+            diff.other_updates.extend(_tl.UpdateNewMessage(
                 message=m,
                 pts=NO_SEQ,
                 pts_count=NO_SEQ,
             ) for m in diff.new_messages)
-            chat_hashes.extend(diff.users, diff.chats);
-            self.reset_channel_deadline(channel_id, timeout)
+            chat_hashes.extend(diff.users, diff.chats)
+            self.reset_channel_deadline(entry, None)
 
-            (diff.updates, diff.users, diff.chats)
+            return diff.other_updates, diff.users, diff.chats
 
     # endregion Getting and applying channel difference.