Avoid parsing of last mutated message

It safe some CPU cycles. Still main reason is to avoid
coverage noise from parser and custom post processors.

2 files changed, 62 insertions, 0 deletions

diff --git a/src/libfuzzer/libfuzzer_macro.cc b/src/libfuzzer/libfuzzer_macro.cc
index 11730cf..4e506cb 100644
--- a/src/libfuzzer/libfuzzer_macro.cc
+++ b/src/libfuzzer/libfuzzer_macro.cc
@@ -15,6 +15,8 @@
 #include "src/libfuzzer/libfuzzer_macro.h"
 
 #include <algorithm>
+#include <memory>
+#include <vector>
 
 #include "src/binary_format.h"
 #include "src/libfuzzer/libfuzzer_mutator.h"
@@ -91,6 +93,35 @@ class BinaryOutputWriter : public OutputWriter {
   }
 };
 
+class LastMutationCache {
+ public:
+  void Store(const uint8_t* data, size_t size, protobuf::Message* message) {
+    if (!message_) message_.reset(message->New());
+    message->GetReflection()->Swap(message, message_.get());
+    data_.assign(data, data + size);
+  }
+
+  bool LoadIfSame(const uint8_t* data, size_t size,
+                  protobuf::Message* message) {
+    if (!message_ || size != data_.size() ||
+        !std::equal(data_.begin(), data_.end(), data))
+      return false;
+
+    message->GetReflection()->Swap(message, message_.get());
+    message_.reset();
+    return true;
+  }
+
+ private:
+  std::vector<uint8_t> data_;
+  std::unique_ptr<protobuf::Message> message_;
+};
+
+LastMutationCache* GetCache() {
+  static LastMutationCache cache;
+  return &cache;
+}
+
 Mutator* GetMutator() {
   static Mutator mutator;
   return &mutator;
@@ -111,6 +142,7 @@ size_t MutateMessage(unsigned int seed, const InputReader& input,
   GetMutator()->Mutate(message, max_size);
   if (size_t new_size = output->Write(*message)) {
     assert(new_size <= output->size());
+    GetCache()->Store(output->data(), new_size, message);
     return new_size;
   }
   return 0;
@@ -127,6 +159,7 @@ size_t CrossOverMessages(unsigned int seed, const InputReader& input1,
   GetMutator()->CrossOver(*message2, message1, max_size);
   if (size_t new_size = output->Write(*message1)) {
     assert(new_size <= output->size());
+    GetCache()->Store(output->data(), new_size, message1);
     return new_size;
   }
   return 0;
@@ -189,6 +222,7 @@ size_t CustomProtoCrossOver(bool binary, const uint8_t* data1, size_t size1,
 
 bool LoadProtoInput(bool binary, const uint8_t* data, size_t size,
                     protobuf::Message* input) {
+  if (GetCache()->LoadIfSame(data, size, input)) return true;
   auto result = binary ? ParseBinaryMessage(data, size, input)
                        : ParseTextMessage(data, size, input);
   if (!result) return false;
diff --git a/src/libfuzzer/libfuzzer_test.cc b/src/libfuzzer/libfuzzer_test.cc
index f14071e..f19eb12 100644
--- a/src/libfuzzer/libfuzzer_test.cc
+++ b/src/libfuzzer/libfuzzer_test.cc
@@ -66,4 +66,32 @@ TEST(LibFuzzerTest, LLVMFuzzerTestOneInput) {
       mock, TestOneInput(AllOf(IsMessageEq(std::cref(msg)), IsInitialized())));
   LLVMFuzzerTestOneInput((const uint8_t*)"", 0);
 }
+
+TEST(LibFuzzerTest, LLVMFuzzerCustomMutator) {
+  testing::StrictMock<MockFuzzer> mock;
+  protobuf_mutator::Msg msg;
+  EXPECT_CALL(mock, PostProcess(_, _)).WillOnce(SaveArgPointee<0>(&msg));
+  EXPECT_CALL(
+      mock, TestOneInput(AllOf(IsMessageEq(std::cref(msg)), IsInitialized())));
+
+  uint8_t buff[1024] = {};
+  size_t size = LLVMFuzzerCustomMutator(buff, 0, sizeof(buff), 5);
+  ASSERT_GT(size, 0U);
+  LLVMFuzzerTestOneInput(buff, size);
+}
+
+TEST(LibFuzzerTest, LLVMFuzzerCustomCrossOver) {
+  testing::StrictMock<MockFuzzer> mock;
+  protobuf_mutator::Msg msg;
+  EXPECT_CALL(mock, PostProcess(_, _)).WillOnce(SaveArgPointee<0>(&msg));
+  EXPECT_CALL(
+      mock, TestOneInput(AllOf(IsMessageEq(std::cref(msg)), IsInitialized())));
+
+  uint8_t buff1[1024] = {};
+  uint8_t buff2[1024] = {};
+  uint8_t buff3[1024] = {};
+  size_t size =
+      LLVMFuzzerCustomCrossOver(buff1, 0, buff2, 0, buff3, sizeof(buff3), 6);
+  ASSERT_GT(size, 0U);
+  LLVMFuzzerTestOneInput(buff3, size);
 }