From c0e47a82c10326d2970f4cfc5365731c6d9ae852 Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Thu, 16 Jun 2016 20:17:41 +0000
Subject: [PATCH] [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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@272948 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Fuzzer/FuzzerExtFunctions.def    |  4 ++
 lib/Fuzzer/FuzzerLoop.cpp            | 56 ++++++++++++----------------
 lib/Fuzzer/test/fuzzer-threaded.test | 11 ++----
 3 files changed, 32 insertions(+), 39 deletions(-)

diff --git a/lib/Fuzzer/FuzzerExtFunctions.def b/lib/Fuzzer/FuzzerExtFunctions.def
index 2810a0046e3..f7dcf9b54fb 100644
--- a/lib/Fuzzer/FuzzerExtFunctions.def
+++ b/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/lib/Fuzzer/FuzzerLoop.cpp b/lib/Fuzzer/FuzzerLoop.cpp
index c7917aca6d7..51cb8703c26 100644
--- a/lib/Fuzzer/FuzzerLoop.cpp
+++ b/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/lib/Fuzzer/test/fuzzer-threaded.test b/lib/Fuzzer/test/fuzzer-threaded.test
index e2c823cda1a..c58a33456cc 100644
--- a/lib/Fuzzer/test/fuzzer-threaded.test
+++ b/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