From efec8787d89f1bb14b8213f6cb5228be9ca84c29 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Mon, 11 Jan 2021 23:31:25 +0300
Subject: [PATCH] Add MessagesManager::check_pts_update.

---
 td/telegram/MessagesManager.cpp | 89 ++++++++++++++-------------------
 td/telegram/MessagesManager.h   |  4 +-
 2 files changed, 40 insertions(+), 53 deletions(-)

diff --git a/td/telegram/MessagesManager.cpp b/td/telegram/MessagesManager.cpp
index 70bdbe720..19e94d9b4 100644
--- a/td/telegram/MessagesManager.cpp
+++ b/td/telegram/MessagesManager.cpp
@@ -6087,7 +6087,7 @@ bool MessagesManager::is_allowed_useless_update(const tl_object_ptr<telegram_api
   return false;
 }
 
-bool MessagesManager::check_update_dialog_id(const tl_object_ptr<telegram_api::Update> &update, DialogId dialog_id) {
+bool MessagesManager::check_pts_update_dialog_id(DialogId dialog_id) {
   switch (dialog_id.get_type()) {
     case DialogType::User:
     case DialogType::Chat:
@@ -6095,7 +6095,6 @@ bool MessagesManager::check_update_dialog_id(const tl_object_ptr<telegram_api::U
     case DialogType::Channel:
     case DialogType::SecretChat:
     case DialogType::None:
-      LOG(ERROR) << "Receive update in wrong " << dialog_id << ": " << oneline(to_string(update));
       return false;
     default:
       UNREACHABLE();
@@ -6160,6 +6159,39 @@ int32 MessagesManager::get_min_pending_pts() const {
   return result;
 }
 
+bool MessagesManager::check_pts_update(const tl_object_ptr<telegram_api::Update> &update) {
+  CHECK(update != nullptr);
+  switch (update->get_id()) {
+    case dummyUpdate::ID:
+    case updateSentMessage::ID:
+    case telegram_api::updateReadMessagesContents::ID:
+    case telegram_api::updateDeleteMessages::ID:
+      return true;
+    case telegram_api::updateNewMessage::ID: {
+      auto update_new_message = static_cast<const telegram_api::updateNewMessage *>(update.get());
+      return check_pts_update_dialog_id(get_message_dialog_id(update_new_message->message_));
+    }
+    case telegram_api::updateReadHistoryInbox::ID: {
+      auto update_read_history_inbox = static_cast<const telegram_api::updateReadHistoryInbox *>(update.get());
+      return check_pts_update_dialog_id(DialogId(update_read_history_inbox->peer_));
+    }
+    case telegram_api::updateReadHistoryOutbox::ID: {
+      auto update_read_history_outbox = static_cast<const telegram_api::updateReadHistoryOutbox *>(update.get());
+      return check_pts_update_dialog_id(DialogId(update_read_history_outbox->peer_));
+    }
+    case telegram_api::updateEditMessage::ID: {
+      auto update_edit_message = static_cast<const telegram_api::updateEditMessage *>(update.get());
+      return check_pts_update_dialog_id(get_message_dialog_id(update_edit_message->message_));
+    }
+    case telegram_api::updatePinnedMessages::ID: {
+      auto update_pinned_messages = static_cast<const telegram_api::updatePinnedMessages *>(update.get());
+      return check_pts_update_dialog_id(DialogId(update_pinned_messages->peer_));
+    }
+    default:
+      return false;
+  }
+}
+
 void MessagesManager::add_pending_update(tl_object_ptr<telegram_api::Update> &&update, int32 new_pts, int32 pts_count,
                                          bool force_apply, Promise<Unit> &&promise, const char *source) {
   // do not try to run getDifference from this function
@@ -6176,56 +6208,9 @@ void MessagesManager::add_pending_update(tl_object_ptr<telegram_api::Update> &&u
   // TODO need to save all updates that can change result of running queries not associated with pts (for example
   // getHistory) and apply them to result of this queries
 
-  switch (update->get_id()) {
-    case dummyUpdate::ID:
-    case updateSentMessage::ID:
-    case telegram_api::updateReadMessagesContents::ID:
-    case telegram_api::updateDeleteMessages::ID:
-      // nothing to check
-      break;
-    case telegram_api::updateNewMessage::ID: {
-      auto update_new_message = static_cast<const telegram_api::updateNewMessage *>(update.get());
-      DialogId dialog_id = get_message_dialog_id(update_new_message->message_);
-      if (!check_update_dialog_id(update, dialog_id)) {
-        return promise.set_value(Unit());
-      }
-      break;
-    }
-    case telegram_api::updateReadHistoryInbox::ID: {
-      auto update_read_history_inbox = static_cast<const telegram_api::updateReadHistoryInbox *>(update.get());
-      auto dialog_id = DialogId(update_read_history_inbox->peer_);
-      if (!check_update_dialog_id(update, dialog_id)) {
-        return promise.set_value(Unit());
-      }
-      break;
-    }
-    case telegram_api::updateReadHistoryOutbox::ID: {
-      auto update_read_history_outbox = static_cast<const telegram_api::updateReadHistoryOutbox *>(update.get());
-      auto dialog_id = DialogId(update_read_history_outbox->peer_);
-      if (!check_update_dialog_id(update, dialog_id)) {
-        return promise.set_value(Unit());
-      }
-      break;
-    }
-    case telegram_api::updateEditMessage::ID: {
-      auto update_edit_message = static_cast<const telegram_api::updateEditMessage *>(update.get());
-      DialogId dialog_id = get_message_dialog_id(update_edit_message->message_);
-      if (!check_update_dialog_id(update, dialog_id)) {
-        return promise.set_value(Unit());
-      }
-      break;
-    }
-    case telegram_api::updatePinnedMessages::ID: {
-      auto update_pinned_messages = static_cast<const telegram_api::updatePinnedMessages *>(update.get());
-      auto dialog_id = DialogId(update_pinned_messages->peer_);
-      if (!check_update_dialog_id(update, dialog_id)) {
-        return promise.set_value(Unit());
-      }
-      break;
-    }
-    default:
-      LOG(ERROR) << "Receive unexpected update " << oneline(to_string(update)) << "from " << source;
-      return;
+  if (!check_pts_update(update)) {
+    LOG(ERROR) << "Receive wrong pts update from " << source << ": " << oneline(to_string(update));
+    return promise.set_value(Unit());
   }
 
   if (force_apply) {
diff --git a/td/telegram/MessagesManager.h b/td/telegram/MessagesManager.h
index 9417fce5e..a7824d7a6 100644
--- a/td/telegram/MessagesManager.h
+++ b/td/telegram/MessagesManager.h
@@ -1826,7 +1826,9 @@ class MessagesManager : public Actor {
 
   bool can_set_game_score(DialogId dialog_id, const Message *m) const;
 
-  bool check_update_dialog_id(const tl_object_ptr<telegram_api::Update> &update, DialogId dialog_id);
+  static bool check_pts_update_dialog_id(DialogId dialog_id);
+
+  static bool check_pts_update(const tl_object_ptr<telegram_api::Update> &update);
 
   void process_update(tl_object_ptr<telegram_api::Update> &&update);