[libFuzzer] more refactoring: make sure CurrentUnitData is awlays a valid pointer to read from

llvm-svn: 270942

3 files changed, 26 insertions, 15 deletions

diff --git a/llvm/lib/Fuzzer/FuzzerInternal.h b/llvm/lib/Fuzzer/FuzzerInternal.h
index c26aa5f4b21..4e2c8e53d1d 100644
--- a/llvm/lib/Fuzzer/FuzzerInternal.h
+++ b/llvm/lib/Fuzzer/FuzzerInternal.h
@@ -441,9 +441,10 @@ private:
   void DumpCurrentUnit(const char *Prefix);
   void DeathCallback();
 
-  void SetCurrentUnit(const uint8_t *Data, size_t Size);
+  void SetCurrentUnit(size_t Size);
   size_t GetCurrentUnitNoThreadCheck(const uint8_t **Data) const;
-  const uint8_t *CurrentUnitData = nullptr;
+  void LazyAllocateCurrentUnitData();
+  uint8_t *CurrentUnitData = nullptr;
   size_t CurrentUnitSize = 0;
   bool InOOMState = false;
 
@@ -455,7 +456,6 @@ private:
 
   std::vector<Unit> Corpus;
   std::unordered_set<std::string> UnitHashesAddedToCorpus;
-  std::vector<uint8_t> MutateInPlaceHere;
 
   std::piecewise_constant_distribution<double> CorpusDistribution;
   UserCallback CB;
diff --git a/llvm/lib/Fuzzer/FuzzerLoop.cpp b/llvm/lib/Fuzzer/FuzzerLoop.cpp
index 6aaa09fddc0..3838cdde1d9 100644
--- a/llvm/lib/Fuzzer/FuzzerLoop.cpp
+++ b/llvm/lib/Fuzzer/FuzzerLoop.cpp
@@ -159,6 +159,11 @@ Fuzzer::Fuzzer(UserCallback CB, MutationDispatcher &MD, FuzzingOptions Options)
   IsMyThread = true;
 }
 
+void Fuzzer::LazyAllocateCurrentUnitData() {
+  if (CurrentUnitData || Options.MaxLen == 0) return;
+  CurrentUnitData = new uint8_t[Options.MaxLen];
+}
+
 void Fuzzer::SetDeathCallback() {
   CHECK_WEAK_API_FUNCTION(__sanitizer_set_death_callback);
   __sanitizer_set_death_callback(StaticDeathCallback);
@@ -501,10 +506,9 @@ void __sanitizer_free_hook(void *ptr) {
 }
 }  // extern "C"
 
-void Fuzzer::SetCurrentUnit(const uint8_t *Data, size_t Size) {
+void Fuzzer::SetCurrentUnit(size_t Size) {
   assert(InFuzzingThread());
   CurrentUnitSize = Size;
-  CurrentUnitData = Data;
 }
 
 size_t Fuzzer::GetCurrentUnitNoThreadCheck(const uint8_t **Data) const {
@@ -518,18 +522,21 @@ size_t Fuzzer::GetCurrentUnitInFuzzingThead(const uint8_t **Data) const {
 }
 
 void Fuzzer::ExecuteCallback(const uint8_t *Data, size_t Size) {
+  LazyAllocateCurrentUnitData();
   UnitStartTime = system_clock::now();
   // We copy the contents of Unit into a separate heap buffer
   // so that we reliably find buffer overflows in it.
   std::unique_ptr<uint8_t[]> DataCopy(new uint8_t[Size]);
   memcpy(DataCopy.get(), Data, Size);
+  if (CurrentUnitData && CurrentUnitData != Data)
+    memcpy(CurrentUnitData, Data, Size);
   AssignTaintLabels(DataCopy.get(), Size);
-  SetCurrentUnit(DataCopy.get(), Size);
+  SetCurrentUnit(Size);
   AllocTracer.Start();
   int Res = CB(DataCopy.get(), Size);
   (void)Res;
   HasMoreMallocsThanFrees = AllocTracer.Stop();
-  SetCurrentUnit(nullptr, 0);
+  SetCurrentUnit(0);
   assert(Res == 0);
 }
 
@@ -689,7 +696,7 @@ void Fuzzer::TryDetectingAMemoryLeak(const uint8_t *Data, size_t Size,
     if (DuringInitialCorpusExecution)
       Printf("\nINFO: a leak has been found in the initial corpus.\n\n");
     Printf("INFO: to ignore leaks on libFuzzer side use -detect_leaks=0.\n\n");
-    SetCurrentUnit(Data, Size);
+    SetCurrentUnit(Size);
     DumpCurrentUnit("leak-");
     PrintFinalStats();
     _Exit(Options.ErrorExitCode);  // not exit() to disable lsan further on.
@@ -697,32 +704,33 @@ void Fuzzer::TryDetectingAMemoryLeak(const uint8_t *Data, size_t Size,
 }
 
 void Fuzzer::MutateAndTestOne() {
+  LazyAllocateCurrentUnitData();
   MD.StartMutationSequence();
 
   auto &U = ChooseUnitToMutate();
-  MutateInPlaceHere.resize(Options.MaxLen);
+  assert(CurrentUnitData);
   size_t Size = U.size();
   assert(Size <= Options.MaxLen && "Oversized Unit");
-  memcpy(MutateInPlaceHere.data(), U.data(), Size);
+  memcpy(CurrentUnitData, U.data(), Size);
 
   for (int i = 0; i < Options.MutateDepth; i++) {
     size_t NewSize = 0;
     if (LLVMFuzzerCustomMutator)
-      NewSize = LLVMFuzzerCustomMutator(MutateInPlaceHere.data(), Size,
+      NewSize = LLVMFuzzerCustomMutator(CurrentUnitData, Size,
                                         Options.MaxLen, MD.GetRand().Rand());
     else
-      NewSize = MD.Mutate(MutateInPlaceHere.data(), Size, Options.MaxLen);
+      NewSize = MD.Mutate(CurrentUnitData, Size, Options.MaxLen);
     assert(NewSize > 0 && "Mutator returned empty unit");
     assert(NewSize <= Options.MaxLen &&
            "Mutator return overisized unit");
     Size = NewSize;
     if (Options.OnlyASCII)
-      ToASCII(MutateInPlaceHere.data(), Size);
+      ToASCII(CurrentUnitData, Size);
     if (i == 0)
       StartTraceRecording();
-    RunOneAndUpdateCorpus(MutateInPlaceHere.data(), Size);
+    RunOneAndUpdateCorpus(CurrentUnitData, Size);
     StopTraceRecording();
-    TryDetectingAMemoryLeak(MutateInPlaceHere.data(), Size,
+    TryDetectingAMemoryLeak(CurrentUnitData, Size,
                             /*DuringInitialCorpusExecution*/ false);
   }
 }
diff --git a/llvm/lib/Fuzzer/test/fuzzer-leak.test b/llvm/lib/Fuzzer/test/fuzzer-leak.test
index 61410aae401..6b4b484b0bf 100644
--- a/llvm/lib/Fuzzer/test/fuzzer-leak.test
+++ b/llvm/lib/Fuzzer/test/fuzzer-leak.test
@@ -14,6 +14,9 @@ RUN: not LLVMFuzzer-LeakTest -runs=100000                 2>&1 | FileCheck %s --
 LEAK_AFTER: Done 100000 runs in
 LEAK_AFTER: ERROR: LeakSanitizer: detected memory leaks
 
+RUN: not LLVMFuzzer-LeakTest -runs=100000 -max_len=1 2>&1 | FileCheck %s --check-prefix=MAX_LEN_1
+MAX_LEN_1: Test unit written to ./leak-7cf184f4c67ad58283ecb19349720b0cae756829
+
 RUN: not LLVMFuzzer-LeakTimeoutTest -timeout=1 2>&1 | FileCheck %s --check-prefix=LEAK_TIMEOUT
 LEAK_TIMEOUT: ERROR: libFuzzer: timeout after
 LEAK_TIMEOUT-NOT: LeakSanitizer