From 90fd69b2029ea2989c4668a52fe39a561b43bb08 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Sat, 14 Oct 2023 22:48:16 +0300
Subject: [PATCH] Add UpdatesManager::is_acceptable_message_media.

---
 td/telegram/UpdatesManager.cpp | 115 +++++++++++++++++----------------
 td/telegram/UpdatesManager.h   |   2 +
 2 files changed, 62 insertions(+), 55 deletions(-)

diff --git a/td/telegram/UpdatesManager.cpp b/td/telegram/UpdatesManager.cpp
index 89a824d4d..53e8f4d77 100644
--- a/td/telegram/UpdatesManager.cpp
+++ b/td/telegram/UpdatesManager.cpp
@@ -807,6 +807,64 @@ bool UpdatesManager::is_acceptable_message_forward_header(
   return true;
 }
 
+bool UpdatesManager::is_acceptable_message_media(
+    const telegram_api::object_ptr<telegram_api::MessageMedia> &media_ptr) const {
+  if (media_ptr == nullptr) {
+    return true;
+  }
+  switch (media_ptr->get_id()) {
+    case telegram_api::messageMediaContact::ID: {
+      auto message_media = static_cast<const telegram_api::messageMediaContact *>(media_ptr.get());
+      UserId user_id(message_media->user_id_);
+      return user_id == UserId() || is_acceptable_user(user_id);
+    }
+    case telegram_api::messageMediaStory::ID: {
+      auto message_media = static_cast<const telegram_api::messageMediaStory *>(media_ptr.get());
+      return is_acceptable_peer(message_media->peer_);
+    }
+    case telegram_api::messageMediaPoll::ID:
+      /*
+      // the users and chats are always min, so no need to check
+      auto message_media_poll = static_cast<const telegram_api::messageMediaPoll *>(media_ptr.get());
+      for (auto recent_voter : message_media_poll->results_->recent_voters_) {
+        if (!is_acceptable_peer(recent_voter)) {
+          return false;
+        }
+      }
+      */
+      return true;
+    case telegram_api::messageMediaWebPage::ID:
+      /*
+      // the channel is always min, so no need to check
+      auto message_media_web_page = static_cast<const telegram_api::messageMediaWebPage *>(media_ptr.get());
+      if (message_media_web_page->webpage_->get_id() == telegram_api::webPage::ID) {
+        auto web_page = static_cast<const telegram_api::webPage *>(message_media_web_page->webpage_.get());
+        if (web_page->cached_page_ != nullptr) {
+          const vector<tl_object_ptr<telegram_api::PageBlock>> *page_blocks = nullptr;
+          downcast_call(*web_page->cached_page_, [&page_blocks](auto &page) { page_blocks = &page.blocks_; });
+          CHECK(page_blocks != nullptr);
+          for (auto &page_block : *page_blocks) {
+            if (page_block->get_id() == telegram_api::pageBlockChannel::ID) {
+              auto page_block_channel = static_cast<const telegram_api::pageBlockChannel *>(page_block.get());
+              auto channel_id = ContactsManager::get_channel_id(page_block_channel->channel_);
+              if (channel_id.is_valid()) {
+                if (!is_acceptable_channel(channel_id)) {
+                  return false;
+                }
+              } else {
+                LOG(ERROR) << "Receive wrong channel " << to_string(page_block_channel->channel_);
+              }
+            }
+          }
+        }
+      }
+      */
+      return true;
+    default:
+      return true;
+  }
+}
+
 bool UpdatesManager::is_acceptable_message(const telegram_api::Message *message_ptr) const {
   CHECK(message_ptr != nullptr);
   int32 constructor_id = message_ptr->get_id();
@@ -840,61 +898,8 @@ bool UpdatesManager::is_acceptable_message(const telegram_api::Message *message_
         return false;
       }
 
-      if (message->media_ != nullptr) {
-        auto media_id = message->media_->get_id();
-        if (media_id == telegram_api::messageMediaContact::ID) {
-          auto message_media = static_cast<const telegram_api::messageMediaContact *>(message->media_.get());
-          UserId user_id(message_media->user_id_);
-          if (user_id != UserId() && !is_acceptable_user(user_id)) {
-            return false;
-          }
-        }
-        if (media_id == telegram_api::messageMediaStory::ID) {
-          auto message_media = static_cast<const telegram_api::messageMediaStory *>(message->media_.get());
-          if (!is_acceptable_peer(message_media->peer_)) {
-            return false;
-          }
-        }
-        /*
-        // the users and chats are always min, so no need to check
-        if (media_id == telegram_api::messageMediaPoll::ID) {
-          auto message_media_poll = static_cast<const telegram_api::messageMediaPoll *>(message->media_.get());
-          for (auto recent_voter : message_media_poll->results_->recent_voters_) {
-            if (!is_acceptable_peer(recent_voter)) {
-              return false;
-            }
-          }
-        }
-        */
-        /*
-        // the channel is always min, so no need to check
-        if (media_id == telegram_api::messageMediaWebPage::ID) {
-          auto message_media_web_page = static_cast<const telegram_api::messageMediaWebPage *>(message->media_.get());
-          if (message_media_web_page->webpage_->get_id() == telegram_api::webPage::ID) {
-            auto web_page = static_cast<const telegram_api::webPage *>(message_media_web_page->webpage_.get());
-            if (web_page->cached_page_ != nullptr) {
-              const vector<tl_object_ptr<telegram_api::PageBlock>> *page_blocks = nullptr;
-              downcast_call(*web_page->cached_page_, [&page_blocks](auto &page) { page_blocks = &page.blocks_; });
-              CHECK(page_blocks != nullptr);
-              for (auto &page_block : *page_blocks) {
-                if (page_block->get_id() == telegram_api::pageBlockChannel::ID) {
-                  auto page_block_channel = static_cast<const telegram_api::pageBlockChannel *>(page_block.get());
-                  auto channel_id = ContactsManager::get_channel_id(page_block_channel->channel_);
-                  if (channel_id.is_valid()) {
-                    if (!is_acceptable_channel(channel_id)) {
-                      return false;
-                    }
-                  } else {
-                    LOG(ERROR) << "Receive wrong channel " << to_string(page_block_channel->channel_);
-                  }
-                }
-              }
-            }
-          }
-        }
-        */
-      } else {
-        CHECK(message->media_ == nullptr);
+      if (!is_acceptable_message_media(message->media_)) {
+        return false;
       }
 
       /*
diff --git a/td/telegram/UpdatesManager.h b/td/telegram/UpdatesManager.h
index 6ccce6c30..6e1182c5d 100644
--- a/td/telegram/UpdatesManager.h
+++ b/td/telegram/UpdatesManager.h
@@ -461,6 +461,8 @@ class UpdatesManager final : public Actor {
   bool is_acceptable_message_forward_header(
       const telegram_api::object_ptr<telegram_api::messageFwdHeader> &header) const;
 
+  bool is_acceptable_message_media(const telegram_api::object_ptr<telegram_api::MessageMedia> &media_ptr) const;
+
   bool is_acceptable_message(const telegram_api::Message *message_ptr) const;
 
   bool is_acceptable_update(const telegram_api::Update *update) const;