Support file reference repair for stories.

td/telegram/FileReferenceManager.cpp

@@ -17,6 +17,7 @@
 #include "td/telegram/NotificationSettingsManager.h"
 #include "td/telegram/StickerSetId.h"
 #include "td/telegram/StickersManager.h"
+#include "td/telegram/StoryManager.h"
 #include "td/telegram/Td.h"
 #include "td/telegram/WebPageId.h"
 #include "td/telegram/WebPagesManager.h"
@@ -77,6 +78,7 @@ fileSourceSavedRingtones = FileSource;                                   // repa
 fileSourceUserFull = FileSource;                                         // repaired with users.getFullUser
 fileSourceAttachmentMenuBot user_id:int53 = FileSource;                  // repaired with messages.getAttachMenuBot
 fileSourceWebApp user_id:int53 short_name:string = FileSource;           // repaired with messages.getAttachMenuBot
+fileSourceStory chat_id:int53 story_id:int32 = FileSource;               // repaired with stories.getStoriesByID
 */
 
 FileSourceId FileReferenceManager::get_current_file_source_id() const {
@@ -161,6 +163,11 @@ FileSourceId FileReferenceManager::create_web_app_file_source(UserId user_id, co
   return add_file_source_id(source, PSLICE() << "Web App " << user_id << '/' << short_name);
 }
 
+FileSourceId FileReferenceManager::create_story_file_source(StoryFullId story_full_id) {
+  FileSourceStory source{story_full_id};
+  return add_file_source_id(source, PSLICE() << story_full_id);
+}
+
 FileReferenceManager::Node &FileReferenceManager::add_node(NodeId node_id) {
   CHECK(node_id.is_valid());
   auto &node = nodes_[node_id];
@@ -370,6 +377,9 @@ void FileReferenceManager::send_query(Destination dest, FileSourceId file_source
       [&](const FileSourceWebApp &source) {
         send_closure_later(G()->attach_menu_manager(), &AttachMenuManager::reload_web_app, source.user_id,
                            source.short_name, std::move(promise));
+      },
+      [&](const FileSourceStory &source) {
+        send_closure_later(G()->story_manager(), &StoryManager::reload_story, source.story_full_id, std::move(promise));
       }));
 }
 

td/telegram/FileReferenceManager.h

@@ -14,6 +14,7 @@
 #include "td/telegram/FullMessageId.h"
 #include "td/telegram/PhotoSizeSource.h"
 #include "td/telegram/SetWithPosition.h"
+#include "td/telegram/StoryFullId.h"
 #include "td/telegram/td_api.h"
 #include "td/telegram/UserId.h"
 
@@ -64,6 +65,7 @@ class FileReferenceManager final : public Actor {
   FileSourceId create_user_full_file_source(UserId user_id);
   FileSourceId create_attach_menu_bot_file_source(UserId user_id);
   FileSourceId create_web_app_file_source(UserId user_id, const string &short_name);
+  FileSourceId create_story_file_source(StoryFullId story_full_id);
 
   using NodeId = FileId;
   void repair_file_reference(NodeId node_id, Promise<> promise);
@@ -169,13 +171,16 @@ class FileReferenceManager final : public Actor {
     UserId user_id;
     string short_name;
   };
+  struct FileSourceStory {
+    StoryFullId story_full_id;
+  };
 
   // append only
   using FileSource =
       Variant<FileSourceMessage, FileSourceUserPhoto, FileSourceChatPhoto, FileSourceChannelPhoto, FileSourceWallpapers,
               FileSourceWebPage, FileSourceSavedAnimations, FileSourceRecentStickers, FileSourceFavoriteStickers,
               FileSourceBackground, FileSourceChatFull, FileSourceChannelFull, FileSourceAppConfig,
-              FileSourceSavedRingtones, FileSourceUserFull, FileSourceAttachMenuBot, FileSourceWebApp>;
+              FileSourceSavedRingtones, FileSourceUserFull, FileSourceAttachMenuBot, FileSourceWebApp, FileSourceStory>;
   WaitFreeVector<FileSource> file_sources_;
 
   int64 query_generation_{0};

td/telegram/FileReferenceManager.hpp

@@ -18,6 +18,8 @@
 #include "td/telegram/MessagesManager.h"
 #include "td/telegram/NotificationSettingsManager.h"
 #include "td/telegram/StickersManager.h"
+#include "td/telegram/StoryFullId.h"
+#include "td/telegram/StoryManager.h"
 #include "td/telegram/Td.h"
 #include "td/telegram/UserId.h"
 #include "td/telegram/WebPagesManager.h"
@@ -58,7 +60,8 @@ void FileReferenceManager::store_file_source(FileSourceId file_source_id, Storer
                           [&](const FileSourceWebApp &source) {
                             td::store(source.user_id, storer);
                             td::store(source.short_name, storer);
-                          }));
+                          },
+                          [&](const FileSourceStory &source) { td::store(source.story_full_id, storer); }));
 }
 
 template <class ParserT>
@@ -141,6 +144,11 @@ FileSourceId FileReferenceManager::parse_file_source(Td *td, ParserT &parser) {
       td::parse(short_name, parser);
       return td->attach_menu_manager_->get_web_app_file_source_id(user_id, short_name);
     }
+    case 17: {
+      StoryFullId story_full_id;
+      td::parse(story_full_id, parser);
+      return td->story_manager_->get_story_file_source_id(story_full_id);
+    }
     default:
       parser.set_error("Invalid type in FileSource");
       return FileSourceId();

td/telegram/StoryManager.cpp

@@ -9,12 +9,14 @@
 #include "td/telegram/AuthManager.h"
 #include "td/telegram/ContactsManager.h"
 #include "td/telegram/files/FileManager.h"
+#include "td/telegram/FileReferenceManager.h"
 #include "td/telegram/Global.h"
 #include "td/telegram/MessageEntity.h"
 #include "td/telegram/StoryContent.h"
 #include "td/telegram/StoryContentType.h"
 #include "td/telegram/Td.h"
 
+#include "td/utils/algorithm.h"
 #include "td/utils/buffer.h"
 #include "td/utils/logging.h"
 #include "td/utils/Status.h"
@@ -110,7 +112,9 @@ td_api::object_ptr<td_api::story> StoryManager::get_story_object(StoryFullId sto
 }
 
 vector<FileId> StoryManager::get_story_file_ids(const Story *story) const {
-  CHECK(story != nullptr);
+  if (story == nullptr) {
+    return {};
+  }
   return get_story_content_file_ids(td_, story->content_.get());
 }
 
@@ -120,6 +124,25 @@ void StoryManager::delete_story_files(const Story *story) const {
   }
 }
 
+void StoryManager::change_story_files(StoryFullId story_full_id, const Story *story,
+                                      const vector<FileId> &old_file_ids) {
+  auto new_file_ids = get_story_file_ids(story);
+  if (new_file_ids == old_file_ids) {
+    return;
+  }
+
+  for (auto file_id : old_file_ids) {
+    if (!td::contains(new_file_ids, file_id)) {
+      send_closure(G()->file_manager(), &FileManager::delete_file, file_id, Promise<Unit>(), "change_story_files");
+    }
+  }
+
+  auto file_source_id = get_story_file_source_id(story_full_id);
+  if (file_source_id.is_valid()) {
+    td_->file_manager_->change_files_source(file_source_id, old_file_ids, new_file_ids);
+  }
+}
+
 StoryId StoryManager::on_get_story(DialogId owner_dialog_id,
                                    telegram_api::object_ptr<telegram_api::storyItem> &&story_item) {
   CHECK(story_item != nullptr);
@@ -149,6 +172,7 @@ StoryId StoryManager::on_get_story(DialogId owner_dialog_id,
     return StoryId();
   }
   auto content_type = content->get_type();
+  auto old_file_ids = get_story_file_ids(story);
   if (story->content_ == nullptr || story->content_->get_type() != content_type) {
     story->content_ = std::move(content);
     is_changed = true;
@@ -177,6 +201,8 @@ StoryId StoryManager::on_get_story(DialogId owner_dialog_id,
   }
 
   if (is_changed || need_save_to_database) {
+    change_story_files(story_full_id, story, old_file_ids);
+
     // save_story(story, story_id);
   }
   if (is_changed) {
@@ -216,13 +242,31 @@ std::pair<int32, vector<StoryId>> StoryManager::on_get_stories(
   return {total_count, std::move(story_ids)};
 }
 
+FileSourceId StoryManager::get_story_file_source_id(StoryFullId story_full_id) {
+  if (td_->auth_manager_->is_bot()) {
+    return FileSourceId();
+  }
+
+  auto dialog_id = story_full_id.get_dialog_id();
+  auto story_id = story_full_id.get_story_id();
+  if (!dialog_id.is_valid() || !story_id.is_valid()) {
+    return FileSourceId();
+  }
+
+  auto &file_source_id = story_full_id_to_file_source_id_[story_full_id];
+  if (!file_source_id.is_valid()) {
+    file_source_id = td_->file_reference_manager_->create_story_file_source(story_full_id);
+  }
+  return file_source_id;
+}
+
 void StoryManager::reload_story(StoryFullId story_full_id, Promise<Unit> &&promise) {
   auto dialog_id = story_full_id.get_dialog_id();
   if (dialog_id.get_type() != DialogType::User) {
     return promise.set_error(Status::Error(400, "Unsupported story owner"));
   }
   auto user_id = dialog_id.get_user_id();
-  td_->create_handler<GetStoriesByIDQuery>(std::move(promise))->send(user_id, {story_full_if.get_story_id().get()});
+  td_->create_handler<GetStoriesByIDQuery>(std::move(promise))->send(user_id, {story_full_id.get_story_id().get()});
 }
 
 }  // namespace td

td/telegram/StoryManager.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/MessageEntity.h"
 #include "td/telegram/StoryFullId.h"
 #include "td/telegram/StoryId.h"
@@ -45,6 +46,8 @@ class StoryManager final : public Actor {
 
   td_api::object_ptr<td_api::story> get_story_object(StoryFullId story_full_id) const;
 
+  FileSourceId get_story_file_source_id(StoryFullId story_full_id);
+
   void reload_story(StoryFullId story_full_id, Promise<Unit> &&promise);
 
  private:
@@ -74,8 +77,12 @@ class StoryManager final : public Actor {
 
   void delete_story_files(const Story *story) const;
 
+  void change_story_files(StoryFullId story_full_id, const Story *story, const vector<FileId> &old_file_ids);
+
   static bool is_local_story_id(StoryId story_id);
 
+  WaitFreeHashMap<StoryFullId, FileSourceId, StoryFullIdHash> story_full_id_to_file_source_id_;
+
   WaitFreeHashMap<StoryFullId, unique_ptr<Story>, StoryFullIdHash> stories_;
 
   Td *td_;