Add FileSourceRichMessage.

Author: levlam

td/telegram/FileReferenceManager.cpp

@@ -16,6 +16,7 @@
 #include "td/telegram/DraftMessageManager.h"
 #include "td/telegram/files/FileManager.h"
 #include "td/telegram/Global.h"
+#include "td/telegram/MessageQueryManager.h"
 #include "td/telegram/MessagesManager.h"
 #include "td/telegram/NotificationSettingsManager.h"
 #include "td/telegram/QuickReplyManager.h"
@@ -126,6 +127,7 @@ fileSourceBotMediaPreviewInfo bot_user_id:int53 language_code:string = FileSourc
 fileSourceStoryAlbum chat_id:int53 story_album_id:int32 = FileSource;                      // stories.getAlbums, not reliable
 fileSourceSavedMusic user_id:int53 file_id:int32 = FileSource;                             // users.getSavedMusicByID
 fileSourceDraftMessage chat_id:int53 topic:MessageTopic = FileSource;                      // messages.getPeerDialogs/messages.getForumTopicsByID/messages.getSavedDialogsByID
+fileSourceRichMessage chat_id:int53 message_id:int53 = FileSource;                         // messages.getRichMessage
 */
 
 FileSourceId FileReferenceManager::get_current_file_source_id() const {
@@ -253,6 +255,11 @@ FileSourceId FileReferenceManager::create_draft_message_file_source(DialogId dia
   return add_file_source_id(source, PSLICE() << "draft message in " << topic << " of " << dialog_id);
 }
 
+FileSourceId FileReferenceManager::create_rich_message_file_source(MessageFullId message_full_id) {
+  FileSourceRichMessage source{message_full_id};
+  return add_file_source_id(source, PSLICE() << "rich " << message_full_id);
+}
+
 FileReferenceManager::Node &FileReferenceManager::add_node(NodeId node_id) {
   CHECK(node_id.is_valid());
   auto &node = nodes_[node_id];
@@ -498,6 +505,10 @@ void FileReferenceManager::send_query(Destination dest, FileSourceId file_source
       [&](const FileSourceDraftMessage &source) {
         send_closure_later(G()->draft_message_manager(), &DraftMessageManager::reload_draft_message, source.dialog_id,
                            source.topic, std::move(promise));
+      },
+      [&](const FileSourceRichMessage &source) {
+        send_closure_later(G()->message_query_manager(), &MessageQueryManager::reload_full_rich_message,
+                           source.message_full_id, std::move(promise));
       }));
 }
 

td/telegram/FileReferenceManager.h

@@ -82,6 +82,7 @@ class FileReferenceManager final : public Actor {
   FileSourceId create_story_album_file_source(StoryAlbumFullId story_album_full_id);
   FileSourceId create_user_saved_music_file_source(UserId user_id, int64 document_id, int64 access_hash);
   FileSourceId create_draft_message_file_source(DialogId dialog_id, MessageTopic topic);
+  FileSourceId create_rich_message_file_source(MessageFullId message_full_id);
 
   using NodeId = FileId;
   void repair_file_reference(NodeId node_id, Promise<> promise);
@@ -217,6 +218,9 @@ class FileReferenceManager final : public Actor {
     DialogId dialog_id;
     MessageTopic topic;
   };
+  struct FileSourceRichMessage {
+    MessageFullId message_full_id;
+  };
 
   // append only
   using FileSource =
@@ -225,7 +229,8 @@ class FileReferenceManager final : public Actor {
               FileSourceBackground, FileSourceChatFull, FileSourceChannelFull, FileSourceAppConfig,
               FileSourceSavedRingtones, FileSourceUserFull, FileSourceAttachMenuBot, FileSourceWebApp, FileSourceStory,
               FileSourceQuickReplyMessage, FileSourceStarTransaction, FileSourceBotMediaPreview,
-              FileSourceBotMediaPreviewInfo, FileSourceStoryAlbum, FileSourceUserSavedMusic, FileSourceDraftMessage>;
+              FileSourceBotMediaPreviewInfo, FileSourceStoryAlbum, FileSourceUserSavedMusic, FileSourceDraftMessage,
+              FileSourceRichMessage>;
   WaitFreeVector<FileSource> file_sources_;
 
   int64 query_generation_{0};

td/telegram/FileReferenceManager.hpp

@@ -17,6 +17,7 @@
 #include "td/telegram/FileReferenceManager.h"
 #include "td/telegram/files/FileSourceId.h"
 #include "td/telegram/MessageFullId.h"
+#include "td/telegram/MessageQueryManager.h"
 #include "td/telegram/MessagesManager.h"
 #include "td/telegram/MessageTopic.hpp"
 #include "td/telegram/NotificationSettingsManager.h"
@@ -90,7 +91,8 @@ void FileReferenceManager::store_file_source(FileSourceId file_source_id, Storer
                           [&](const FileSourceDraftMessage &source) {
                             td::store(source.dialog_id, storer);
                             td::store(source.topic, storer);
-                          }));
+                          },
+                          [&](const FileSourceRichMessage &source) { td::store(source.message_full_id, storer); }));
 }
 
 template <class ParserT>
@@ -225,6 +227,11 @@ FileSourceId FileReferenceManager::parse_file_source(Td *td, ParserT &parser) {
       td::parse(topic, parser);
       return td->draft_message_manager_->get_draft_message_file_source_id(dialog_id, topic);
     }
+    case 25: {
+      MessageFullId message_full_id;
+      td::parse(message_full_id, parser);
+      return td->message_query_manager_->get_rich_message_file_source_id(message_full_id);
+    }
     default:
       parser.set_error("Invalid type in FileSource");
       return FileSourceId();

td/telegram/MessageQueryManager.cpp

@@ -16,6 +16,7 @@
 #include "td/telegram/DialogParticipant.h"
 #include "td/telegram/DialogParticipantManager.h"
 #include "td/telegram/FactCheck.h"
+#include "td/telegram/FileReferenceManager.h"
 #include "td/telegram/files/FileManager.h"
 #include "td/telegram/FolderId.h"
 #include "td/telegram/ForumTopicManager.h"
@@ -86,7 +87,6 @@ class GetRichMessageQuery final : public Td::ResultHandler {
       return on_error(Status::Error(400, "Chat not found"));
     }
     message_id_ = message_full_id.get_message_id();
-    CHECK(message_id_.is_server());
     send_query(G()->net_query_creator().create(
         telegram_api::messages_getRichMessage(std::move(input_peer), message_id_.get_server_message_id().get())));
   }
@@ -2145,9 +2145,42 @@ void MessageQueryManager::on_get_affected_history(DialogId dialog_id, AffectedHi
 
 void MessageQueryManager::get_full_rich_message(MessageFullId message_full_id,
                                                 Promise<td_api::object_ptr<td_api::richMessage>> &&promise) {
+  if (!message_full_id.get_message_id().is_server()) {
+    return promise.set_error(400, "Invalid message specified");
+  }
   td_->create_handler<GetRichMessageQuery>(std::move(promise))->send(message_full_id);
 }
 
+void MessageQueryManager::reload_full_rich_message(MessageFullId message_full_id, Promise<Unit> &&promise) {
+  return get_full_rich_message(
+      message_full_id, PromiseCreator::lambda([promise = std::move(promise)](
+                                                  Result<td_api::object_ptr<td_api::richMessage>> result) mutable {
+        if (result.is_error()) {
+          promise.set_error(result.move_as_error());
+        } else {
+          promise.set_value(Unit());
+        }
+      }));
+}
+
+FileSourceId MessageQueryManager::get_rich_message_file_source_id(MessageFullId message_full_id) {
+  if (td_->auth_manager_->is_bot()) {
+    return FileSourceId();
+  }
+
+  auto dialog_id = message_full_id.get_dialog_id();
+  auto message_id = message_full_id.get_message_id();
+  if (!dialog_id.is_valid() || dialog_id.get_type() == DialogType::SecretChat || !message_id.is_server()) {
+    return FileSourceId();
+  }
+
+  auto &file_source_id = rich_message_full_id_to_file_source_id_[message_full_id];
+  if (!file_source_id.is_valid()) {
+    file_source_id = td_->file_reference_manager_->create_rich_message_file_source(message_full_id);
+  }
+  return file_source_id;
+}
+
 void MessageQueryManager::upload_message_covers(BusinessConnectionId business_connection_id, DialogId dialog_id,
                                                 vector<MessageCover> covers, Promise<Unit> &&promise) {
   TRY_STATUS_PROMISE(promise, G()->close_status());

td/telegram/MessageQueryManager.h

@@ -13,6 +13,7 @@
 #include "td/telegram/DialogListId.h"
 #include "td/telegram/EmojiGameInfo.h"
 #include "td/telegram/files/FileId.h"
+#include "td/telegram/files/FileSourceId.h"
 #include "td/telegram/files/FileUploadId.h"
 #include "td/telegram/ForumTopicId.h"
 #include "td/telegram/MessageCover.h"
@@ -35,6 +36,7 @@
 #include "td/utils/FlatHashSet.h"
 #include "td/utils/Promise.h"
 #include "td/utils/Status.h"
+#include "td/utils/WaitFreeHashMap.h"
 
 #include <functional>
 #include <memory>
@@ -57,6 +59,10 @@ class MessageQueryManager final : public Actor {
 
   void get_full_rich_message(MessageFullId message_full_id, Promise<td_api::object_ptr<td_api::richMessage>> &&promise);
 
+  void reload_full_rich_message(MessageFullId message_full_id, Promise<Unit> &&promise);
+
+  FileSourceId get_rich_message_file_source_id(MessageFullId message_full_id);
+
   void upload_message_covers(BusinessConnectionId business_connection_id, DialogId dialog_id,
                              vector<MessageCover> covers, Promise<Unit> &&promise);
 
@@ -390,6 +396,8 @@ class MessageQueryManager final : public Actor {
   double emoji_game_info_receive_time_ = 0.0;
   EmojiGameInfo emoji_game_info_;
 
+  WaitFreeHashMap<MessageFullId, FileSourceId, MessageFullIdHash> rich_message_full_id_to_file_source_id_;
+
   MultiTimeout send_message_view_metrics_timeout_{"SendMessageViewMetricsTimeout"};
 
   Td *td_;

td/telegram/MessagesManager.cpp

@@ -14994,7 +14994,7 @@ void MessagesManager::get_full_rich_message(MessageFullId message_full_id,
   if (m == nullptr) {
     return promise.set_error(400, "Message not found");
   }
-  if (!m->message_id.is_server() || m->content->get_type() != MessageContentType::RichText) {
+  if (m->content->get_type() != MessageContentType::RichText) {
     return promise.set_error(400, "Invalid message specified");
   }
   td_->message_query_manager_->get_full_rich_message(message_full_id, std::move(promise));