[sanitizers] introduce yet another API function: __sanitizer_install_malloc_and_free_hooks

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@272943 91177308-0d34-0410-b5e6-96231b3b80d8

8 files changed, 83 insertions, 14 deletions

diff --git a/lib/asan/asan_internal.h b/lib/asan/asan_internal.h
index 8a815e3c1..91101722c 100644
--- a/lib/asan/asan_internal.h
+++ b/lib/asan/asan_internal.h
@@ -110,10 +110,16 @@ bool PlatformHasDifferentMemcpyAndMemmove();
 
 // Add convenient macro for interface functions that may be represented as
 // weak hooks.
-#define ASAN_MALLOC_HOOK(ptr, size) \
-  if (&__sanitizer_malloc_hook) __sanitizer_malloc_hook(ptr, size)
-#define ASAN_FREE_HOOK(ptr) \
-  if (&__sanitizer_free_hook) __sanitizer_free_hook(ptr)
+#define ASAN_MALLOC_HOOK(ptr, size)                                   \
+  do {                                                                \
+    if (&__sanitizer_malloc_hook) __sanitizer_malloc_hook(ptr, size); \
+    RunMallocHooks(ptr, size);                                        \
+  } while (false)
+#define ASAN_FREE_HOOK(ptr)                                 \
+  do {                                                      \
+    if (&__sanitizer_free_hook) __sanitizer_free_hook(ptr); \
+    RunFreeHooks(ptr);                                      \
+  } while (false)
 #define ASAN_ON_ERROR() \
   if (&__asan_on_error) __asan_on_error()
 
diff --git a/lib/asan/asan_win_dll_thunk.cc b/lib/asan/asan_win_dll_thunk.cc
index 22d1d88bb..d63548feb 100644
--- a/lib/asan/asan_win_dll_thunk.cc
+++ b/lib/asan/asan_win_dll_thunk.cc
@@ -343,6 +343,7 @@ INTERFACE_FUNCTION(__sanitizer_unaligned_store16)
 INTERFACE_FUNCTION(__sanitizer_unaligned_store32)
 INTERFACE_FUNCTION(__sanitizer_unaligned_store64)
 INTERFACE_FUNCTION(__sanitizer_verify_contiguous_container)
+INTERFACE_FUNCTION(__sanitizer_install_malloc_and_free_hooks)
 
 // TODO(timurrrr): Add more interface functions on the as-needed basis.
 
diff --git a/lib/lsan/lsan_allocator.cc b/lib/lsan/lsan_allocator.cc
index 0a3678132..a5220f1a3 100644
--- a/lib/lsan/lsan_allocator.cc
+++ b/lib/lsan/lsan_allocator.cc
@@ -99,11 +99,13 @@ void *Allocate(const StackTrace &stack, uptr size, uptr alignment,
     memset(p, 0, size);
   RegisterAllocation(stack, p, size);
   if (&__sanitizer_malloc_hook) __sanitizer_malloc_hook(p, size);
+  RunMallocHooks(p, size);
   return p;
 }
 
 void Deallocate(void *p) {
   if (&__sanitizer_free_hook) __sanitizer_free_hook(p);
+  RunFreeHooks(p);
   RegisterDeallocation(p);
   allocator.Deallocate(&cache, p);
 }
diff --git a/lib/msan/msan.h b/lib/msan/msan.h
index be74f6ea1..1f2ff59ca 100644
--- a/lib/msan/msan.h
+++ b/lib/msan/msan.h
@@ -309,15 +309,21 @@ void MsanTSDDtor(void *tsd);
 
 }  // namespace __msan
 
-#define MSAN_MALLOC_HOOK(ptr, size)     \
-  if (&__sanitizer_malloc_hook) {       \
-    UnpoisonParam(2);                   \
-    __sanitizer_malloc_hook(ptr, size); \
-  }
-#define MSAN_FREE_HOOK(ptr)     \
-  if (&__sanitizer_free_hook) { \
-    UnpoisonParam(1);           \
-    __sanitizer_free_hook(ptr); \
-  }
+#define MSAN_MALLOC_HOOK(ptr, size)       \
+  do {                                    \
+    if (&__sanitizer_malloc_hook) {       \
+      UnpoisonParam(2);                   \
+      __sanitizer_malloc_hook(ptr, size); \
+    }                                     \
+    RunMallocHooks(ptr, size);            \
+  } while (false)
+#define MSAN_FREE_HOOK(ptr)       \
+  do {                            \
+    if (&__sanitizer_free_hook) { \
+      UnpoisonParam(1);           \
+      __sanitizer_free_hook(ptr); \
+    }                             \
+    RunFreeHooks(ptr);            \
+  } while (false)
 
 #endif  // MSAN_H
diff --git a/lib/sanitizer_common/sanitizer_allocator_interface.h b/lib/sanitizer_common/sanitizer_allocator_interface.h
index 2cd924c99..797c38a79 100644
--- a/lib/sanitizer_common/sanitizer_allocator_interface.h
+++ b/lib/sanitizer_common/sanitizer_allocator_interface.h
@@ -29,6 +29,10 @@ SANITIZER_INTERFACE_ATTRIBUTE uptr __sanitizer_get_heap_size();
 SANITIZER_INTERFACE_ATTRIBUTE uptr __sanitizer_get_free_bytes();
 SANITIZER_INTERFACE_ATTRIBUTE uptr __sanitizer_get_unmapped_bytes();
 
+SANITIZER_INTERFACE_ATTRIBUTE int __sanitizer_install_malloc_and_free_hooks(
+    void (*malloc_hook)(const void *, uptr),
+    void (*free_hook)(const void *));
+
 SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE
     /* OPTIONAL */ void __sanitizer_malloc_hook(void *ptr, uptr size);
 SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE
diff --git a/lib/sanitizer_common/sanitizer_common.cc b/lib/sanitizer_common/sanitizer_common.cc
index fb68e5299..8aea7bcdb 100644
--- a/lib/sanitizer_common/sanitizer_common.cc
+++ b/lib/sanitizer_common/sanitizer_common.cc
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "sanitizer_common.h"
+#include "sanitizer_allocator_interface.h"
 #include "sanitizer_allocator_internal.h"
 #include "sanitizer_flags.h"
 #include "sanitizer_libc.h"
@@ -421,6 +422,44 @@ void PrintCmdline() {
   Printf("\n\n");
 }
 
+// Malloc hooks.
+static const int kMaxMallocFreeHooks = 5;
+struct MallocFreeHook {
+  void (*malloc_hook)(const void *, uptr);
+  void (*free_hook)(const void *);
+};
+
+static MallocFreeHook MFHooks[kMaxMallocFreeHooks];
+
+void RunMallocHooks(const void *ptr, uptr size) {
+  for (int i = 0; i < kMaxMallocFreeHooks; i++) {
+    auto hook = MFHooks[i].malloc_hook;
+    if (!hook) return;
+    hook(ptr, size);
+  }
+}
+
+void RunFreeHooks(const void *ptr) {
+  for (int i = 0; i < kMaxMallocFreeHooks; i++) {
+    auto hook = MFHooks[i].free_hook;
+    if (!hook) return;
+    hook(ptr);
+  }
+}
+
+static int InstallMallocFreeHooks(void (*malloc_hook)(const void *, uptr),
+                                  void (*free_hook)(const void *)) {
+  if (!malloc_hook || !free_hook) return 0;
+  for (int i = 0; i < kMaxMallocFreeHooks; i++) {
+    if (MFHooks[i].malloc_hook == nullptr) {
+      MFHooks[i].malloc_hook = malloc_hook;
+      MFHooks[i].free_hook = free_hook;
+      return i + 1;
+    }
+  }
+  return 0;
+}
+
 } // namespace __sanitizer
 
 using namespace __sanitizer;  // NOLINT
@@ -443,4 +482,11 @@ SANITIZER_INTERFACE_ATTRIBUTE
 void __sanitizer_set_death_callback(void (*callback)(void)) {
   SetUserDieCallback(callback);
 }
+
+SANITIZER_INTERFACE_ATTRIBUTE
+int __sanitizer_install_malloc_and_free_hooks(void (*malloc_hook)(const void *,
+                                                                  uptr),
+                                              void (*free_hook)(const void *)) {
+  return InstallMallocFreeHooks(malloc_hook, free_hook);
+}
 } // extern "C"
diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index f78f50745..24ef93a9c 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -103,6 +103,8 @@ void NoHugePagesInRegion(uptr addr, uptr length);
 void DontDumpShadowMemory(uptr addr, uptr length);
 // Check if the built VMA size matches the runtime one.
 void CheckVMASize();
+void RunMallocHooks(const void *ptr, uptr size);
+void RunFreeHooks(const void *ptr);
 
 // InternalScopedBuffer can be used instead of large stack arrays to
 // keep frame size low.
diff --git a/lib/tsan/rtl/tsan_mman.cc b/lib/tsan/rtl/tsan_mman.cc
index aa7198200..7693077f6 100644
--- a/lib/tsan/rtl/tsan_mman.cc
+++ b/lib/tsan/rtl/tsan_mman.cc
@@ -228,6 +228,7 @@ void invoke_malloc_hook(void *ptr, uptr size) {
   if (ctx == 0 || !ctx->initialized || thr->ignore_interceptors)
     return;
   __sanitizer_malloc_hook(ptr, size);
+  RunMallocHooks(ptr, size);
 }
 
 void invoke_free_hook(void *ptr) {
@@ -235,6 +236,7 @@ void invoke_free_hook(void *ptr) {
   if (ctx == 0 || !ctx->initialized || thr->ignore_interceptors)
     return;
   __sanitizer_free_hook(ptr);
+  RunFreeHooks(ptr);
 }
 
 void *internal_alloc(MBlockType typ, uptr sz) {