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.
7 #include <fuzzer/FuzzedDataProvider.h>
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"
23 // Sends a fully working model and configuration to the calling observer.
25 : public optimization_guide::TestOptimizationGuideModelProvider {
27 ModelProvider() = default;
28 ~ModelProvider() override = default;
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);
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());
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")
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)
66 observer->OnModelUpdated(
67 optimization_guide::proto::OPTIMIZATION_TARGET_PAGE_TOPICS_V2,
72 // An AnnotatorImpl that never unloads the model, thus keeping the executions
74 class TestAnnotator : public browsing_topics::AnnotatorImpl {
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),
86 void UnloadModel() override {
87 // Do nothing so that the model stays loaded.
91 // Static test fixture to maintain the state needed to run repeated fuzz tests.
92 class AnnotatorFuzzerTest {
94 explicit AnnotatorFuzzerTest()
95 : annotator_(std::make_unique<TestAnnotator>(
97 task_environment_.GetMainThreadTaskRunner(),
99 scoped_feature_list_.InitAndDisableFeature(
100 optimization_guide::features::kPreventLongRunningPredictionModels);
102 ~AnnotatorFuzzerTest() {
104 // To prevent ASAN issues, ensure the ModelExecutor gets destroyed since
105 // that is done using a DeleteSoon PostTask.
106 task_environment_.RunUntilIdle();
109 browsing_topics::Annotator* annotator() { return annotator_.get(); }
111 void RunUntilIdle() { task_environment_.RunUntilIdle(); }
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_;
120 extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
121 static AnnotatorFuzzerTest test;
122 static bool had_successful_run = false;
124 std::string input(reinterpret_cast<const char*>(data), size);
126 base::RunLoop run_loop;
127 test.annotator()->BatchAnnotate(
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.";
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