Move the Browsing Topics Annotator fuzz test
Moves the fuzz test into the browsing_topics component.
Verified this is fully working locally.
Bug: b/278162907
Change-Id: I4c26d117436316939145cc3c136f7a51f1df727a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4495631
Reviewed-by: Yao Xiao <[email protected]>
Commit-Queue: Robert Ogden <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1139126}
diff --git a/components/browsing_topics/BUILD.gn b/components/browsing_topics/BUILD.gn
index 8ea72a76..78f94e31a 100644
--- a/components/browsing_topics/BUILD.gn
+++ b/components/browsing_topics/BUILD.gn
@@ -1,6 +1,7 @@
# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
+import("//testing/libfuzzer/fuzzer_test.gni")
source_set("browsing_topics") {
sources = [
@@ -106,3 +107,16 @@
"//third_party/zlib/google:compression_utils",
]
}
+
+fuzzer_test("browsing_topics_annotator_fuzzer") {
+ sources = [ "annotator_fuzzer.cc" ]
+ deps = [
+ ":browsing_topics",
+ "//base:base",
+ "//base/test:test_support",
+ "//components/optimization_guide/core:core",
+ "//components/optimization_guide/core:features",
+ "//components/optimization_guide/core:test_support",
+ "//components/optimization_guide/proto:optimization_guide_proto",
+ ]
+}
diff --git a/components/browsing_topics/annotator_fuzzer.cc b/components/browsing_topics/annotator_fuzzer.cc
new file mode 100644
index 0000000..45e13e9
--- /dev/null
+++ b/components/browsing_topics/annotator_fuzzer.cc
@@ -0,0 +1,146 @@
+// Copyright 2023 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <string>
+
+#include <fuzzer/FuzzedDataProvider.h>
+
+#include "base/files/file_util.h"
+#include "base/path_service.h"
+#include "base/task/sequenced_task_runner.h"
+#include "base/task/thread_pool.h"
+#include "base/test/scoped_feature_list.h"
+#include "base/test/task_environment.h"
+#include "components/browsing_topics/annotator_impl.h"
+#include "components/optimization_guide/core/optimization_guide_features.h"
+#include "components/optimization_guide/core/optimization_guide_model_provider.h"
+#include "components/optimization_guide/core/test_model_info_builder.h"
+#include "components/optimization_guide/core/test_optimization_guide_model_provider.h"
+#include "components/optimization_guide/proto/models.pb.h"
+#include "components/optimization_guide/proto/page_topics_model_metadata.pb.h"
+
+// Sends a fully working model and configuration to the calling observer.
+class ModelProvider
+ : public optimization_guide::TestOptimizationGuideModelProvider {
+ public:
+ ModelProvider() = default;
+ ~ModelProvider() override = default;
+
+ // optimization_guide::TestOptimizationGuideModelProvider:
+ void AddObserverForOptimizationTargetModel(
+ optimization_guide::proto::OptimizationTarget optimization_target,
+ const absl::optional<optimization_guide::proto::Any>& model_metadata,
+ optimization_guide::OptimizationTargetModelObserver* observer) override {
+ optimization_guide::proto::Any any_metadata;
+ any_metadata.set_type_url(
+ "type.googleapis.com/com.foo.PageTopicsModelMetadata");
+ optimization_guide::proto::PageTopicsModelMetadata
+ page_topics_model_metadata;
+ page_topics_model_metadata.set_version(123);
+ page_topics_model_metadata.add_supported_output(
+ optimization_guide::proto::PAGE_TOPICS_SUPPORTED_OUTPUT_CATEGORIES);
+ auto* output_params =
+ page_topics_model_metadata.mutable_output_postprocessing_params();
+ auto* category_params = output_params->mutable_category_params();
+ category_params->set_max_categories(5);
+ category_params->set_min_none_weight(0.8);
+ category_params->set_min_category_weight(0.1);
+ category_params->set_min_normalized_weight_within_top_n(0.1);
+ page_topics_model_metadata.SerializeToString(any_metadata.mutable_value());
+
+ base::FilePath source_root_dir;
+ base::PathService::Get(base::DIR_SOURCE_ROOT, &source_root_dir);
+ base::FilePath model_file_path =
+ source_root_dir.AppendASCII("components")
+ .AppendASCII("test")
+ .AppendASCII("data")
+ .AppendASCII("browsing_topics")
+ .AppendASCII("golden_data_model.tflite");
+ std::unique_ptr<optimization_guide::ModelInfo> model_info =
+ optimization_guide::TestModelInfoBuilder()
+ .SetModelFilePath(model_file_path)
+ .SetModelMetadata(any_metadata)
+ .Build();
+
+ observer->OnModelUpdated(
+ optimization_guide::proto::OPTIMIZATION_TARGET_PAGE_TOPICS_V2,
+ *model_info);
+ }
+};
+
+// An AnnotatorImpl that never unloads the model, thus keeping the executions
+// per second high.
+class TestAnnotator : public browsing_topics::AnnotatorImpl {
+ public:
+ TestAnnotator(
+ optimization_guide::OptimizationGuideModelProvider* model_provider,
+ scoped_refptr<base::SequencedTaskRunner> background_task_runner,
+ const absl::optional<optimization_guide::proto::Any>& model_metadata)
+ : browsing_topics::AnnotatorImpl(model_provider,
+ std::move(background_task_runner),
+ model_metadata) {}
+
+ protected:
+ // AnnotatorImpl:
+ void UnloadModel() override {
+ // Do nothing so that the model stays loaded.
+ }
+};
+
+// Static test fixture to maintain the state needed to run repeated fuzz tests.
+class AnnotatorFuzzerTest {
+ public:
+ explicit AnnotatorFuzzerTest()
+ : annotator_(std::make_unique<TestAnnotator>(
+ &model_provider_,
+ task_environment_.GetMainThreadTaskRunner(),
+ absl::nullopt)) {
+ scoped_feature_list_.InitAndDisableFeature(
+ optimization_guide::features::kPreventLongRunningPredictionModels);
+ }
+ ~AnnotatorFuzzerTest() { task_environment_.RunUntilIdle(); }
+
+ browsing_topics::Annotator* annotator() { return annotator_.get(); }
+
+ void RunUntilIdle() { task_environment_.RunUntilIdle(); }
+
+ private:
+ base::test::TaskEnvironment task_environment_;
+ base::test::ScopedFeatureList scoped_feature_list_;
+ ModelProvider model_provider_;
+ std::unique_ptr<browsing_topics::AnnotatorImpl> annotator_;
+};
+
+extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
+ static AnnotatorFuzzerTest test;
+ static bool had_successful_run = false;
+
+ std::string input(reinterpret_cast<const char*>(data), size);
+
+ base::RunLoop run_loop;
+ test.annotator()->BatchAnnotate(
+ base::BindOnce(
+ [](base::RunLoop* run_loop,
+ const std::vector<browsing_topics::Annotation>& annotations) {
+ if (!annotations[0].topics.empty() && !had_successful_run) {
+ had_successful_run = true;
+ // Print a single debug message so that its obvious things are
+ // working (or able to at least once) when running locally.
+ LOG(INFO) << "Congrats! Got a successful model execution. This "
+ "message will not be printed again.";
+ }
+
+ run_loop->Quit();
+ },
+ &run_loop),
+ {input});
+ run_loop.Run();
+
+ // The model executor does some PostTask'ing to manage its state. While these
+ // tasks are not important for fuzzing, we don't want to queue up a ton of
+ // them.
+ test.RunUntilIdle();
+
+ return 0;
+}
diff --git a/components/browsing_topics/annotator_impl.h b/components/browsing_topics/annotator_impl.h
index d958528..87d63b4 100644
--- a/components/browsing_topics/annotator_impl.h
+++ b/components/browsing_topics/annotator_impl.h
@@ -77,6 +77,10 @@
absl::optional<std::vector<int32_t>> ExtractCategoriesFromModelOutput(
const std::vector<tflite::task::core::Category>& model_output) const;
+ protected:
+ // optimization_guide::BertModelHandler:
+ void UnloadModel() override;
+
private:
// Sets the |override_list_| after it was loaded on a background thread and
// calls |StartBatchAnnotate|.
@@ -104,9 +108,6 @@
BatchAnnotationCallback callback,
std::unique_ptr<std::vector<Annotation>> annotations_ptr);
- // optimization_guide::BertModelHandler:
- void UnloadModel() override;
-
// Sets |annotation.topics| from the output of the model, calling
// |ExtractCategoriesFromModelOutput| in the process.
void PostprocessCategoriesToBatchAnnotationResult(
diff --git a/components/optimization_guide/core/BUILD.gn b/components/optimization_guide/core/BUILD.gn
index 494611d..a2d21f6 100644
--- a/components/optimization_guide/core/BUILD.gn
+++ b/components/optimization_guide/core/BUILD.gn
@@ -488,32 +488,6 @@
}
if (build_with_tflite_lib) {
- template("bert_model_fuzzer") {
- forward_variables_from(invoker, "*")
-
- fuzzer_test("optimization_guide_${target_name}_fuzzer") {
- sources = [ "bert_model_fuzzer.cc" ]
- defines = [
- "OPTIMIZATION_TARGET=${optimization_target}",
- "MODEL_FILE_BASENAME=\"${model_file_basename}\"",
- ]
- deps = [
- ":core",
- ":features",
- ":machine_learning",
- ":model_executor",
- "//base:base",
- "//base/test:test_support",
- "//components/optimization_guide/proto:optimization_guide_proto",
- ]
- }
- }
-
- bert_model_fuzzer("page_topics") {
- optimization_target = "OPTIMIZATION_TARGET_PAGE_TOPICS_V2"
- model_file_basename = "bert_page_topics_model.tflite"
- }
-
fuzzer_test("optimization_guide_page_visibility_model_fuzzer") {
sources = [ "page_visibility_model_fuzzer.cc" ]
deps = [
diff --git a/components/optimization_guide/core/bert_model_fuzzer.cc b/components/optimization_guide/core/bert_model_fuzzer.cc
deleted file mode 100644
index 7aa32ef..0000000
--- a/components/optimization_guide/core/bert_model_fuzzer.cc
+++ /dev/null
@@ -1,81 +0,0 @@
-// Copyright 2022 The Chromium Authors
-// Use of this source code is governed by a BSD-style license that can be
-// found in the LICENSE file.
-
-#include <string>
-
-#include <fuzzer/FuzzedDataProvider.h>
-
-#include "base/path_service.h"
-#include "base/test/task_environment.h"
-#include "components/optimization_guide/core/bert_model_executor.h"
-
-// This fuzzer relies on two macros so that the test code can be easily reused
-// across multiple model inputs. Make sure these are defined by the build rule.
-//
-// OPTIMIZATION_TARGET should be something like
-// OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD.
-//
-// MODEL_FILE_BASENAME should be something like "model.tflite", quotes included!
-// Given model file base names use the OptGuide test data dir in //components.
-class BertModelExecutorFuzzer {
- public:
- BertModelExecutorFuzzer()
- : model_executor_(optimization_guide::BertModelExecutor(
- optimization_guide::proto::OPTIMIZATION_TARGET)) {
- model_executor_.InitializeAndMoveToExecutionThread(
- // This is an arbitrarily long time since we don't need to test the
- // timeout behavior here, libfuzzer will take care of hangs.
- /*model_inference_timeout=*/base::Minutes(60),
- optimization_guide::proto::OPTIMIZATION_TARGET,
- /*execution_task_runner=*/task_environment_.GetMainThreadTaskRunner(),
- /*reply_task_runner=*/task_environment_.GetMainThreadTaskRunner());
-
- base::FilePath source_root_dir;
- base::PathService::Get(base::DIR_SOURCE_ROOT, &source_root_dir);
- base::FilePath model_file_path = source_root_dir.AppendASCII("components")
- .AppendASCII("test")
- .AppendASCII("data")
- .AppendASCII("optimization_guide")
- .AppendASCII(MODEL_FILE_BASENAME);
-
- model_executor_.UpdateModelFile(model_file_path);
- model_executor_.SetShouldUnloadModelOnComplete(false);
- }
- ~BertModelExecutorFuzzer() { task_environment_.RunUntilIdle(); }
-
- optimization_guide::BertModelExecutor* executor() { return &model_executor_; }
-
- void RunUntilIdle() { task_environment_.RunUntilIdle(); }
-
- private:
- base::test::TaskEnvironment task_environment_;
- optimization_guide::BertModelExecutor model_executor_;
-};
-
-extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
- static BertModelExecutorFuzzer test;
- static bool had_successful_run = false;
-
- std::string input(reinterpret_cast<const char*>(data), size);
- test.executor()->SendForExecution(
- base::BindOnce(
- [](const absl::optional<std::vector<tflite::task::core::Category>>&
- output) {
- if (output && !had_successful_run) {
- had_successful_run = true;
- // Print a single debug message so that its obvious things are
- // working (or able to at least once) when running locally.
- LOG(INFO) << "Congrats! Got a successful model execution. This "
- "message will not be printed again.";
- }
- }),
- base::TimeTicks(), input);
-
- // The model executor does some PostTask'ing to manage its state. While these
- // tasks are not important for fuzzing, we don't want to queue up a ton of
- // them.
- test.RunUntilIdle();
-
- return 0;
-}
\ No newline at end of file
