[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

1 files changed, 24 insertions, 32 deletions

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;