diff --git a/td/telegram/files/FileLoadManager.cpp b/td/telegram/files/FileLoadManager.cpp index 9549cba3f..6cc51c818 100644 --- a/td/telegram/files/FileLoadManager.cpp +++ b/td/telegram/files/FileLoadManager.cpp @@ -146,6 +146,20 @@ void FileLoadManager::unlink_file(string file_path, Promise promise) { promise.set_value(Unit()); } +void FileLoadManager::check_full_local_location(FullLocalLocationInfo local_info, bool skip_file_size_checks, + Promise promise) { + promise.set_result(::td::check_full_local_location(std::move(local_info), skip_file_size_checks)); +} + +void FileLoadManager::check_partial_local_location(PartialLocalFileLocation partial, Promise promise) { + auto status = ::td::check_partial_local_location(partial); + if (status.is_error()) { + promise.set_error(std::move(status)); + } else { + promise.set_value(Unit()); + } +} + // void upload_reload_parts(QueryId id, vector parts); // void upload_restart(QueryId id); void FileLoadManager::cancel(QueryId id) { diff --git a/td/telegram/files/FileLoadManager.h b/td/telegram/files/FileLoadManager.h index 43ab076ae..76b79b7a7 100644 --- a/td/telegram/files/FileLoadManager.h +++ b/td/telegram/files/FileLoadManager.h @@ -10,6 +10,7 @@ #include "td/telegram/files/FileEncryptionKey.h" #include "td/telegram/files/FileFromBytes.h" #include "td/telegram/files/FileHashUploader.h" +#include "td/telegram/files/FileLoaderUtils.h" #include "td/telegram/files/FileLocation.h" #include "td/telegram/files/FileType.h" #include "td/telegram/files/FileUploader.h" @@ -64,6 +65,11 @@ class FileLoadManager final : public Actor { void unlink_file(string file_path, Promise promise); + void check_full_local_location(FullLocalLocationInfo local_info, bool skip_file_size_checks, + Promise promise); + + void check_partial_local_location(PartialLocalFileLocation partial, Promise promise); + private: struct Node { QueryId query_id_; diff --git a/td/telegram/files/FileManager.cpp b/td/telegram/files/FileManager.cpp index b0cf7591c..8d3bfee1d 100644 --- a/td/telegram/files/FileManager.cpp +++ b/td/telegram/files/FileManager.cpp @@ -943,13 +943,86 @@ Status FileManager::check_local_location(FileNodePtr node, bool skip_file_size_c status = check_partial_local_location(node->local_.partial()); } if (status.is_error()) { - send_closure(G()->download_manager(), &DownloadManager::remove_file_if_finished, node->main_file_id_); - node->drop_local_location(); - try_flush_node(node, "check_local_location"); + on_failed_check_local_location(node); } return status; } +void FileManager::on_failed_check_local_location(FileNodePtr node) { + send_closure(G()->download_manager(), &DownloadManager::remove_file_if_finished, node->main_file_id_); + node->drop_local_location(); + try_flush_node(node, "on_failed_check_local_location"); +} + +void FileManager::check_local_location_async(FileNodePtr node, bool skip_file_size_checks, Promise promise) { + if (node->local_.type() == LocalFileLocation::Type::Empty) { + return promise.set_value(Unit()); + } + + if (node->local_.type() == LocalFileLocation::Type::Full) { + send_closure(file_load_manager_, &FileLoadManager::check_full_local_location, + FullLocalLocationInfo{node->local_.full(), node->size_}, skip_file_size_checks, + PromiseCreator::lambda([actor_id = actor_id(this), file_id = node->main_file_id_, + checked_location = node->local_, + promise = std::move(promise)](Result result) mutable { + send_closure(actor_id, &FileManager::on_check_full_local_location, file_id, + std::move(checked_location), std::move(result), std::move(promise)); + })); + } else { + CHECK(node->local_.type() == LocalFileLocation::Type::Partial); + send_closure(file_load_manager_, &FileLoadManager::check_partial_local_location, node->local_.partial(), + PromiseCreator::lambda([actor_id = actor_id(this), file_id = node->main_file_id_, + checked_location = node->local_, + promise = std::move(promise)](Result result) mutable { + send_closure(actor_id, &FileManager::on_check_partial_local_location, file_id, + std::move(checked_location), std::move(result), std::move(promise)); + })); + } +} + +void FileManager::on_check_full_local_location(FileId file_id, LocalFileLocation checked_location, + Result r_info, Promise promise) { + TRY_STATUS_PROMISE(promise, G()->close_status()); + auto node = get_file_node(file_id); + CHECK(node); + if (node->local_ != checked_location) { + LOG(INFO) << "Full location changed while being checked; ignore check result"; + return promise.set_value(Unit()); + } + Status status; + if (r_info.is_error()) { + status = r_info.move_as_error(); + } else if (bad_paths_.count(r_info.ok().location_.path_) != 0) { + status = Status::Error(400, "Sending of internal database files is forbidden"); + } else if (r_info.ok().location_ != node->local_.full() || r_info.ok().size_ != node->size_) { + LOG(ERROR) << "Local location changed from " << node->local_.full() << " with size " << node->size_ << " to " + << r_info.ok().location_ << " with size " << r_info.ok().size_; + } + if (status.is_error()) { + on_failed_check_local_location(node); + promise.set_error(std::move(status)); + } else { + promise.set_value(Unit()); + } +} + +void FileManager::on_check_partial_local_location(FileId file_id, LocalFileLocation checked_location, + Result result, Promise promise) { + TRY_STATUS_PROMISE(promise, G()->close_status()); + auto node = get_file_node(file_id); + CHECK(node); + if (node->local_ != checked_location) { + LOG(INFO) << "Partial location changed while being checked; ignore check result"; + return promise.set_value(Unit()); + } + if (result.is_error()) { + on_failed_check_local_location(node); + promise.set_error(result.move_as_error()); + } else { + promise.set_value(Unit()); + } +} + bool FileManager::try_fix_partial_local_location(FileNodePtr node) { LOG(INFO) << "Trying to fix partial local location"; if (node->local_.type() != LocalFileLocation::Type::Partial) { @@ -1992,8 +2065,7 @@ void FileManager::get_content(FileId file_id, Promise promise) { if (!node) { return promise.set_error(Status::Error("Unknown file_id")); } - auto status = check_local_location(node, true); - status.ignore(); + check_local_location(node, true).ignore(); auto file_view = FileView(node); if (!file_view.has_local_location()) { diff --git a/td/telegram/files/FileManager.h b/td/telegram/files/FileManager.h index 8bb08bf41..836620330 100644 --- a/td/telegram/files/FileManager.h +++ b/td/telegram/files/FileManager.h @@ -617,6 +617,13 @@ class FileManager final : public FileLoadManager::Callback { FileId register_pmc_file_data(FileData &&data); Status check_local_location(FileNodePtr node, bool skip_file_size_checks); + void on_failed_check_local_location(FileNodePtr node); + void check_local_location_async(FileNodePtr node, bool skip_file_size_checks, Promise promise); + void on_check_full_local_location(FileId file_id, LocalFileLocation checked_location, + Result r_info, Promise promise); + void on_check_partial_local_location(FileId file_id, LocalFileLocation checked_location, Result result, + Promise promise); + static bool try_fix_partial_local_location(FileNodePtr node); void try_flush_node_full(FileNodePtr node, bool new_remote, bool new_local, bool new_generate, FileDbId other_pmc_id); void try_flush_node(FileNodePtr node, const char *source);