From f90cdf2ffb7777683d69cf8b6066da9bc19b895f Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Tue, 31 May 2022 11:58:38 +0200
Subject: [PATCH] Fix apply_difference should not end get diff for secret if
 not active

---
 telethon/_updates/messagebox.py | 39 +++++++++++++++++++++++----------
 1 file changed, 28 insertions(+), 11 deletions(-)

diff --git a/telethon/_updates/messagebox.py b/telethon/_updates/messagebox.py
index 66e4ac5b..9aeeef2b 100644
--- a/telethon/_updates/messagebox.py
+++ b/telethon/_updates/messagebox.py
@@ -344,7 +344,8 @@ class MessageBox:
         try:
             self.getting_diff_for.remove(entry)
         except KeyError:
-            pass
+            raise RuntimeError('Called end_get_diff on an entry which was not getting diff for')
+
         self.reset_deadline(entry, next_updates_deadline())
         assert entry not in self.possible_gaps, "gaps shouldn't be created while getting difference"
 
@@ -537,25 +538,41 @@ class MessageBox:
         diff,
         chat_hashes,
     ):
+        finish = None
+        result = None
+
         if isinstance(diff, tl.updates.DifferenceEmpty):
+            finish = True
             self.date = diff.date
             self.seq = diff.seq
-            self.end_get_diff(ENTRY_ACCOUNT)
-            self.end_get_diff(ENTRY_SECRET)
-            return [], [], []
+            result = [], [], []
         elif isinstance(diff, tl.updates.Difference):
-            self.end_get_diff(ENTRY_ACCOUNT)
-            self.end_get_diff(ENTRY_SECRET)
+            finish = True
             chat_hashes.extend(diff.users, diff.chats)
-            return self.apply_difference_type(diff, chat_hashes)
+            result = self.apply_difference_type(diff, chat_hashes)
         elif isinstance(diff, tl.updates.DifferenceSlice):
+            finish = False
             chat_hashes.extend(diff.users, diff.chats)
-            return self.apply_difference_type(diff, chat_hashes)
+            result = self.apply_difference_type(diff, chat_hashes)
         elif isinstance(diff, tl.updates.DifferenceTooLong):
+            finish = True
             self.map[ENTRY_ACCOUNT].pts = diff.pts  # the deadline will be reset once the diff ends
-            self.end_get_diff(ENTRY_ACCOUNT)
-            self.end_get_diff(ENTRY_SECRET)
-            return [], [], []
+            result = [], [], []
+
+        if finish:
+            account = ENTRY_ACCOUNT in self.getting_diff_for
+            secret = ENTRY_SECRET in self.getting_diff_for
+
+            if not account and not secret:
+                raise RuntimeWarning('Should not be applying the difference when neither account or secret was diff was active')
+
+            # Both may be active if both expired at the same time.
+            if account:
+                self.end_get_diff(ENTRY_ACCOUNT)
+            if secret:
+                self.end_get_diff(ENTRY_SECRET)
+
+        return result
 
     def apply_difference_type(
         self,