Support paid media message contents with multiple files in MessagesManager.

td/telegram/MessageContent.cpp

@@ -3537,7 +3537,9 @@ tl_object_ptr<telegram_api::InputMedia> get_input_media(const MessageContent *co
   auto file_references = FileManager::extract_file_references(input_media);
   for (size_t i = 0; i < file_references.size(); i++) {
     if (file_references[i] == FileReferenceView::invalid_file_reference()) {
-      auto file_id = get_message_content_any_file_id(content);
+      auto file_ids = get_message_content_any_file_ids(content);
+      CHECK(file_ids.size() == file_references.size());
+      auto file_id = file_ids[i];
       if (!force) {
         LOG(INFO) << "File " << file_id << " has invalid file reference";
         return nullptr;
@@ -7799,12 +7801,27 @@ FileId get_message_content_upload_file_id(const MessageContent *content) {
       return static_cast<const MessageVideoNote *>(content)->file_id;
     case MessageContentType::VoiceNote:
       return static_cast<const MessageVoiceNote *>(content)->file_id;
+    case MessageContentType::PaidMedia:
+      UNREACHABLE();
+      break;
     default:
       break;
   }
   return FileId();
 }
 
+vector<FileId> get_message_content_upload_file_ids(const MessageContent *content) {
+  if (content->get_type() == MessageContentType::PaidMedia) {
+    return transform(static_cast<const MessagePaidMedia *>(content)->media,
+                     [](const MessageExtendedMedia &media) { return media.get_upload_file_id(); });
+  }
+  auto file_id = get_message_content_upload_file_id(content);
+  if (file_id.is_valid()) {
+    return {file_id};
+  }
+  return {};
+}
+
 FileId get_message_content_any_file_id(const MessageContent *content) {
   FileId result = get_message_content_upload_file_id(content);
   if (!result.is_valid()) {
@@ -7817,6 +7834,18 @@ FileId get_message_content_any_file_id(const MessageContent *content) {
   return result;
 }
 
+vector<FileId> get_message_content_any_file_ids(const MessageContent *content) {
+  if (content->get_type() == MessageContentType::PaidMedia) {
+    return transform(static_cast<const MessagePaidMedia *>(content)->media,
+                     [](const MessageExtendedMedia &media) { return media.get_any_file_id(); });
+  }
+  auto file_id = get_message_content_any_file_id(content);
+  if (file_id.is_valid()) {
+    return {file_id};
+  }
+  return {};
+}
+
 void update_message_content_file_id_remote(MessageContent *content, FileId file_id) {
   if (file_id.get_remote() == 0) {
     return;
@@ -7837,6 +7866,9 @@ void update_message_content_file_id_remote(MessageContent *content, FileId file_
         return &static_cast<MessageVideoNote *>(content)->file_id;
       case MessageContentType::VoiceNote:
         return &static_cast<MessageVoiceNote *>(content)->file_id;
+      case MessageContentType::PaidMedia:
+        UNREACHABLE();
+        return static_cast<FileId *>(nullptr);
       default:
         return static_cast<FileId *>(nullptr);
     }
@@ -7846,6 +7878,23 @@ void update_message_content_file_id_remote(MessageContent *content, FileId file_
   }
 }
 
+void update_message_content_file_id_remotes(MessageContent *content, const vector<FileId> &file_ids) {
+  if (content->get_type() == MessageContentType::PaidMedia) {
+    auto &media = static_cast<MessagePaidMedia *>(content)->media;
+    if (file_ids.size() != media.size()) {
+      return;
+    }
+    for (size_t i = 0; i < file_ids.size(); i++) {
+      media[i].update_file_id_remote(file_ids[i]);
+    }
+    return;
+  }
+  if (file_ids.size() != 1) {
+    return;
+  }
+  update_message_content_file_id_remote(content, file_ids[0]);
+}
+
 FileId get_message_content_thumbnail_file_id(const MessageContent *content, const Td *td) {
   switch (content->get_type()) {
     case MessageContentType::Animation:
@@ -7870,12 +7919,27 @@ FileId get_message_content_thumbnail_file_id(const MessageContent *content, cons
           static_cast<const MessageVideoNote *>(content)->file_id);
     case MessageContentType::VoiceNote:
       return FileId();
+    case MessageContentType::PaidMedia:
+      UNREACHABLE();
+      return FileId();
     default:
       break;
   }
   return FileId();
 }
 
+vector<FileId> get_message_content_thumbnail_file_ids(const MessageContent *content, const Td *td) {
+  if (content->get_type() == MessageContentType::PaidMedia) {
+    return transform(static_cast<const MessagePaidMedia *>(content)->media,
+                     [td](const MessageExtendedMedia &media) { return media.get_thumbnail_file_id(td); });
+  }
+  auto file_id = get_message_content_thumbnail_file_id(content, td);
+  if (file_id.is_valid()) {
+    return {file_id};
+  }
+  return {};
+}
+
 vector<FileId> get_message_content_file_ids(const MessageContent *content, const Td *td) {
   CHECK(content != nullptr);
   switch (content->get_type()) {

td/telegram/MessageContent.h

@@ -276,12 +276,20 @@ const Photo *get_message_content_photo(const MessageContent *content);
 
 FileId get_message_content_upload_file_id(const MessageContent *content);
 
+vector<FileId> get_message_content_upload_file_ids(const MessageContent *content);
+
 FileId get_message_content_any_file_id(const MessageContent *content);
 
+vector<FileId> get_message_content_any_file_ids(const MessageContent *content);
+
 void update_message_content_file_id_remote(MessageContent *content, FileId file_id);
 
+void update_message_content_file_id_remotes(MessageContent *content, const vector<FileId> &file_ids);
+
 FileId get_message_content_thumbnail_file_id(const MessageContent *content, const Td *td);
 
+vector<FileId> get_message_content_thumbnail_file_ids(const MessageContent *content, const Td *td);
+
 vector<FileId> get_message_content_file_ids(const MessageContent *content, const Td *td);
 
 StoryFullId get_message_content_story_full_id(const Td *td, const MessageContent *content);

td/telegram/MessageExtendedMedia.cpp

@@ -316,6 +316,15 @@ FileId MessageExtendedMedia::get_thumbnail_file_id(const Td *td) const {
   return FileId();
 }
 
+void MessageExtendedMedia::update_file_id_remote(FileId file_id) {
+  if (file_id.get_remote() == 0 || type_ != Type::Video) {
+    return;
+  }
+  if (video_file_id_ == file_id && video_file_id_.get_remote() == 0) {
+    video_file_id_ = file_id;
+  }
+}
+
 telegram_api::object_ptr<telegram_api::InputMedia> MessageExtendedMedia::get_input_media(
     Td *td, tl_object_ptr<telegram_api::InputFile> input_file,
     tl_object_ptr<telegram_api::InputFile> input_thumbnail) const {

td/telegram/MessageExtendedMedia.h

@@ -96,6 +96,8 @@ class MessageExtendedMedia {
 
   FileId get_thumbnail_file_id(const Td *td) const;
 
+  void update_file_id_remote(FileId file_id);
+
   telegram_api::object_ptr<telegram_api::InputMedia> get_input_media(
       Td *td, tl_object_ptr<telegram_api::InputFile> input_file,
       tl_object_ptr<telegram_api::InputFile> input_thumbnail) const;

td/telegram/MessagesManager.cpp

@@ -23298,14 +23298,18 @@ vector<FileId> MessagesManager::get_message_file_ids(const Message *m) const {
 }
 
 void MessagesManager::cancel_upload_message_content_files(const MessageContent *content) {
-  auto file_id = get_message_content_upload_file_id(content);
+  auto file_ids = get_message_content_upload_file_ids(content);
   // always cancel file upload, it should be a no-op in the worst case
-  if (being_uploaded_files_.erase(file_id) || file_id.is_valid()) {
+  for (auto file_id : file_ids) {
-    cancel_upload_file(file_id, "cancel_upload_message_content_files");
+    if (being_uploaded_files_.erase(file_id) || file_id.is_valid()) {
+      cancel_upload_file(file_id, "cancel_upload_message_content_files");
+    }
   }
-  file_id = get_message_content_thumbnail_file_id(content, td_);
+  file_ids = get_message_content_thumbnail_file_ids(content, td_);
-  if (being_uploaded_thumbnails_.erase(file_id) || file_id.is_valid()) {
+  for (auto file_id : file_ids) {
-    cancel_upload_file(file_id, "cancel_upload_message_content_files");
+    if (being_uploaded_thumbnails_.erase(file_id) || file_id.is_valid()) {
+      cancel_upload_file(file_id, "cancel_upload_message_content_files");
+    }
   }
 }
 
@@ -23963,10 +23967,18 @@ void MessagesManager::do_send_message(DialogId dialog_id, const Message *m, vect
     return;
   }
 
-  FileId file_id = get_message_content_any_file_id(content);  // any_file_id, because it could be a photo sent by ID
+  auto file_ids = get_message_content_any_file_ids(content);  // any_file_ids, because it could be a photo sent by ID
-  FileId thumbnail_file_id = get_message_content_thumbnail_file_id(content, td_);
+  auto thumbnail_file_ids = get_message_content_thumbnail_file_ids(content, td_);
-  LOG(DEBUG) << "Need to send file " << file_id << " with thumbnail " << thumbnail_file_id;
+  LOG(DEBUG) << "Need to send files " << file_ids << " with thumbnails " << thumbnail_file_ids;
+  if (file_ids.size() != thumbnail_file_ids.size()) {
+    CHECK(file_ids.size() == 1u);
+    CHECK(thumbnail_file_ids.empty());
+  }
   if (is_secret) {
+    CHECK(file_ids.size() <= 1);
+    auto file_id = file_ids.empty() ? FileId() : file_ids[0];
+    auto thumbnail_file_id = thumbnail_file_ids.empty() ? FileId() : thumbnail_file_ids[0];
+
     CHECK(!is_edit);
     auto layer = td_->user_manager_->get_secret_chat_layer(dialog_id.get_secret_chat_id());
     auto secret_input_media = get_secret_input_media(content, td_, nullptr, BufferSlice(), layer);
@@ -23993,23 +24005,32 @@ void MessagesManager::do_send_message(DialogId dialog_id, const Message *m, vect
           content_type == MessageContentType::Story) {
         return;
       }
-      CHECK(file_id.is_valid());
+      CHECK(!file_ids.empty());
-      FileView file_view = td_->file_manager_->get_file_view(file_id);
+      for (size_t i = 0; i < file_ids.size(); i++) {
-      if (get_main_file_type(file_view.get_type()) == FileType::Photo) {
+        auto file_id = file_ids[i];
-        thumbnail_file_id = FileId();
+        CHECK(file_id.is_valid());
-      }
 
-      LOG(INFO) << "Ask to upload file " << file_id << " with bad parts " << bad_parts;
+        auto thumbnail_file_id = i < thumbnail_file_ids.size() ? thumbnail_file_ids[i] : FileId();
-      CHECK(file_id.is_valid());
+        FileView file_view = td_->file_manager_->get_file_view(file_id);
-      bool is_inserted =
+        if (get_main_file_type(file_view.get_type()) == FileType::Photo) {
-          being_uploaded_files_
+          thumbnail_file_id = FileId();
-              .emplace(file_id, std::make_pair(MessageFullId(dialog_id, m->message_id), thumbnail_file_id))
+        }
-              .second;
+
-      CHECK(is_inserted);
+        LOG(INFO) << "Ask to upload file " << file_id << " with bad parts " << bad_parts;
-      // need to call resume_upload synchronously to make upload process consistent with being_uploaded_files_
+        bool is_inserted =
-      // and to send is_uploading_active == true in the updates
+            being_uploaded_files_
-      td_->file_manager_->resume_upload(file_id, std::move(bad_parts), upload_media_callback_, 1, m->message_id.get());
+                .emplace(file_id, std::make_pair(MessageFullId(dialog_id, m->message_id), thumbnail_file_id))
+                .second;
+        CHECK(is_inserted);
+        // need to call resume_upload synchronously to make upload process consistent with being_uploaded_files_
+        // and to send is_uploading_active == true in the updates
+        td_->file_manager_->resume_upload(file_id, std::move(bad_parts), upload_media_callback_, 1,
+                                          m->message_id.get());
+      }
     } else {
+      CHECK(file_ids.size() <= 1);  // TODO PaidMedia
+      auto file_id = file_ids.empty() ? FileId() : file_ids[0];
+      auto thumbnail_file_id = thumbnail_file_ids.empty() ? FileId() : thumbnail_file_ids[0];
       on_message_media_uploaded(dialog_id, m, std::move(input_media), file_id, thumbnail_file_id);
     }
   }
@@ -31106,7 +31127,7 @@ void MessagesManager::on_send_dialog_action_timeout(DialogId dialog_id) {
   }
   CHECK(m->message_id.is_yet_unsent());
   if (m->forward_info != nullptr || m->had_forward_info || m->is_copy || m->message_id.is_scheduled() ||
-      m->sender_dialog_id.is_valid()) {
+      m->sender_dialog_id.is_valid() || m->content->get_type() == MessageContentType::PaidMedia) {
     return;
   }
 
@@ -33927,8 +33948,7 @@ bool MessagesManager::update_message_content(DialogId dialog_id, Message *old_me
   MessageContentType old_content_type = old_content->get_type();
   MessageContentType new_content_type = new_content->get_type();
 
-  auto old_file_id = get_message_content_any_file_id(old_content.get());
+  auto old_file_ids = get_message_content_any_file_ids(old_content.get());
-  bool need_finish_upload = old_file_id.is_valid() && need_merge_files;
   if (old_content_type != new_content_type) {
     if (old_message->ttl.is_valid() && old_message->ttl_expires_at > 0 &&
         is_expired_message_content(new_content_type) &&
@@ -33941,18 +33961,25 @@ bool MessagesManager::update_message_content(DialogId dialog_id, Message *old_me
       old_message->is_content_secret = old_message->ttl.is_secret_message_content(new_content->get_type());
     }
 
-    if (need_merge_files && old_file_id.is_valid()) {
+    if (need_merge_files) {
-      td_->file_manager_->try_merge_documents(old_file_id, get_message_content_any_file_id(new_content.get()));
+      auto new_file_ids = get_message_content_any_file_ids(new_content.get());
+      if (new_file_ids.size() == old_file_ids.size()) {
+        for (size_t i = 0; i < old_file_ids.size(); i++) {
+          td_->file_manager_->try_merge_documents(old_file_ids[i], new_file_ids[i]);
+        }
+      }
     }
   } else {
     merge_message_contents(td_, old_content.get(), new_content.get(), need_message_changed_warning(old_message),
                            dialog_id, need_merge_files, is_content_changed, need_update);
     compare_message_contents(td_, old_content.get(), new_content.get(), is_content_changed, need_update);
   }
-  if (need_finish_upload) {
+  if (need_merge_files) {
     // the file is likely to be already merged with a server file, but if not we need to
     // cancel file upload of the main file to allow next upload with the same file to succeed
-    cancel_upload_file(old_file_id, "update_message_content");
+    for (auto old_file_id : old_file_ids) {
+      cancel_upload_file(old_file_id, "update_message_content");
+    }
   }
 
   if (is_content_changed || need_update) {
@@ -33962,9 +33989,9 @@ bool MessagesManager::update_message_content(DialogId dialog_id, Message *old_me
     }
     old_content = std::move(new_content);
     old_message->last_edit_pts = 0;
-    update_message_content_file_id_remote(old_content.get(), old_file_id);
+    update_message_content_file_id_remotes(old_content.get(), old_file_ids);
   } else {
-    update_message_content_file_id_remote(old_content.get(), get_message_content_any_file_id(new_content.get()));
+    update_message_content_file_id_remotes(old_content.get(), get_message_content_any_file_ids(new_content.get()));
   }
   if (is_content_changed && !need_update) {
     LOG(INFO) << "Content of " << old_message->message_id << " in " << dialog_id << " has changed";