components/browsing_topics/annotator_fuzzer.cc

   1 // Copyright 2023 The Chromium Authors
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include <string>
   6
   7 #include <fuzzer/FuzzedDataProvider.h>
   8
   9 #include "base/files/file_util.h"
  10 #include "base/path_service.h"
  11 #include "base/task/sequenced_task_runner.h"
  12 #include "base/task/thread_pool.h"
  13 #include "base/test/scoped_feature_list.h"
  14 #include "base/test/task_environment.h"
  15 #include "components/browsing_topics/annotator_impl.h"
  16 #include "components/optimization_guide/core/optimization_guide_features.h"
  17 #include "components/optimization_guide/core/optimization_guide_model_provider.h"
  18 #include "components/optimization_guide/core/test_model_info_builder.h"
  19 #include "components/optimization_guide/core/test_optimization_guide_model_provider.h"
  20 #include "components/optimization_guide/proto/models.pb.h"
  21 #include "components/optimization_guide/proto/page_topics_model_metadata.pb.h"
  22
  23 // Sends a fully working model and configuration to the calling observer.
  24 class ModelProvider
  25     : public optimization_guide::TestOptimizationGuideModelProvider {
  26  public:
  27   ModelProvider() = default;
  28   ~ModelProvider() override = default;
  29
  30   // optimization_guide::TestOptimizationGuideModelProvider:
  31   void AddObserverForOptimizationTargetModel(
  32       optimization_guide::proto::OptimizationTarget optimization_target,
  33       const absl::optional<optimization_guide::proto::Any>& model_metadata,
  34       optimization_guide::OptimizationTargetModelObserver* observer) override {
  35     optimization_guide::proto::Any any_metadata;
  36     any_metadata.set_type_url(
  37         "type.googleapis.com/com.foo.PageTopicsModelMetadata");
  38     optimization_guide::proto::PageTopicsModelMetadata
  39         page_topics_model_metadata;
  40     page_topics_model_metadata.set_version(123);
  41     page_topics_model_metadata.add_supported_output(
  42         optimization_guide::proto::PAGE_TOPICS_SUPPORTED_OUTPUT_CATEGORIES);
  43     auto* output_params =
  44         page_topics_model_metadata.mutable_output_postprocessing_params();
  45     auto* category_params = output_params->mutable_category_params();
  46     category_params->set_max_categories(5);
  47     category_params->set_min_none_weight(0.8);
  48     category_params->set_min_category_weight(0.1);
  49     category_params->set_min_normalized_weight_within_top_n(0.1);
  50     page_topics_model_metadata.SerializeToString(any_metadata.mutable_value());
  51
  52     base::FilePath source_root_dir;
  53     base::PathService::Get(base::DIR_SRC_TEST_DATA_ROOT, &source_root_dir);
  54     base::FilePath model_file_path =
  55         source_root_dir.AppendASCII("components")
  56             .AppendASCII("test")
  57             .AppendASCII("data")
  58             .AppendASCII("browsing_topics")
  59             .AppendASCII("golden_data_model.tflite");
  60     std::unique_ptr<optimization_guide::ModelInfo> model_info =
  61         optimization_guide::TestModelInfoBuilder()
  62             .SetModelFilePath(model_file_path)
  63             .SetModelMetadata(any_metadata)
  64             .Build();
  65
  66     observer->OnModelUpdated(
  67         optimization_guide::proto::OPTIMIZATION_TARGET_PAGE_TOPICS_V2,
  68         *model_info);
  69   }
  70 };
  71
  72 // An AnnotatorImpl that never unloads the model, thus keeping the executions
  73 // per second high.
  74 class TestAnnotator : public browsing_topics::AnnotatorImpl {
  75  public:
  76   TestAnnotator(
  77       optimization_guide::OptimizationGuideModelProvider* model_provider,
  78       scoped_refptr<base::SequencedTaskRunner> background_task_runner,
  79       const absl::optional<optimization_guide::proto::Any>& model_metadata)
  80       : browsing_topics::AnnotatorImpl(model_provider,
  81                                        std::move(background_task_runner),
  82                                        model_metadata) {}
  83
  84  protected:
  85   // AnnotatorImpl:
  86   void UnloadModel() override {
  87     // Do nothing so that the model stays loaded.
  88   }
  89 };
  90
  91 // Static test fixture to maintain the state needed to run repeated fuzz tests.
  92 class AnnotatorFuzzerTest {
  93  public:
  94   explicit AnnotatorFuzzerTest()
  95       : annotator_(std::make_unique<TestAnnotator>(
  96             &model_provider_,
  97             task_environment_.GetMainThreadTaskRunner(),
  98             absl::nullopt)) {
  99     scoped_feature_list_.InitAndDisableFeature(
 100         optimization_guide::features::kPreventLongRunningPredictionModels);
 101   }
 102   ~AnnotatorFuzzerTest() {
 103     annotator_.reset();
 104     // To prevent ASAN issues, ensure the ModelExecutor gets destroyed since
 105     // that is done using a DeleteSoon PostTask.
 106     task_environment_.RunUntilIdle();
 107   }
 108
 109   browsing_topics::Annotator* annotator() { return annotator_.get(); }
 110
 111   void RunUntilIdle() { task_environment_.RunUntilIdle(); }
 112
 113  private:
 114   base::test::TaskEnvironment task_environment_;
 115   base::test::ScopedFeatureList scoped_feature_list_;
 116   ModelProvider model_provider_;
 117   std::unique_ptr<browsing_topics::AnnotatorImpl> annotator_;
 118 };
 119
 120 extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
 121   static AnnotatorFuzzerTest test;
 122   static bool had_successful_run = false;
 123
 124   std::string input(reinterpret_cast<const char*>(data), size);
 125
 126   base::RunLoop run_loop;
 127   test.annotator()->BatchAnnotate(
 128       base::BindOnce(
 129           [](base::RunLoop* run_loop,
 130              const std::vector<browsing_topics::Annotation>& annotations) {
 131             if (!annotations[0].topics.empty() && !had_successful_run) {
 132               had_successful_run = true;
 133               // Print a single debug message so that its obvious things are
 134               // working (or able to at least once) when running locally.
 135               LOG(INFO) << "Congrats! Got a successful model execution. This "
 136                            "message will not be printed again.";
 137             }
 138
 139             run_loop->Quit();
 140           },
 141           &run_loop),
 142       {input});
 143   run_loop.Run();
 144
 145   // The model executor does some PostTask'ing to manage its state. While these
 146   // tasks are not important for fuzzing, we don't want to queue up a ton of
 147   // them.
 148   test.RunUntilIdle();
 149
 150   return 0;
 151 }