[sanitizers] introduce yet another API function: __sanitizer_install_malloc_and_free_hooks

llvm-svn: 272943

compiler-rt/include/sanitizer/allocator_interface.h

@@ -59,6 +59,23 @@ extern "C" {
        deallocation of "ptr". */
   void __sanitizer_malloc_hook(const volatile void *ptr, size_t size);
   void __sanitizer_free_hook(const volatile void *ptr);
+
+  /* Installs a pair of hooks for malloc/free.
+     Several (currently, 5) hook pairs may be installed, they are executed
+     in the order they were installed and after calling
+     __sanitizer_malloc_hook/__sanitizer_free_hook.
+     Unlike __sanitizer_malloc_hook/__sanitizer_free_hook these hooks can be
+     chained and do not rely on weak symbols working on the platform, but
+     require __sanitizer_install_malloc_and_free_hooks to be called at startup
+     and thus will not be called on malloc/free very early in the process.
+     Returns the number of hooks currently installed or 0 on failure.
+     Not thread-safe, should be called in the main thread before starting
+     other threads.
+  */
+  int __sanitizer_install_malloc_and_free_hooks(
+      void (*malloc_hook)(const volatile void *, size_t),
+      void (*free_hook)(const volatile void *));
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

compiler-rt/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()
 

compiler-rt/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.
 

compiler-rt/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);
 }

compiler-rt/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

compiler-rt/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

compiler-rt/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"

compiler-rt/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.

compiler-rt/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) {

compiler-rt/test/sanitizer_common/TestCases/malloc_hook.cc

@@ -10,22 +10,28 @@
 extern "C" {
 const volatile void *global_ptr;
 
+#define WRITE(s) write(1, s, sizeof(s))
+
 // Note: avoid calling functions that allocate memory in malloc/free
 // to avoid infinite recursion.
 void __sanitizer_malloc_hook(const volatile void *ptr, size_t sz) {
   if (__sanitizer_get_ownership(ptr) && sz == 4) {
-    write(1, "MallocHook\n", sizeof("MallocHook\n"));
+    WRITE("MallocHook\n");
     global_ptr = ptr;
   }
 }
 void __sanitizer_free_hook(const volatile void *ptr) {
   if (__sanitizer_get_ownership(ptr) && ptr == global_ptr)
-    write(1, "FreeHook\n", sizeof("FreeHook\n"));
+    WRITE("FreeHook\n");
 }
 }  // extern "C"
 
 volatile int *x;
 
+void MallocHook1(const volatile void *ptr, size_t sz) { WRITE("MH1\n"); }
+void MallocHook2(const volatile void *ptr, size_t sz) { WRITE("MH2\n"); }
+void FreeHook1(const volatile void *ptr) { WRITE("FH1\n"); }
+void FreeHook2(const volatile void *ptr) { WRITE("FH2\n"); }
 // Call this function with uninitialized arguments to poison
 // TLS shadow for function parameters before calling operator
 // new and, eventually, user-provided hook.
@@ -34,9 +40,13 @@ __attribute__((noinline)) void allocate(int *unused1, int *unused2) {
 }
 
 int main() {
+  __sanitizer_install_malloc_and_free_hooks(MallocHook1, FreeHook1);
+  __sanitizer_install_malloc_and_free_hooks(MallocHook2, FreeHook2);
   int *undef1, *undef2;
   allocate(undef1, undef2);
   // CHECK: MallocHook
+  // CHECK: MH1
+  // CHECK: MH2
   // Check that malloc hook was called with correct argument.
   if (global_ptr != (void*)x) {
     _exit(1);
@@ -44,5 +54,7 @@ int main() {
   *x = 0;
   delete x;
   // CHECK: FreeHook
+  // CHECK: FH1
+  // CHECK: FH2
   return 0;
 }