diff --git a/td/telegram/files/FileLoader.cpp b/td/telegram/files/FileLoader.cpp index 9ebd6f8e0..3bdeaa6c7 100644 --- a/td/telegram/files/FileLoader.cpp +++ b/td/telegram/files/FileLoader.cpp @@ -54,7 +54,12 @@ void FileLoader::update_local_file_location(const LocalFileLocation &local) { return; } auto prefix_info = r_prefix_info.move_as_ok(); - parts_manager_.set_known_prefix(narrow_cast(prefix_info.size), prefix_info.is_ready); + auto status = parts_manager_.set_known_prefix(narrow_cast(prefix_info.size), prefix_info.is_ready); + if (status.is_error()) { + on_error(std::move(status)); + stop_flag_ = true; + return; + } loop(); } @@ -67,11 +72,12 @@ void FileLoader::start_up() { } auto file_info = r_file_info.ok(); auto size = file_info.size; + auto expected_size = std::max(size, file_info.expected_size); bool is_size_final = file_info.is_size_final; auto part_size = file_info.part_size; auto &ready_parts = file_info.ready_parts; auto use_part_count_limit = file_info.use_part_count_limit; - auto status = parts_manager_.init(size, is_size_final, part_size, ready_parts, use_part_count_limit); + auto status = parts_manager_.init(size, expected_size, is_size_final, part_size, ready_parts, use_part_count_limit); if (status.is_error()) { on_error(std::move(status)); stop_flag_ = true; diff --git a/td/telegram/files/FileLoader.h b/td/telegram/files/FileLoader.h index 7a4d35d70..10d97fab6 100644 --- a/td/telegram/files/FileLoader.h +++ b/td/telegram/files/FileLoader.h @@ -48,6 +48,7 @@ class FileLoader : public FileLoaderActor { }; struct FileInfo { int64 size; + int64 expected_size = 0; bool is_size_final; int32 part_size; std::vector ready_parts; diff --git a/td/telegram/files/FileManager.cpp b/td/telegram/files/FileManager.cpp index eb5aabaa5..8b6ba545a 100644 --- a/td/telegram/files/FileManager.cpp +++ b/td/telegram/files/FileManager.cpp @@ -639,6 +639,9 @@ Result FileManager::register_file(FileData data, FileLocationSource file if (status.is_error()) { LOG(WARNING) << "Invalid local location: " << status << " from " << source; data.local_ = LocalFileLocation(); + if (data.remote_.type() == RemoteFileLocation::Type::Partial) { + data.remote_ = {}; + } if (!has_remote && !has_generate) { return std::move(status); @@ -2178,7 +2181,7 @@ void FileManager::on_error_impl(FileNode *node, FileManager::Query::Type type, b << ". File type is " << file_type_name[static_cast(FileView(node).get_type())]; if (status.code() == 0) { // Remove partial locations - if (node->local_.type() == LocalFileLocation::Type::Partial) { + if (node->local_.type() == LocalFileLocation::Type::Partial && status.message() != "FILE_UPLOAD_RESTART") { LOG(INFO) << "Unlink file " << node->local_.partial().path_; unlink(node->local_.partial().path_).ignore(); node->set_local_location(LocalFileLocation(), 0); diff --git a/td/telegram/files/FileUploader.cpp b/td/telegram/files/FileUploader.cpp index 26f16fda3..635e5c893 100644 --- a/td/telegram/files/FileUploader.cpp +++ b/td/telegram/files/FileUploader.cpp @@ -223,7 +223,8 @@ Result> FileUploader::start_part(Part part, int32 p NetQueryPtr net_query; if (big_flag_) { - auto query = telegram_api::upload_saveBigFilePart(file_id_, part.id, part_count, std::move(bytes)); + auto query = + telegram_api::upload_saveBigFilePart(file_id_, part.id, local_is_ready_ ? part_count : -1, std::move(bytes)); net_query = G()->net_query_creator().create(create_storer(query), DcId::main(), NetQuery::Type::Upload, NetQuery::AuthFlag::On, NetQuery::GzipFlag::Off); } else { diff --git a/td/telegram/files/PartsManager.cpp b/td/telegram/files/PartsManager.cpp index eae9ddd15..47e85dc2a 100644 --- a/td/telegram/files/PartsManager.cpp +++ b/td/telegram/files/PartsManager.cpp @@ -16,6 +16,14 @@ namespace td { /*** PartsManager ***/ + +namespace { +int64 calc_parts_count(int64 size, int64 part_size) { + CHECK(part_size != 0); + return (size + part_size - 1) / part_size; +} +} // namespace + Status PartsManager::init_known_prefix(int64 known_prefix, size_t part_size, const std::vector &ready_parts) { known_prefix_flag_ = true; known_prefix_size_ = known_prefix; @@ -32,6 +40,14 @@ Status PartsManager::init_no_size(size_t part_size, const std::vector &read part_size_ = part_size; } else { part_size_ = 32 * (1 << 10); + while (use_part_count_limit_ && calc_parts_count(expected_size_, part_size_) > MAX_PART_COUNT) { + part_size_ *= 2; + CHECK(part_size_ <= MAX_PART_SIZE); + } + // just in case if expected_size_ is wrong + if (part_size_ < MAX_PART_SIZE) { + part_size_ *= 2; + } } part_count_ = 0; if (known_prefix_flag_) { @@ -44,40 +60,42 @@ Status PartsManager::init_no_size(size_t part_size, const std::vector &read return Status::OK(); } -Status PartsManager::init(int64 size, bool is_size_final, size_t part_size, const std::vector &ready_parts, - bool use_part_count_limit) { +Status PartsManager::init(int64 size, int64 expected_size, bool is_size_final, size_t part_size, + const std::vector &ready_parts, bool use_part_count_limit) { + CHECK(expected_size >= size); use_part_count_limit_ = use_part_count_limit; + expected_size_ = expected_size; + if (expected_size_ > MAX_FILE_SIZE) { + return Status::Error("Too big file"); + } if (!is_size_final) { return init_known_prefix(size, part_size, ready_parts); } if (size == 0) { return init_no_size(part_size, ready_parts); } - if (size > MAX_FILE_SIZE) { - return Status::Error("Too big file"); - } CHECK(size > 0) << tag("size", size); unknown_size_flag_ = false; size_ = size; if (part_size != 0) { part_size_ = part_size; - if (use_part_count_limit_ && (size_ + part_size_ - 1) / part_size_ > MAX_PART_COUNT) { + if (use_part_count_limit_ && calc_parts_count(expected_size_, part_size_) > MAX_PART_COUNT) { return Status::Error("FILE_UPLOAD_RESTART"); } } else { // TODO choose part_size_ depending on size part_size_ = 64 * (1 << 10); - while (use_part_count_limit && (size_ + part_size_ - 1) / part_size_ > MAX_PART_COUNT) { + while (use_part_count_limit && calc_parts_count(expected_size_, part_size_) > MAX_PART_COUNT) { part_size_ *= 2; CHECK(part_size_ <= MAX_PART_SIZE); } } CHECK(1 <= size_) << tag("size_", size_); - CHECK(!use_part_count_limit || (size_ + part_size_ - 1) / part_size_ <= MAX_PART_COUNT) - << tag("size_", size_) << tag("is_size_final", is_size_final) << tag("part_size_", part_size_) - << tag("ready_parts", ready_parts.size()); - part_count_ = static_cast((size + part_size_ - 1) / part_size_); + CHECK(!use_part_count_limit || calc_parts_count(expected_size_, part_size_) <= MAX_PART_COUNT) + << tag("size_", size_) << tag("expected_size", size_) << tag("is_size_final", is_size_final) + << tag("part_size_", part_size_) << tag("ready_parts", ready_parts.size()); + part_count_ = static_cast(calc_parts_count(size_, part_size_)); init_common(ready_parts); return Status::OK(); @@ -147,14 +165,15 @@ Result PartsManager::start_part() { return get_part(id); } -void PartsManager::set_known_prefix(size_t size, bool is_ready) { +Status PartsManager::set_known_prefix(size_t size, bool is_ready) { CHECK(known_prefix_flag_); CHECK(size >= static_cast(known_prefix_size_)); known_prefix_size_ = narrow_cast(size); + expected_size_ = std::max(known_prefix_size_, expected_size_); CHECK(static_cast(part_count_) == part_status_.size()); if (is_ready) { - part_count_ = static_cast((size + part_size_ - 1) / part_size_); + part_count_ = static_cast(calc_parts_count(size, part_size_)); size_ = narrow_cast(size); unknown_size_flag_ = false; @@ -164,6 +183,10 @@ void PartsManager::set_known_prefix(size_t size, bool is_ready) { CHECK(static_cast(part_count_) >= part_status_.size()) << size << " " << is_ready << " " << part_count_ << " " << part_size_ << " " << part_status_.size(); part_status_.resize(part_count_); + if (use_part_count_limit_ && calc_parts_count(expected_size_, part_size_) > MAX_PART_COUNT) { + return Status::Error("FILE_UPLOAD_RESTART"); + } + return Status::OK(); } Status PartsManager::on_part_ok(int32 id, size_t part_size, size_t actual_size) { diff --git a/td/telegram/files/PartsManager.h b/td/telegram/files/PartsManager.h index 91ba116fd..26c31d517 100644 --- a/td/telegram/files/PartsManager.h +++ b/td/telegram/files/PartsManager.h @@ -20,8 +20,8 @@ struct Part { class PartsManager { public: - Status init(int64 size, bool is_size_final, size_t part_size, const std::vector &ready_parts, - bool use_part_count_limit) TD_WARN_UNUSED_RESULT; + Status init(int64 size, int64 expected_size, bool is_size_final, size_t part_size, + const std::vector &ready_parts, bool use_part_count_limit) TD_WARN_UNUSED_RESULT; bool ready(); bool unchecked_ready(); Status finish() TD_WARN_UNUSED_RESULT; @@ -30,7 +30,7 @@ class PartsManager { Result start_part() TD_WARN_UNUSED_RESULT; Status on_part_ok(int32 id, size_t part_size, size_t actual_size) TD_WARN_UNUSED_RESULT; void on_part_failed(int32 id); - void set_known_prefix(size_t size, bool is_ready); + Status set_known_prefix(size_t size, bool is_ready); void set_need_check(); void set_checked_prefix_size(int64 size); @@ -58,6 +58,7 @@ class PartsManager { int64 known_prefix_size_; int64 size_; + int64 expected_size_; int64 min_size_; int64 max_size_; bool unknown_size_flag_;