[libFuzzer] add -trace_malloc= flag

llvm-svn: 284149

7 files changed, 70 insertions, 5 deletions

diff --git a/llvm/lib/Fuzzer/FuzzerDriver.cpp b/llvm/lib/Fuzzer/FuzzerDriver.cpp
index 974bbca5600..f15c79fcd4d 100644
--- a/llvm/lib/Fuzzer/FuzzerDriver.cpp
+++ b/llvm/lib/Fuzzer/FuzzerDriver.cpp
@@ -405,6 +405,7 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
   Options.OnlyASCII = Flags.only_ascii;
   Options.OutputCSV = Flags.output_csv;
   Options.DetectLeaks = Flags.detect_leaks;
+  Options.TraceMalloc = Flags.trace_malloc;
   Options.RssLimitMb = Flags.rss_limit_mb;
   if (Flags.runs >= 0)
     Options.MaxNumberOfRuns = Flags.runs;
diff --git a/llvm/lib/Fuzzer/FuzzerFlags.def b/llvm/lib/Fuzzer/FuzzerFlags.def
index 43712764f51..d3b8b98ca3a 100644
--- a/llvm/lib/Fuzzer/FuzzerFlags.def
+++ b/llvm/lib/Fuzzer/FuzzerFlags.def
@@ -91,6 +91,8 @@ FUZZER_FLAG_INT(close_fd_mask, 0, "If 1, close stdout at startup; "
     "Be careful, this will also close e.g. asan's stderr/stdout.")
 FUZZER_FLAG_INT(detect_leaks, 1, "If 1, and if LeakSanitizer is enabled "
     "try to detect memory leaks during fuzzing (i.e. not only at shut down).")
+FUZZER_FLAG_INT(trace_malloc, 0, "If >= 1 will print all mallocs/frees. "
+                                 "If >= 2 will also print stack traces.")
 FUZZER_FLAG_INT(rss_limit_mb, 2048, "If non-zero, the fuzzer will exit upon"
     "reaching this limit of RSS memory usage.")
 FUZZER_FLAG_STRING(exit_on_src_pos, "Exit if a newly found PC originates"
diff --git a/llvm/lib/Fuzzer/FuzzerLoop.cpp b/llvm/lib/Fuzzer/FuzzerLoop.cpp
index 1b5e0265e13..37db6ee09e5 100644
--- a/llvm/lib/Fuzzer/FuzzerLoop.cpp
+++ b/llvm/lib/Fuzzer/FuzzerLoop.cpp
@@ -111,23 +111,46 @@ bool Fuzzer::RecordMaxCoverage(Fuzzer::Coverage *C) {
 // Leak detection is expensive, so we first check if there were more mallocs
 // than frees (using the sanitizer malloc hooks) and only then try to call lsan.
 struct MallocFreeTracer {
-  void Start() {
+  void Start(int TraceLevel) {
+    this->TraceLevel = TraceLevel;
+    if (TraceLevel)
+      Printf("MallocFreeTracer: START\n");
     Mallocs = 0;
     Frees = 0;
   }
   // Returns true if there were more mallocs than frees.
-  bool Stop() { return Mallocs > Frees; }
+  bool Stop() { 
+    if (TraceLevel)
+      Printf("MallocFreeTracer: STOP %zd %zd (%s)\n", Mallocs.load(),
+             Frees.load(), Mallocs == Frees ? "same" : "DIFFERENT");
+    bool Result = Mallocs > Frees;
+    Mallocs = 0;
+    Frees = 0;
+    TraceLevel = 0;
+    return Result;
+  }
   std::atomic<size_t> Mallocs;
   std::atomic<size_t> Frees;
+  int TraceLevel = 0;
 };
 
 static MallocFreeTracer AllocTracer;
 
 void MallocHook(const volatile void *ptr, size_t size) {
-  AllocTracer.Mallocs++;
+  size_t N = AllocTracer.Mallocs++;
+  if (int TraceLevel = AllocTracer.TraceLevel) {
+    Printf("MALLOC[%zd] %p %zd\n", N, ptr, size);
+    if (TraceLevel >= 2 && EF)
+      EF->__sanitizer_print_stack_trace();
+  }
 }
 void FreeHook(const volatile void *ptr) {
-  AllocTracer.Frees++;
+  size_t N = AllocTracer.Frees++;
+  if (int TraceLevel = AllocTracer.TraceLevel) {
+    Printf("FREE[%zd]   %p %zd\n", N, ptr);
+    if (TraceLevel >= 2 && EF)
+      EF->__sanitizer_print_stack_trace();
+  }
 }
 
 Fuzzer::Fuzzer(UserCallback CB, InputCorpus &Corpus, MutationDispatcher &MD,
@@ -486,7 +509,7 @@ void Fuzzer::ExecuteCallback(const uint8_t *Data, size_t Size) {
   if (CurrentUnitData && CurrentUnitData != Data)
     memcpy(CurrentUnitData, Data, Size);
   CurrentUnitSize = Size;
-  AllocTracer.Start();
+  AllocTracer.Start(Options.TraceMalloc);
   UnitStartTime = system_clock::now();
   ResetCounters();  // Reset coverage right before the callback.
   TPC.ResetMaps();
diff --git a/llvm/lib/Fuzzer/FuzzerOptions.h b/llvm/lib/Fuzzer/FuzzerOptions.h
index ed31fa22b34..d0bac9cee83 100644
--- a/llvm/lib/Fuzzer/FuzzerOptions.h
+++ b/llvm/lib/Fuzzer/FuzzerOptions.h
@@ -51,6 +51,7 @@ struct FuzzingOptions {
   bool PrintCorpusStats = false;
   bool PrintCoverage = false;
   bool DetectLeaks = true;
+  int  TraceMalloc = 0;
 };
 
 }  // namespace fuzzer
diff --git a/llvm/lib/Fuzzer/test/CMakeLists.txt b/llvm/lib/Fuzzer/test/CMakeLists.txt
index d2ccc23fa17..d52c282032b 100644
--- a/llvm/lib/Fuzzer/test/CMakeLists.txt
+++ b/llvm/lib/Fuzzer/test/CMakeLists.txt
@@ -106,6 +106,7 @@ set(Tests
   ThreadedLeakTest
   ThreadedTest
   TimeoutTest
+  TraceMallocTest
   )
 
 if(APPLE)
diff --git a/llvm/lib/Fuzzer/test/TraceMallocTest.cpp b/llvm/lib/Fuzzer/test/TraceMallocTest.cpp
new file mode 100644
index 00000000000..43e6950e185
--- /dev/null
+++ b/llvm/lib/Fuzzer/test/TraceMallocTest.cpp
@@ -0,0 +1,27 @@
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+
+// Tests -trace_malloc
+#include <assert.h>
+#include <cstdint>
+#include <cstdlib>
+#include <cstddef>
+#include <iostream>
+
+int *Ptr;
+
+extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
+  if (!Size) return 0;
+  if (*Data == 1) {
+    delete Ptr;
+    Ptr = nullptr;
+  } else if (*Data == 2) {
+    delete Ptr;
+    Ptr = new int;
+  } else if (*Data == 3) {
+    if (!Ptr)
+      Ptr = new int;
+  }
+  return 0;
+}
+
diff --git a/llvm/lib/Fuzzer/test/trace-malloc.test b/llvm/lib/Fuzzer/test/trace-malloc.test
new file mode 100644
index 00000000000..c8ef833a56d
--- /dev/null
+++ b/llvm/lib/Fuzzer/test/trace-malloc.test
@@ -0,0 +1,10 @@
+RUN: LLVMFuzzer-TraceMallocTest -seed=1 -trace_malloc=1 -runs=1000 2>&1 | FileCheck %s
+CHECK-DAG: MallocFreeTracer: STOP 0 0 (same)
+CHECK-DAG: MallocFreeTracer: STOP 0 1 (DIFFERENT)
+CHECK-DAG: MallocFreeTracer: STOP 1 0 (DIFFERENT)
+CHECK-DAG: MallocFreeTracer: STOP 1 1 (same)
+
+RUN: LLVMFuzzer-TraceMallocTest -seed=1 -trace_malloc=2 -runs=1000 2>&1 | FileCheck %s --check-prefix=TRACE2
+TRACE2-DAG: FREE[0]
+TRACE2-DAG: MALLOC[0]
+TRACE2-DAG: in LLVMFuzzerTestOneInput