compiler-rt/test/memprof/TestCases/malloc_hook.cpp

   1 // Check that MemProf correctly handles malloc and free hooks.
   2 // RUN: %clangxx_memprof -O2 %s -o %t && %run %t 2>&1 | FileCheck %s
   3
   4 #include <sanitizer/allocator_interface.h>
   5 #include <stdlib.h>
   6 #include <unistd.h>
   7
   8 extern "C" {
   9 const volatile void *global_ptr;
  10
  11 #define WRITE(s) write(1, s, sizeof(s))
  12
  13 // Note: avoid calling functions that allocate memory in malloc/free
  14 // to avoid infinite recursion.
  15 void __sanitizer_malloc_hook(const volatile void *ptr, size_t sz) {
  16   if (__sanitizer_get_ownership(ptr) && sz == 4) {
  17     WRITE("MallocHook\n");
  18     global_ptr = ptr;
  19   }
  20 }
  21 void __sanitizer_free_hook(const volatile void *ptr) {
  22   if (__sanitizer_get_ownership(ptr) && ptr == global_ptr)
  23     WRITE("FreeHook\n");
  24 }
  25 } // extern "C"
  26
  27 volatile int *x;
  28
  29 void MallocHook1(const volatile void *ptr, size_t sz) { WRITE("MH1\n"); }
  30 void MallocHook2(const volatile void *ptr, size_t sz) { WRITE("MH2\n"); }
  31 void FreeHook1(const volatile void *ptr) { WRITE("FH1\n"); }
  32 void FreeHook2(const volatile void *ptr) { WRITE("FH2\n"); }
  33 // Call this function with uninitialized arguments to poison
  34 // TLS shadow for function parameters before calling operator
  35 // new and, eventually, user-provided hook.
  36 __attribute__((noinline)) void allocate(int *unused1, int *unused2) {
  37   x = new int;
  38 }
  39
  40 int main() {
  41   __sanitizer_install_malloc_and_free_hooks(MallocHook1, FreeHook1);
  42   __sanitizer_install_malloc_and_free_hooks(MallocHook2, FreeHook2);
  43   int *undef1, *undef2;
  44   allocate(undef1, undef2);
  45   // CHECK: MallocHook
  46   // CHECK: MH1
  47   // CHECK: MH2
  48   // Check that malloc hook was called with correct argument.
  49   if (global_ptr != (void *)x) {
  50     _exit(1);
  51   }
  52   *x = 0;
  53   delete x;
  54   // CHECK: FreeHook
  55   // CHECK: FH1
  56   // CHECK: FH2
  57   return 0;
  58 }