[msan] Allocator statistics interface and malloc hooks.

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

5 files changed, 155 insertions, 0 deletions

diff --git a/lib/msan/lit_tests/malloc_hook.cc b/lib/msan/lit_tests/malloc_hook.cc
new file mode 100644
index 000000000..5e7e7dc85
--- /dev/null
+++ b/lib/msan/lit_tests/malloc_hook.cc
@@ -0,0 +1,36 @@
+// RUN: %clangxx_msan -O2 %s -o %t
+// RUN: %t 2>&1 | FileCheck %s
+#include <stdlib.h>
+#include <unistd.h>
+
+extern "C" {
+bool __msan_get_ownership(const void *p);
+
+void *global_ptr;
+
+// Note: avoid calling functions that allocate memory in malloc/free
+// to avoid infinite recursion.
+void __msan_malloc_hook(void *ptr, size_t sz) {
+  if (__msan_get_ownership(ptr)) {
+    write(1, "MallocHook\n", sizeof("MallocHook\n"));
+    global_ptr = ptr;
+  }
+}
+void __msan_free_hook(void *ptr) {
+  if (__msan_get_ownership(ptr) && ptr == global_ptr)
+    write(1, "FreeHook\n", sizeof("FreeHook\n"));
+}
+}  // extern "C"
+
+int main() {
+  volatile int *x = new int;
+  // CHECK: MallocHook
+  // Check that malloc hook was called with correct argument.
+  if (global_ptr != (void*)x) {
+    _exit(1);
+  }
+  *x = 0;
+  delete x;
+  // CHECK: FreeHook
+  return 0;
+}
diff --git a/lib/msan/msan.h b/lib/msan/msan.h
index 1432d99bd..b4bf64253 100644
--- a/lib/msan/msan.h
+++ b/lib/msan/msan.h
@@ -87,4 +87,9 @@ void UnpoisonParam(uptr n);
 
 }  // namespace __msan
 
+#define MSAN_MALLOC_HOOK(ptr, size) \
+  if (&__msan_malloc_hook) __msan_malloc_hook(ptr, size)
+#define MSAN_FREE_HOOK(ptr) \
+  if (&__msan_free_hook) __msan_free_hook(ptr)
+
 #endif  // MSAN_H
diff --git a/lib/msan/msan_allocator.cc b/lib/msan/msan_allocator.cc
index 7435843ce..7b989decb 100644
--- a/lib/msan/msan_allocator.cc
+++ b/lib/msan/msan_allocator.cc
@@ -61,14 +61,17 @@ static void *MsanAllocate(StackTrace *stack, uptr size,
     CHECK_EQ((stack_id >> 31), 0);  // Higher bit is occupied by stack origins.
     __msan_set_origin(res, size, stack_id);
   }
+  MSAN_MALLOC_HOOK(res, size);
   return res;
 }
 
 void MsanDeallocate(void *p) {
   CHECK(p);
   Init();
+  MSAN_FREE_HOOK(p);
   Metadata *meta = reinterpret_cast<Metadata*>(allocator.GetMetaData(p));
   uptr size = meta->requested_size;
+  meta->requested_size = 0;
   // This memory will not be reused by anyone else, so we are free to keep it
   // poisoned.
   __msan_poison(p, size);
@@ -104,4 +107,52 @@ void *MsanReallocate(StackTrace *stack, void *old_p, uptr new_size,
   return new_p;
 }
 
+static uptr AllocationSize(const void *p) {
+  if (p == 0)
+    return 0;
+  const void *beg = allocator.GetBlockBegin(p);
+  if (beg != p)
+    return 0;
+  Metadata *b = (Metadata*)allocator.GetMetaData(p);
+  return b->requested_size;
+}
+
 }  // namespace __msan
+
+using namespace __msan;
+
+uptr __msan_get_current_allocated_bytes() {
+  u64 stats[AllocatorStatCount];
+  allocator.GetStats(stats);
+  u64 m = stats[AllocatorStatMalloced];
+  u64 f = stats[AllocatorStatFreed];
+  return m >= f ? m - f : 1;
+}
+
+uptr __msan_get_heap_size() {
+  u64 stats[AllocatorStatCount];
+  allocator.GetStats(stats);
+  u64 m = stats[AllocatorStatMmapped];
+  u64 f = stats[AllocatorStatUnmapped];
+  return m >= f ? m - f : 1;
+}
+
+uptr __msan_get_free_bytes() {
+  return 1;
+}
+
+uptr __msan_get_unmapped_bytes() {
+  return 1;
+}
+
+uptr __msan_get_estimated_allocated_size(uptr size) {
+  return size;
+}
+
+bool __msan_get_ownership(const void *p) {
+  return AllocationSize(p) != 0;
+}
+
+uptr __msan_get_allocated_size(const void *p) {
+  return AllocationSize(p);
+}
diff --git a/lib/msan/msan_interface_internal.h b/lib/msan/msan_interface_internal.h
index a5502eef1..5d643e862 100644
--- a/lib/msan/msan_interface_internal.h
+++ b/lib/msan/msan_interface_internal.h
@@ -141,6 +141,33 @@ void __sanitizer_unaligned_store32(uu32 *p, u32 x);
 
 SANITIZER_INTERFACE_ATTRIBUTE
 void __sanitizer_unaligned_store64(uu64 *p, u64 x);
+
+SANITIZER_INTERFACE_ATTRIBUTE
+uptr __msan_get_estimated_allocated_size(uptr size);
+
+SANITIZER_INTERFACE_ATTRIBUTE
+bool __msan_get_ownership(const void *p);
+
+SANITIZER_INTERFACE_ATTRIBUTE
+uptr __msan_get_allocated_size(const void *p);
+
+SANITIZER_INTERFACE_ATTRIBUTE
+uptr __msan_get_current_allocated_bytes();
+
+SANITIZER_INTERFACE_ATTRIBUTE
+uptr __msan_get_heap_size();
+
+SANITIZER_INTERFACE_ATTRIBUTE
+uptr __msan_get_free_bytes();
+
+SANITIZER_INTERFACE_ATTRIBUTE
+uptr __msan_get_unmapped_bytes();
+
+SANITIZER_WEAK_ATTRIBUTE SANITIZER_INTERFACE_ATTRIBUTE
+/* OPTIONAL */ void __msan_malloc_hook(void *ptr, uptr size);
+
+SANITIZER_WEAK_ATTRIBUTE SANITIZER_INTERFACE_ATTRIBUTE
+/* OPTIONAL */ void __msan_free_hook(void *ptr);
 }  // extern "C"
 
 #endif  // MSAN_INTERFACE_INTERNAL_H
diff --git a/lib/msan/tests/msan_test.cc b/lib/msan/tests/msan_test.cc
index 42db3697c..520fa7bb4 100644
--- a/lib/msan/tests/msan_test.cc
+++ b/lib/msan/tests/msan_test.cc
@@ -2785,3 +2785,39 @@ TEST(MemorySanitizer, CallocOverflow) {
 TEST(MemorySanitizerStress, DISABLED_MallocStackTrace) {
   RecursiveMalloc(22);
 }
+
+TEST(MemorySanitizerAllocator, get_estimated_allocated_size) {
+  size_t sizes[] = {0, 20, 5000, 1<<20};
+  for (size_t i = 0; i < sizeof(sizes) / sizeof(*sizes); ++i) {
+    size_t alloc_size = __msan_get_estimated_allocated_size(sizes[i]);
+    EXPECT_EQ(alloc_size, sizes[i]);
+  }
+}
+
+TEST(MemorySanitizerAllocator, get_allocated_size_and_ownership) {
+  char *array = reinterpret_cast<char*>(malloc(100));
+  int *int_ptr = new int;
+
+  EXPECT_TRUE(__msan_get_ownership(array));
+  EXPECT_EQ(100, __msan_get_allocated_size(array));
+
+  EXPECT_TRUE(__msan_get_ownership(int_ptr));
+  EXPECT_EQ(sizeof(*int_ptr), __msan_get_allocated_size(int_ptr));
+
+  void *wild_addr = reinterpret_cast<void*>(0x1);
+  EXPECT_FALSE(__msan_get_ownership(wild_addr));
+  EXPECT_EQ(0, __msan_get_allocated_size(wild_addr));
+
+  EXPECT_FALSE(__msan_get_ownership(array + 50));
+  EXPECT_EQ(0, __msan_get_allocated_size(array + 50));
+
+  // NULL is a valid argument for GetAllocatedSize but is not owned.                                                  
+  EXPECT_FALSE(__msan_get_ownership(NULL));
+  EXPECT_EQ(0, __msan_get_allocated_size(NULL));
+ 
+  free(array);
+  EXPECT_FALSE(__msan_get_ownership(array));
+  EXPECT_EQ(0, __msan_get_allocated_size(array));
+
+  delete int_ptr;
+}