Add support for rejecting invalid inputs

README.md

@@ -88,6 +88,38 @@ DEFINE_PROTO_FUZZER(const MyMessageType& input) {
 
 Please see [libfuzzer_example.cc](/examples/libfuzzer/libfuzzer_example.cc) as an example.
 
+### Reject unwanted inputs
+To reject adding unwanted inputs in corpus:
+
+Either define `LIBPROTOBUFMUTATOR_ENABLE_RETURN_CODE` macro in your code and return `-1` for inputs to be rejected:
+
+```
+#define LIBPROTOBUFMUTATOR_ENABLE_RETURN_CODE
+#include "src/libfuzzer/libfuzzer_macro.h"
+
+DEFINE_PROTO_FUZZER(const MyMessageType& input) {
+  // Code which needs to be fuzzed.
+  ConsumeMyMessageType(input);
+  // If 'input' is to be rejected
+  return -1;
+}
+```
+
+or use another macro `DEFINE_PROTO_FUZZER_RET` supporting return code:
+
+```
+#include "src/libfuzzer/libfuzzer_macro.h"
+
+DEFINE_PROTO_FUZZER_RET(const MyMessageType& input) {
+  // Code which needs to be fuzzed.
+  ConsumeMyMessageType(input);
+  // If 'input' is to be rejected
+  return -1;
+}
+```
+
+Refer https://llvm.org/docs/LibFuzzer.html#rejecting-unwanted-inputs for details.
+
 ### Mutation post-processing (experimental)
 Sometimes it's necessary to keep particular values in some fields without which the proto
 is going to be rejected by fuzzed code. E.g. code may expect consistency between some fields

src/libfuzzer/CMakeLists.txt

@@ -29,14 +29,20 @@ install(TARGETS protobuf-mutator-libfuzzer
         INCLUDES DESTINATION include/libprotobuf-mutator include/libprotobuf-mutator/src)
 
 if (LIB_PROTO_MUTATOR_TESTING)
-  add_executable(libfuzzer_test
-                 libfuzzer_test.cc)
-  target_link_libraries(libfuzzer_test
-                        protobuf-mutator
-                        protobuf-mutator-libfuzzer
-                        mutator-test-proto
-                        ${GTEST_BOTH_LIBRARIES}
-                        ${CMAKE_THREAD_LIBS_INIT})
-  add_test(test.protobuf_libfuzzer_test libfuzzer_test --gtest_color=yes AUTO)
-  add_dependencies(check libfuzzer_test)
+  list(APPEND TARGETS libfuzzer_test libfuzzer_return_enabled_test libfuzzer_return_overload_test)
+  list(APPEND FLAGS "" LIBPROTOBUFMUTATOR_ENABLE_RETURN_CODE TEST_RETURN_OVERLOAD)
+
+  foreach(TARGET FLAG IN ZIP_LISTS TARGETS FLAGS)
+      add_executable(${TARGET}
+                     libfuzzer_test.cc)
+      target_compile_definitions(${TARGET} PRIVATE ${FLAG})
+      target_link_libraries(${TARGET}
+                            protobuf-mutator
+                            protobuf-mutator-libfuzzer
+                            mutator-test-proto
+                            ${GTEST_BOTH_LIBRARIES}
+                            ${CMAKE_THREAD_LIBS_INIT})
+      add_test(test.protobuf_${TARGET} ${TARGET} --gtest_color=yes AUTO)
+      add_dependencies(check ${TARGET})
+  endforeach()
 endif()

src/libfuzzer/libfuzzer_macro.h

@@ -26,13 +26,23 @@
 // Defines custom mutator, crossover and test functions using default
 // serialization format. Default is text.
 #define DEFINE_PROTO_FUZZER(arg) DEFINE_TEXT_PROTO_FUZZER(arg)
+// Same as 'DEFINE_PROTO_FUZZER' but supports a return code to reject invalid
+// inputs.
+#define DEFINE_PROTO_FUZZER_RET(arg) DEFINE_TEXT_PROTO_FUZZER_RET(arg)
 // Defines custom mutator, crossover and test functions using text
 // serialization. This format is more convenient to read.
 #define DEFINE_TEXT_PROTO_FUZZER(arg) DEFINE_PROTO_FUZZER_IMPL(false, arg)
+// Same as 'DEFINE_TEXT_PROTO_FUZZER' but supports a return code to reject invalid
+// inputs.
+#define DEFINE_TEXT_PROTO_FUZZER_RET(arg) DEFINE_PROTO_FUZZER_RET_IMPL(false, arg)
 // Defines custom mutator, crossover and test functions using binary
 // serialization. This makes mutations faster. However often test function is
 // significantly slower than mutator, so fuzzing rate may stay unchanged.
 #define DEFINE_BINARY_PROTO_FUZZER(arg) DEFINE_PROTO_FUZZER_IMPL(true, arg)
+// Same as 'DEFINE_BINARY_PROTO_FUZZER' but supports a return code to reject invalid
+// inputs.
+#define DEFINE_BINARY_PROTO_FUZZER_RET(arg) \
+  DEFINE_PROTO_FUZZER_RET_IMPL(true, arg)
 
 // Registers the callback as a potential mutation performed on the parent
 // message of a field. This must be called inside an initialization code block.
@@ -78,20 +88,45 @@
     return 0;                                                               \
   }
 
+#define DEFINE_TEST_ONE_PROTO_INPUT_RET_IMPL(use_binary, Proto)             \
+  extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) { \
+    using protobuf_mutator::libfuzzer::LoadProtoInput;                      \
+    Proto input;                                                            \
+    if (LoadProtoInput(use_binary, data, size, &input))                     \
+      return TestOneProtoInput(input);                                      \
+    return -1;                                                              \
+  }
+
 #define DEFINE_POST_PROCESS_PROTO_MUTATION_IMPL(Proto) \
   using PostProcessorRegistration =                    \
       protobuf_mutator::libfuzzer::PostProcessorRegistration<Proto>;
 
-#define DEFINE_PROTO_FUZZER_IMPL(use_binary, arg)                 \
-  static void TestOneProtoInput(arg);                             \
-  using FuzzerProtoType =                                         \
-      protobuf_mutator::libfuzzer::macro_internal::GetFirstParam< \
-          decltype(&TestOneProtoInput)>::type;                    \
-  DEFINE_CUSTOM_PROTO_MUTATOR_IMPL(use_binary, FuzzerProtoType)   \
-  DEFINE_CUSTOM_PROTO_CROSSOVER_IMPL(use_binary, FuzzerProtoType) \
-  DEFINE_TEST_ONE_PROTO_INPUT_IMPL(use_binary, FuzzerProtoType)   \
-  DEFINE_POST_PROCESS_PROTO_MUTATION_IMPL(FuzzerProtoType)        \
+#ifdef LIBPROTOBUFMUTATOR_ENABLE_RETURN_CODE
+#define DEFINE_PROTO_FUZZER_IMPL(use_binary, arg) \
+  DEFINE_PROTO_FUZZER_RET_IMPL(use_binary, arg)
+#else
+#define DEFINE_PROTO_FUZZER_IMPL(use_binary, arg)                          \
+  static void TestOneProtoInput(arg);                                      \
+  using FuzzerProtoType =                                                  \
+      protobuf_mutator::libfuzzer::macro_internal::GetFirstParam<decltype( \
+          &TestOneProtoInput)>::type;                                      \
+  DEFINE_CUSTOM_PROTO_MUTATOR_IMPL(use_binary, FuzzerProtoType)            \
+  DEFINE_CUSTOM_PROTO_CROSSOVER_IMPL(use_binary, FuzzerProtoType)          \
+  DEFINE_TEST_ONE_PROTO_INPUT_IMPL(use_binary, FuzzerProtoType)            \
+  DEFINE_POST_PROCESS_PROTO_MUTATION_IMPL(FuzzerProtoType)                 \
   static void TestOneProtoInput(arg)
+#endif
+
+#define DEFINE_PROTO_FUZZER_RET_IMPL(use_binary, arg)                      \
+  static int TestOneProtoInput(arg);                                       \
+  using FuzzerProtoType =                                                  \
+      protobuf_mutator::libfuzzer::macro_internal::GetFirstParam<decltype( \
+          &TestOneProtoInput)>::type;                                      \
+  DEFINE_CUSTOM_PROTO_MUTATOR_IMPL(use_binary, FuzzerProtoType)            \
+  DEFINE_CUSTOM_PROTO_CROSSOVER_IMPL(use_binary, FuzzerProtoType)          \
+  DEFINE_TEST_ONE_PROTO_INPUT_RET_IMPL(use_binary, FuzzerProtoType)        \
+  DEFINE_POST_PROCESS_PROTO_MUTATION_IMPL(FuzzerProtoType)                 \
+  static int TestOneProtoInput(arg)
 
 namespace protobuf_mutator {
 namespace libfuzzer {
@@ -129,11 +164,11 @@ namespace macro_internal {
 template <typename T>
 struct GetFirstParam;
 
-template <class Arg>
-struct GetFirstParam<void (*)(Arg)> {
-  using type = typename std::remove_const<
-      typename std::remove_reference<Arg>::type>::type;
-};
+template <class Ret, class Arg>
+  struct GetFirstParam<Ret (*)(Arg)> {
+    using type = typename std::remove_const<
+        typename std::remove_reference<Arg>::type>::type;
+  };
 
 }  // namespace macro_internal
 

src/libfuzzer/libfuzzer_test.cc

@@ -21,6 +21,7 @@ using protobuf_mutator::protobuf::util::MessageDifferencer;
 using ::testing::_;
 using ::testing::AllOf;
 using ::testing::DoAll;
+using ::testing::Eq;
 using ::testing::Ref;
 using ::testing::SaveArg;
 using ::testing::SaveArgPointee;
@@ -42,9 +43,24 @@ protobuf_mutator::libfuzzer::PostProcessorRegistration<protobuf_mutator::Msg>
       mock_fuzzer->PostProcess(message, seed);
     }};
 
-DEFINE_TEXT_PROTO_FUZZER(const protobuf_mutator::Msg& message) {
+#if defined(LIBPROTOBUFMUTATOR_ENABLE_RETURN_CODE) || \
+    defined(TEST_RETURN_OVERLOAD)
+const int returnCode = 5;
+#endif
+
+#ifdef TEST_RETURN_OVERLOAD
+DEFINE_TEXT_PROTO_FUZZER_RET(const protobuf_mutator::Msg& message) {
   mock_fuzzer->TestOneInput(message);
-}
+  return returnCode;
+  }
+#else
+  DEFINE_TEXT_PROTO_FUZZER(const protobuf_mutator::Msg& message) {
+    mock_fuzzer->TestOneInput(message);
+  #ifdef LIBPROTOBUFMUTATOR_ENABLE_RETURN_CODE
+    return returnCode;
+  #endif
+  }
+#endif
 
 MATCHER_P(IsMessageEq, msg, "") {
   return MessageDifferencer::Equals(arg, msg.get());
@@ -59,7 +75,12 @@ TEST(LibFuzzerTest, LLVMFuzzerTestOneInput) {
       .WillOnce(DoAll(SaveArgPointee<0>(&msg), SaveArg<1>(&seed)));
   EXPECT_CALL(
       mock, TestOneInput(AllOf(IsMessageEq(std::cref(msg)), IsInitialized())));
+#if defined(LIBPROTOBUFMUTATOR_ENABLE_RETURN_CODE) || \
+    defined(TEST_RETURN_VARIANT)
+  EXPECT_THAT(LLVMFuzzerTestOneInput((const uint8_t*)"", 0), Eq(returnCode));
+#else
   LLVMFuzzerTestOneInput((const uint8_t*)"", 0);
+#endif
 
   EXPECT_CALL(mock, PostProcess(_, seed)).WillOnce(SaveArgPointee<0>(&msg));
   EXPECT_CALL(