123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161 |
- // Copyright 2022 The Chromium Authors. All rights reserved.
- // Use of this source code is governed by a BSD-style license that can be
- // found in the LICENSE file.
- #include "components/optimization_guide/core/prediction_model_override.h"
- #include "base/files/file_util.h"
- #include "base/strings/string_util.h"
- #include "base/task/sequenced_task_runner.h"
- #include "base/task/thread_pool.h"
- #include "base/threading/thread_task_runner_handle.h"
- #include "components/optimization_guide/core/model_util.h"
- #include "components/optimization_guide/core/prediction_model_download_manager.h"
- #include "components/services/unzip/public/cpp/unzip.h"
- #if BUILDFLAG(IS_IOS)
- #include "components/services/unzip/in_process_unzipper.h" // nogncheck
- #else
- #include "components/services/unzip/content/unzip_service.h" // nogncheck
- #endif
- namespace optimization_guide {
- namespace {
- void OnModelOverrideProcessed(OnPredictionModelBuiltCallback callback,
- std::unique_ptr<proto::PredictionModel> model) {
- std::move(callback).Run(std::move(model));
- }
- std::unique_ptr<proto::PredictionModel> ProcessModelOverrideOnBGThread(
- proto::OptimizationTarget optimization_target,
- const base::FilePath& unzipped_dir_path) {
- // Unpack and verify model info file.
- base::FilePath model_info_path = unzipped_dir_path.Append(
- PredictionModelDownloadManager::ModelInfoFileName());
- std::string binary_model_info_pb;
- if (!base::ReadFileToString(model_info_path, &binary_model_info_pb)) {
- LOG(ERROR) << "Failed to read " << FilePathToString(model_info_path);
- return nullptr;
- }
- proto::ModelInfo model_info;
- if (!model_info.ParseFromString(binary_model_info_pb)) {
- LOG(ERROR) << "Failed to parse " << FilePathToString(model_info_path);
- return nullptr;
- }
- if (!model_info.has_version() || !model_info.has_optimization_target()) {
- LOG(ERROR) << FilePathToString(model_info_path)
- << "is invalid because it does not contain a version and/or "
- "optimization target";
- return nullptr;
- }
- for (int i = 0; i < model_info.additional_files_size(); i++) {
- proto::AdditionalModelFile* additional_file =
- model_info.mutable_additional_files(i);
- base::FilePath additional_file_basename =
- *StringToFilePath(additional_file->file_path());
- base::FilePath additional_file_absolute =
- unzipped_dir_path.Append(additional_file_basename);
- additional_file->set_file_path(FilePathToString(additional_file_absolute));
- }
- std::unique_ptr<proto::PredictionModel> model =
- std::make_unique<proto::PredictionModel>();
- *model->mutable_model_info() = model_info;
- model->mutable_model()->set_download_url(
- FilePathToString(unzipped_dir_path.Append(GetBaseFileNameForModels())));
- return model;
- }
- void OnModelOverrideUnzipped(proto::OptimizationTarget optimization_target,
- const base::FilePath& unzipped_dir_path,
- OnPredictionModelBuiltCallback callback,
- bool success) {
- if (!success) {
- LOG(ERROR) << FilePathToString(unzipped_dir_path) << "failed to unzip";
- std::move(callback).Run(nullptr);
- return;
- }
- base::ThreadPool::PostTaskAndReplyWithResult(
- FROM_HERE, {base::MayBlock(), base::TaskPriority::BEST_EFFORT},
- base::BindOnce(&ProcessModelOverrideOnBGThread, optimization_target,
- unzipped_dir_path),
- base::BindOnce(&OnModelOverrideProcessed, std::move(callback)));
- }
- void OnModelOverrideVerified(
- proto::OptimizationTarget optimization_target,
- const std::string& passed_crx_file_path,
- OnPredictionModelBuiltCallback callback,
- absl::optional<std::pair<base::FilePath, base::FilePath>> src_dst) {
- if (!src_dst) {
- LOG(ERROR) << passed_crx_file_path << " failed verification";
- std::move(callback).Run(nullptr);
- return;
- }
- #if BUILDFLAG(IS_IOS)
- auto unzipper = unzip::LaunchInProcessUnzipper();
- #else
- auto unzipper = unzip::LaunchUnzipper();
- #endif
- unzip::Unzip(std::move(unzipper), src_dst->first, src_dst->second,
- base::BindOnce(&OnModelOverrideUnzipped, optimization_target,
- src_dst->second, std::move(callback)));
- }
- } // namespace
- void BuildPredictionModelFromCommandLineForOptimizationTarget(
- proto::OptimizationTarget optimization_target,
- OnPredictionModelBuiltCallback callback) {
- absl::optional<std::pair<std::string, absl::optional<proto::Any>>>
- model_file_path_and_metadata =
- GetModelOverrideForOptimizationTarget(optimization_target);
- if (!model_file_path_and_metadata) {
- std::move(callback).Run(nullptr);
- return;
- }
- if (base::EndsWith(model_file_path_and_metadata->first, ".crx3")) {
- DVLOG(0) << "Attempting to parse the model override at "
- << model_file_path_and_metadata->first
- << " as a crx model package for "
- << GetStringNameForOptimizationTarget(optimization_target);
- if (model_file_path_and_metadata->second) {
- LOG(ERROR) << "Ignoring the metadata that was passed since a crx package "
- "was given";
- }
- base::ThreadPool::PostTaskAndReplyWithResult(
- FROM_HERE, {base::MayBlock(), base::TaskPriority::BEST_EFFORT},
- base::BindOnce(PredictionModelDownloadManager::VerifyDownload,
- *StringToFilePath(model_file_path_and_metadata->first),
- /*delete_file_on_error=*/false),
- base::BindOnce(&OnModelOverrideVerified, optimization_target,
- model_file_path_and_metadata->first,
- std::move(callback)));
- return;
- }
- std::unique_ptr<proto::PredictionModel> prediction_model =
- std::make_unique<proto::PredictionModel>();
- prediction_model->mutable_model_info()->set_optimization_target(
- optimization_target);
- prediction_model->mutable_model_info()->set_version(123);
- if (model_file_path_and_metadata->second) {
- *prediction_model->mutable_model_info()->mutable_model_metadata() =
- model_file_path_and_metadata->second.value();
- }
- prediction_model->mutable_model()->set_download_url(
- model_file_path_and_metadata->first);
- std::move(callback).Run(std::move(prediction_model));
- }
- } // namespace optimization_guide
|