[libFuzzer] use the new chainable malloc hooks instead of the old un-chainable ones, use atomic for malloc/free counters instead of a thread local counter in the main thread. This should make on-the-spot leak detection in libFuzzer more reliable

llvm-svn: 272948

3 files changed, 32 insertions, 39 deletions

diff --git a/llvm/lib/Fuzzer/FuzzerExtFunctions.def b/llvm/lib/Fuzzer/FuzzerExtFunctions.def
index 2810a0046e3..f7dcf9b54fb 100644
--- a/llvm/lib/Fuzzer/FuzzerExtFunctions.def
+++ b/llvm/lib/Fuzzer/FuzzerExtFunctions.def
@@ -31,6 +31,10 @@ EXT_FUNC(__lsan_disable, void, (), false);
 EXT_FUNC(__lsan_do_recoverable_leak_check, int, (), false);
 EXT_FUNC(__sanitizer_get_coverage_pc_buffer, uintptr_t, (uintptr_t**), true);
 EXT_FUNC(__sanitizer_get_number_of_counters, size_t, (), false);
+EXT_FUNC(__sanitizer_install_malloc_and_free_hooks, int,
+         (void (*malloc_hook)(const volatile void *, size_t),
+          void (*free_hook)(const volatile void *)),
+         false);
 EXT_FUNC(__sanitizer_get_total_unique_caller_callee_pairs, size_t, (), false);
 EXT_FUNC(__sanitizer_get_total_unique_coverage, size_t, (), true);
 EXT_FUNC(__sanitizer_print_memory_profile, int, (size_t), false);
diff --git a/llvm/lib/Fuzzer/FuzzerLoop.cpp b/llvm/lib/Fuzzer/FuzzerLoop.cpp
index c7917aca6d7..51cb8703c26 100644
--- a/llvm/lib/Fuzzer/FuzzerLoop.cpp
+++ b/llvm/lib/Fuzzer/FuzzerLoop.cpp
@@ -124,6 +124,28 @@ struct CoverageController {
   }
 };
 
+// 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() {
+    Mallocs = 0;
+    Frees = 0;
+  }
+  // Returns true if there were more mallocs than frees.
+  bool Stop() { return Mallocs > Frees; }
+  std::atomic<size_t> Mallocs;
+  std::atomic<size_t> Frees;
+};
+
+static MallocFreeTracer AllocTracer;
+
+void MallocHook(const volatile void *ptr, size_t size) {
+  AllocTracer.Mallocs++;
+}
+void FreeHook(const volatile void *ptr) {
+  AllocTracer.Frees++;
+}
+
 Fuzzer::Fuzzer(UserCallback CB, MutationDispatcher &MD, FuzzingOptions Options)
     : CB(CB), MD(MD), Options(Options) {
   SetDeathCallback();
@@ -132,6 +154,8 @@ Fuzzer::Fuzzer(UserCallback CB, MutationDispatcher &MD, FuzzingOptions Options)
   F = this;
   ResetCoverage();
   IsMyThread = true;
+  if (Options.DetectLeaks && EF->__sanitizer_install_malloc_and_free_hooks)
+    EF->__sanitizer_install_malloc_and_free_hooks(MallocHook, FreeHook);
 }
 
 void Fuzzer::LazyAllocateCurrentUnitData() {
@@ -444,38 +468,6 @@ void Fuzzer::RunOneAndUpdateCorpus(const uint8_t *Data, size_t Size) {
     ReportNewCoverage({Data, Data + Size});
 }
 
-// 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() {
-    Mallocs = 0;
-    Frees = 0;
-  }
-  // Returns true if there were more mallocs than frees.
-  bool Stop() { return Mallocs > Frees; }
-  size_t Mallocs;
-  size_t Frees;
-};
-
-static thread_local MallocFreeTracer AllocTracer;
-
-// FIXME: The hooks only count on Linux because
-// on Mac OSX calls to malloc are intercepted before
-// thread local storage is initialised leading to
-// crashes when accessing ``AllocTracer``.
-extern "C" {
-__attribute__((weak))
-void __sanitizer_malloc_hook(void *ptr, size_t size) {
-  if (!LIBFUZZER_APPLE)
-    AllocTracer.Mallocs++;
-}
-__attribute__((weak))
-void __sanitizer_free_hook(void *ptr) {
-  if (!LIBFUZZER_APPLE)
-    AllocTracer.Frees++;
-}
-}  // extern "C"
-
 size_t Fuzzer::GetCurrentUnitInFuzzingThead(const uint8_t **Data) const {
   assert(InFuzzingThread());
   *Data = CurrentUnitData;
diff --git a/llvm/lib/Fuzzer/test/fuzzer-threaded.test b/llvm/lib/Fuzzer/test/fuzzer-threaded.test
index e2c823cda1a..c58a33456cc 100644
--- a/llvm/lib/Fuzzer/test/fuzzer-threaded.test
+++ b/llvm/lib/Fuzzer/test/fuzzer-threaded.test
@@ -1,10 +1,7 @@
 CHECK: Done 1000 runs in
 
-# TODO(kcc): re-enable leak detection here.
-# Currently laak detection makes run counts imprecise.
-
-RUN: LLVMFuzzer-ThreadedTest -use_traces=1 -runs=1000 -detect_leaks=0 2>&1 | FileCheck %s
-RUN: LLVMFuzzer-ThreadedTest -use_traces=1 -runs=1000 -detect_leaks=0 2>&1 | FileCheck %s
-RUN: LLVMFuzzer-ThreadedTest -use_traces=1 -runs=1000 -detect_leaks=0 2>&1 | FileCheck %s
-RUN: LLVMFuzzer-ThreadedTest -use_traces=1 -runs=1000 -detect_leaks=0 2>&1 | FileCheck %s
+RUN: LLVMFuzzer-ThreadedTest -use_traces=1 -runs=1000  2>&1 | FileCheck %s
+RUN: LLVMFuzzer-ThreadedTest -use_traces=1 -runs=1000  2>&1 | FileCheck %s
+RUN: LLVMFuzzer-ThreadedTest -use_traces=1 -runs=1000  2>&1 | FileCheck %s
+RUN: LLVMFuzzer-ThreadedTest -use_traces=1 -runs=1000  2>&1 | FileCheck %s