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