Support file reference repair for files of attachment menu bots.

td/telegram/AttachMenuManager.cpp

@@ -12,6 +12,7 @@
 #include "td/telegram/Dependencies.h"
 #include "td/telegram/Document.h"
 #include "td/telegram/DocumentsManager.h"
+#include "td/telegram/FileReferenceManager.h"
 #include "td/telegram/files/FileId.hpp"
 #include "td/telegram/files/FileManager.h"
 #include "td/telegram/Global.h"
@@ -503,6 +504,21 @@ void AttachMenuManager::init() {
         }
         hash_ = is_cache_outdated ? 0 : attach_menu_bots_log_event.hash_;
         attach_menu_bots_ = std::move(attach_menu_bots_log_event.attach_menu_bots_);
+
+        for (auto attach_menu_bot : attach_menu_bots_) {
+          auto file_source_id = get_attach_menu_bot_file_source_id(attach_menu_bot.user_id_);
+          auto register_file_source = [&](FileId file_id) {
+            if (file_id.is_valid()) {
+              td_->file_manager_->add_file_source(file_id, file_source_id);
+            }
+          };
+          register_file_source(attach_menu_bot.default_icon_file_id_);
+          register_file_source(attach_menu_bot.ios_static_icon_file_id_);
+          register_file_source(attach_menu_bot.ios_animated_icon_file_id_);
+          register_file_source(attach_menu_bot.android_icon_file_id_);
+          register_file_source(attach_menu_bot.macos_icon_file_id_);
+          register_file_source(attach_menu_bot.placeholder_file_id_);
+        }
       } else {
         LOG(ERROR) << "Ignore invalid attachment menu bots log event";
       }
@@ -649,12 +665,14 @@ void AttachMenuManager::close_web_view(int64 query_id, Promise<Unit> &&promise)
 }
 
 Result<AttachMenuManager::AttachMenuBot> AttachMenuManager::get_attach_menu_bot(
-    tl_object_ptr<telegram_api::attachMenuBot> &&bot) const {
+    tl_object_ptr<telegram_api::attachMenuBot> &&bot) {
   UserId user_id(bot->bot_id_);
   if (!td_->contacts_manager_->have_user(user_id)) {
     return Status::Error(PSLICE() << "Have no information about " << user_id);
   }
 
+  auto file_source_id = get_attach_menu_bot_file_source_id(user_id);
+
   AttachMenuBot attach_menu_bot;
   attach_menu_bot.is_added_ = !bot->inactive_;
   attach_menu_bot.user_id_ = user_id;
@@ -704,6 +722,7 @@ Result<AttachMenuManager::AttachMenuBot> AttachMenuManager::get_attach_menu_bot(
       default:
         UNREACHABLE();
     }
+    td_->file_manager_->add_file_source(parsed_document.file_id, file_source_id);
     if (expect_colors) {
       if (icon->colors_.empty()) {
         LOG(ERROR) << "Have no colors for attachment menu bot icon for " << user_id;
@@ -882,6 +901,26 @@ void AttachMenuManager::get_attach_menu_bot(UserId user_id,
   td_->create_handler<GetAttachMenuBotQuery>(std::move(query_promise))->send(std::move(input_user));
 }
 
+void AttachMenuManager::reload_attach_menu_bot(UserId user_id, Promise<Unit> &&promise) {
+  TRY_RESULT_PROMISE(promise, input_user, td_->contacts_manager_->get_input_user(user_id));
+
+  auto wrapped_promise = PromiseCreator::lambda(
+      [promise = std::move(promise)](Result<td_api::object_ptr<td_api::attachmentMenuBot>> result) mutable {
+        if (result.is_error()) {
+          promise.set_error(result.move_as_error());
+        } else {
+          promise.set_value(Unit());
+        }
+      });
+  auto query_promise =
+      PromiseCreator::lambda([actor_id = actor_id(this), user_id, promise = std::move(wrapped_promise)](
+                                 Result<telegram_api::object_ptr<telegram_api::attachMenuBotsBot>> &&result) mutable {
+        send_closure(actor_id, &AttachMenuManager::on_get_attach_menu_bot, user_id, std::move(result),
+                     std::move(promise));
+      });
+  td_->create_handler<GetAttachMenuBotQuery>(std::move(query_promise))->send(std::move(input_user));
+}
+
 void AttachMenuManager::on_get_attach_menu_bot(
     UserId user_id, Result<telegram_api::object_ptr<telegram_api::attachMenuBotsBot>> &&result,
     Promise<td_api::object_ptr<td_api::attachmentMenuBot>> &&promise) {
@@ -921,6 +960,19 @@ void AttachMenuManager::on_get_attach_menu_bot(
   promise.set_value(get_attachment_menu_bot_object(attach_menu_bot));
 }
 
+FileSourceId AttachMenuManager::get_attach_menu_bot_file_source_id(UserId user_id) {
+  if (!user_id.is_valid() || !is_active()) {
+    return FileSourceId();
+  }
+
+  auto &source_id = attach_menu_bot_file_source_ids_[user_id];
+  if (!source_id.is_valid()) {
+    source_id = td_->file_reference_manager_->create_attach_menu_bot_file_source(user_id);
+  }
+  VLOG(file_references) << "Return " << source_id << " for attach menu bot " << user_id;
+  return source_id;
+}
+
 void AttachMenuManager::toggle_bot_is_added_to_attach_menu(UserId user_id, bool is_added, Promise<Unit> &&promise) {
   CHECK(is_active());
 

td/telegram/AttachMenuManager.h

@@ -8,6 +8,7 @@
 
 #include "td/telegram/DialogId.h"
 #include "td/telegram/files/FileId.h"
+#include "td/telegram/files/FileSourceId.h"
 #include "td/telegram/MessageId.h"
 #include "td/telegram/td_api.h"
 #include "td/telegram/telegram_api.h"
@@ -44,6 +45,10 @@ class AttachMenuManager final : public Actor {
 
   void get_attach_menu_bot(UserId user_id, Promise<td_api::object_ptr<td_api::attachmentMenuBot>> &&promise);
 
+  void reload_attach_menu_bot(UserId user_id, Promise<Unit> &&promise);
+
+  FileSourceId get_attach_menu_bot_file_source_id(UserId user_id);
+
   void toggle_bot_is_added_to_attach_menu(UserId user_id, bool is_added, Promise<Unit> &&promise);
 
   void get_current_state(vector<td_api::object_ptr<td_api::Update>> &updates) const;
@@ -117,7 +122,7 @@ class AttachMenuManager final : public Actor {
 
   void schedule_ping_web_view();
 
-  Result<AttachMenuBot> get_attach_menu_bot(tl_object_ptr<telegram_api::attachMenuBot> &&bot) const;
+  Result<AttachMenuBot> get_attach_menu_bot(tl_object_ptr<telegram_api::attachMenuBot> &&bot);
 
   td_api::object_ptr<td_api::attachmentMenuBot> get_attachment_menu_bot_object(const AttachMenuBot &bot) const;
 
@@ -144,6 +149,7 @@ class AttachMenuManager final : public Actor {
   bool is_inited_ = false;
   int64 hash_ = 0;
   vector<AttachMenuBot> attach_menu_bots_;
+  FlatHashMap<UserId, FileSourceId, UserIdHash> attach_menu_bot_file_source_ids_;
 
   struct OpenedWebView {
     DialogId dialog_id_;

td/telegram/FileReferenceManager.cpp

@@ -7,6 +7,7 @@
 #include "td/telegram/FileReferenceManager.h"
 
 #include "td/telegram/AnimationsManager.h"
+#include "td/telegram/AttachMenuManager.h"
 #include "td/telegram/BackgroundManager.h"
 #include "td/telegram/ConfigManager.h"
 #include "td/telegram/ContactsManager.h"
@@ -73,6 +74,8 @@ fileSourceBasicGroupFull basic_group_id:int32 = FileSource;              // repa
 fileSourceSupergroupFull supergroup_id:int32 = FileSource;               // repaired with messages.getFullChannel
 fileSourceAppConfig = FileSource;                                        // repaired with help.getAppConfig, not reliable
 fileSourceSavedRingtones = FileSource;                                   // repaired with account.getSavedRingtones
+fileSourceUserFull = FileSource;                                         // repaired with users.getFullUser
+fileSourceAttachmentMenuBot = FileSource;                                // repaired with messages.getAttachMenuBot
 */
 
 FileSourceId FileReferenceManager::get_current_file_source_id() const {
@@ -122,11 +125,6 @@ FileSourceId FileReferenceManager::create_background_file_source(BackgroundId ba
   return add_file_source_id(source, PSLICE() << background_id);
 }
 
-FileSourceId FileReferenceManager::create_user_full_file_source(UserId user_id) {
-  FileSourceUserFull source{user_id};
-  return add_file_source_id(source, PSLICE() << "full " << user_id);
-}
-
 FileSourceId FileReferenceManager::create_chat_full_file_source(ChatId chat_id) {
   FileSourceChatFull source{chat_id};
   return add_file_source_id(source, PSLICE() << "full " << chat_id);
@@ -142,6 +140,21 @@ FileSourceId FileReferenceManager::create_app_config_file_source() {
   return add_file_source_id(source, "app config");
 }
 
+FileSourceId FileReferenceManager::create_saved_ringtones_file_source() {
+  FileSourceSavedRingtones source;
+  return add_file_source_id(source, "saved notification sounds");
+}
+
+FileSourceId FileReferenceManager::create_user_full_file_source(UserId user_id) {
+  FileSourceUserFull source{user_id};
+  return add_file_source_id(source, PSLICE() << "full " << user_id);
+}
+
+FileSourceId FileReferenceManager::create_attach_menu_bot_file_source(UserId user_id) {
+  FileSourceAttachMenuBot source{user_id};
+  return add_file_source_id(source, PSLICE() << "attachment menu bot " << user_id);
+}
+
 bool FileReferenceManager::add_file_source(NodeId node_id, FileSourceId file_source_id) {
   CHECK(node_id.is_valid());
   bool is_added = nodes_[node_id].file_source_ids.add(file_source_id);
@@ -149,11 +162,6 @@ bool FileReferenceManager::add_file_source(NodeId node_id, FileSourceId file_sou
   return is_added;
 }
 
-FileSourceId FileReferenceManager::create_saved_ringtones_file_source() {
-  FileSourceSavedRingtones source;
-  return add_file_source_id(source, "saved notification sounds");
-}
-
 bool FileReferenceManager::remove_file_source(NodeId node_id, FileSourceId file_source_id) {
   CHECK(node_id.is_valid());
   bool is_removed = nodes_[node_id].file_source_ids.remove(file_source_id);
@@ -336,6 +344,10 @@ void FileReferenceManager::send_query(Destination dest, FileSourceId file_source
       [&](const FileSourceUserFull &source) {
         send_closure_later(G()->contacts_manager(), &ContactsManager::reload_user_full, source.user_id,
                            std::move(promise));
+      },
+      [&](const FileSourceAttachMenuBot &source) {
+        send_closure_later(G()->attach_menu_manager(), &AttachMenuManager::reload_attach_menu_bot, source.user_id,
+                           std::move(promise));
       }));
 }
 

td/telegram/FileReferenceManager.h

@@ -57,11 +57,12 @@ class FileReferenceManager final : public Actor {
   FileSourceId create_recent_stickers_file_source(bool is_attached);
   FileSourceId create_favorite_stickers_file_source();
   FileSourceId create_background_file_source(BackgroundId background_id, int64 access_hash);
-  FileSourceId create_user_full_file_source(UserId user_id);
   FileSourceId create_chat_full_file_source(ChatId chat_id);
   FileSourceId create_channel_full_file_source(ChannelId channel_id);
   FileSourceId create_app_config_file_source();
   FileSourceId create_saved_ringtones_file_source();
+  FileSourceId create_user_full_file_source(UserId user_id);
+  FileSourceId create_attach_menu_bot_file_source(UserId user_id);
 
   using NodeId = FileId;
   void repair_file_reference(NodeId node_id, Promise<> promise);
@@ -160,13 +161,16 @@ class FileReferenceManager final : public Actor {
   struct FileSourceUserFull {
     UserId user_id;
   };
+  struct FileSourceAttachMenuBot {
+    UserId user_id;
+  };
 
   // append only
   using FileSource =
       Variant<FileSourceMessage, FileSourceUserPhoto, FileSourceChatPhoto, FileSourceChannelPhoto, FileSourceWallpapers,
               FileSourceWebPage, FileSourceSavedAnimations, FileSourceRecentStickers, FileSourceFavoriteStickers,
               FileSourceBackground, FileSourceChatFull, FileSourceChannelFull, FileSourceAppConfig,
-              FileSourceSavedRingtones, FileSourceUserFull>;
+              FileSourceSavedRingtones, FileSourceUserFull, FileSourceAttachMenuBot>;
   WaitFreeVector<FileSource> file_sources_;
 
   int64 query_generation_{0};

td/telegram/FileReferenceManager.hpp

@@ -7,6 +7,7 @@
 #pragma once
 
 #include "td/telegram/AnimationsManager.h"
+#include "td/telegram/AttachMenuManager.h"
 #include "td/telegram/BackgroundManager.h"
 #include "td/telegram/ChannelId.h"
 #include "td/telegram/ChatId.h"
@@ -52,7 +53,8 @@ void FileReferenceManager::store_file_source(FileSourceId file_source_id, Storer
                           [&](const FileSourceChatFull &source) { td::store(source.chat_id, storer); },
                           [&](const FileSourceChannelFull &source) { td::store(source.channel_id, storer); },
                           [&](const FileSourceAppConfig &source) {}, [&](const FileSourceSavedRingtones &source) {},
-                          [&](const FileSourceUserFull &source) { td::store(source.user_id, storer); }));
+                          [&](const FileSourceUserFull &source) { td::store(source.user_id, storer); },
+                          [&](const FileSourceAttachMenuBot &source) { td::store(source.user_id, storer); }));
 }
 
 template <class ParserT>
@@ -123,6 +125,11 @@ FileSourceId FileReferenceManager::parse_file_source(Td *td, ParserT &parser) {
       td::parse(user_id, parser);
       return td->contacts_manager_->get_user_full_file_source_id(user_id);
     }
+    case 15: {
+      UserId user_id;
+      td::parse(user_id, parser);
+      return td->attach_menu_manager_->get_attach_menu_bot_file_source_id(user_id);
+    }
     default:
       parser.set_error("Invalid type in FileSource");
       return FileSourceId();

td/telegram/Global.h

@@ -32,6 +32,7 @@
 namespace td {
 
 class AnimationsManager;
+class AttachMenuManager;
 class BackgroundManager;
 class CallManager;
 class ConfigManager;
@@ -182,6 +183,13 @@ class Global final : public ActorContext {
     animations_manager_ = animations_manager;
   }
 
+  ActorId<AttachMenuManager> attach_menu_manager() const {
+    return attach_menu_manager_;
+  }
+  void set_attach_menu_manager(ActorId<AttachMenuManager> attach_menu_manager) {
+    attach_menu_manager_ = attach_menu_manager;
+  }
+
   ActorId<BackgroundManager> background_manager() const {
     return background_manager_;
   }
@@ -454,6 +462,7 @@ class Global final : public ActorContext {
 
   ActorId<Td> td_;
   ActorId<AnimationsManager> animations_manager_;
+  ActorId<AttachMenuManager> attach_menu_manager_;
   ActorId<BackgroundManager> background_manager_;
   ActorId<CallManager> call_manager_;
   ActorId<ConfigManager> config_manager_;

td/telegram/Td.cpp

@@ -3957,6 +3957,7 @@ void Td::init_managers() {
   G()->set_animations_manager(animations_manager_actor_.get());
   attach_menu_manager_ = make_unique<AttachMenuManager>(this, create_reference());
   attach_menu_manager_actor_ = register_actor("AttachMenuManager", attach_menu_manager_.get());
+  G()->set_attach_menu_manager(attach_menu_manager_actor_.get());
   background_manager_ = make_unique<BackgroundManager>(this, create_reference());
   background_manager_actor_ = register_actor("BackgroundManager", background_manager_.get());
   G()->set_background_manager(background_manager_actor_.get());