AddressSanitizer: add support for malloc_usable_size() function

llvm-svn: 148287

7 files changed, 81 insertions, 23 deletions

diff --git a/compiler-rt/lib/asan/asan_allocator.cc b/compiler-rt/lib/asan/asan_allocator.cc
index e2057dced8c..b5de1f85dc0 100644
--- a/compiler-rt/lib/asan/asan_allocator.cc
+++ b/compiler-rt/lib/asan/asan_allocator.cc
@@ -811,19 +811,46 @@ int asan_posix_memalign(void **memptr, size_t alignment, size_t size,
   return 0;
 }
 
-size_t __asan_mz_size(const void *ptr) {
-  return malloc_info.AllocationSize((uintptr_t)ptr);
+static void GetAllocationSizeAndOwnership(const void *ptr, size_t *size,
+                                          bool *owned) {
+  size_t allocation_size = malloc_info.AllocationSize((uintptr_t)ptr);
+  if (size != NULL) {
+    *size = allocation_size;
+  }
+  if (owned != NULL) {
+    *owned = (ptr == NULL) || (allocation_size > 0);
+  }
+}
+
+size_t asan_malloc_usable_size(void *ptr, AsanStackTrace *stack) {
+  size_t usable_size;
+  bool owned;
+  GetAllocationSizeAndOwnership(ptr, &usable_size, &owned);
+  if (!owned) {
+    Report("ERROR: AddressSanitizer attempting to call malloc_usable_size() "
+           "for pointer which is not owned: %p\n", ptr);
+    stack->PrintStack();
+    Describe((uintptr_t)ptr, 1);
+    ShowStatsAndAbort();
+  }
+  return usable_size;
+}
+
+size_t asan_mz_size(const void *ptr) {
+  size_t mz_size;
+  GetAllocationSizeAndOwnership(ptr, &mz_size, NULL);
+  return mz_size;
 }
 
 void DescribeHeapAddress(uintptr_t addr, uintptr_t access_size) {
   Describe(addr, access_size);
 }
 
-void __asan_mz_force_lock() {
+void asan_mz_force_lock() {
   malloc_info.ForceLock();
 }
 
-void __asan_mz_force_unlock() {
+void asan_mz_force_unlock() {
   malloc_info.ForceUnlock();
 }
 
@@ -999,17 +1026,22 @@ size_t __asan_get_estimated_allocated_size(size_t size) {
 }
 
 bool __asan_get_ownership(const void *p) {
-  return (p == NULL) ||
-      (malloc_info.AllocationSize((uintptr_t)p) > 0);
+  bool owned;
+  GetAllocationSizeAndOwnership(p, NULL, &owned);
+  return owned;
 }
 
 size_t __asan_get_allocated_size(const void *p) {
-  if (p == NULL) return 0;
-  size_t allocated_size = malloc_info.AllocationSize((uintptr_t)p);
+  size_t allocated_size;
+  bool owned;
+  GetAllocationSizeAndOwnership(p, &allocated_size, &owned);
   // Die if p is not malloced or if it is already freed.
-  if (allocated_size == 0) {
-    Printf("__asan_get_allocated_size failed, ptr=%p is not owned\n", p);
+  if (!owned) {
+    Report("ERROR: AddressSanitizer attempting to call "
+           "__asan_get_allocated_size() for pointer which is "
+           "not owned: %p\n", p);
     PRINT_CURRENT_STACK();
+    Describe((uintptr_t)p, 1);
     ShowStatsAndAbort();
   }
   return allocated_size;
diff --git a/compiler-rt/lib/asan/asan_allocator.h b/compiler-rt/lib/asan/asan_allocator.h
index 9b691e04ff4..913da0ed0d1 100644
--- a/compiler-rt/lib/asan/asan_allocator.h
+++ b/compiler-rt/lib/asan/asan_allocator.h
@@ -147,10 +147,11 @@ void *asan_pvalloc(size_t size, AsanStackTrace *stack);
 
 int asan_posix_memalign(void **memptr, size_t alignment, size_t size,
                           AsanStackTrace *stack);
+size_t asan_malloc_usable_size(void *ptr, AsanStackTrace *stack);
 
-size_t __asan_mz_size(const void *ptr);
-void __asan_mz_force_lock();
-void __asan_mz_force_unlock();
+size_t asan_mz_size(const void *ptr);
+void asan_mz_force_lock();
+void asan_mz_force_unlock();
 void DescribeHeapAddress(uintptr_t addr, size_t access_size);
 
 }  // namespace __asan
diff --git a/compiler-rt/lib/asan/asan_linux.cc b/compiler-rt/lib/asan/asan_linux.cc
index 67acd524620..2971290be52 100644
--- a/compiler-rt/lib/asan/asan_linux.cc
+++ b/compiler-rt/lib/asan/asan_linux.cc
@@ -233,7 +233,7 @@ bool AsanProcMaps::GetObjectNameAndOffset(uintptr_t addr, uintptr_t *offset,
   return false;
 }
 
-#else // __arm__
+#else  // __arm__
 
 struct DlIterateData {
   int count;
@@ -279,7 +279,7 @@ bool AsanProcMaps::GetObjectNameAndOffset(uintptr_t addr, uintptr_t *offset,
   return false;
 }
 
-#endif // __arm__
+#endif  // __arm__
 
 void AsanThread::SetThreadStackTopAndBottom() {
   if (tid() == 0) {
diff --git a/compiler-rt/lib/asan/asan_malloc_linux.cc b/compiler-rt/lib/asan/asan_malloc_linux.cc
index 9dbc7a127fc..41a472cad8f 100644
--- a/compiler-rt/lib/asan/asan_malloc_linux.cc
+++ b/compiler-rt/lib/asan/asan_malloc_linux.cc
@@ -108,6 +108,12 @@ void* __libc_memalign(size_t align, size_t s)
   __attribute__((alias("memalign")));
 
 INTERCEPTOR_ATTRIBUTE
+size_t malloc_usable_size(void *ptr) {
+  GET_STACK_TRACE_HERE_FOR_MALLOC;
+  return asan_malloc_usable_size(ptr, &stack);
+}
+
+INTERCEPTOR_ATTRIBUTE
 struct mallinfo mallinfo() {
   struct mallinfo res;
   real_memset(&res, 0, sizeof(res));
diff --git a/compiler-rt/lib/asan/asan_malloc_mac.cc b/compiler-rt/lib/asan/asan_malloc_mac.cc
index 8a6f1bc41ff..3bbacf07b6e 100644
--- a/compiler-rt/lib/asan/asan_malloc_mac.cc
+++ b/compiler-rt/lib/asan/asan_malloc_mac.cc
@@ -69,7 +69,7 @@ size_t mz_size(malloc_zone_t* zone, const void* ptr) {
   if (system_malloc_zone) {
     if ((system_malloc_zone->size)(system_malloc_zone, ptr)) return 0;
   }
-  return __asan_mz_size(ptr);
+  return asan_mz_size(ptr);
 }
 
 void *mz_malloc(malloc_zone_t *zone, size_t size) {
@@ -141,7 +141,7 @@ void mz_free(malloc_zone_t *zone, void *ptr) {
     system_purgeable_zone->free(system_purgeable_zone, ptr);
     return;
   }
-  if (__asan_mz_size(ptr)) {
+  if (asan_mz_size(ptr)) {
     GET_STACK_TRACE_HERE_FOR_FREE(ptr);
     asan_free(ptr, &stack);
   } else {
@@ -165,7 +165,7 @@ void cf_free(void *ptr, void *info) {
     system_purgeable_zone->free(system_purgeable_zone, ptr);
     return;
   }
-  if (__asan_mz_size(ptr)) {
+  if (asan_mz_size(ptr)) {
     GET_STACK_TRACE_HERE_FOR_FREE(ptr);
     asan_free(ptr, &stack);
   } else {
@@ -184,7 +184,7 @@ void *mz_realloc(malloc_zone_t *zone, void *ptr, size_t size) {
     GET_STACK_TRACE_HERE_FOR_MALLOC;
     return asan_malloc(size, &stack);
   } else {
-    if (__asan_mz_size(ptr)) {
+    if (asan_mz_size(ptr)) {
       GET_STACK_TRACE_HERE_FOR_MALLOC;
       return asan_realloc(ptr, size, &stack);
     } else {
@@ -207,7 +207,7 @@ void *cf_realloc(void *ptr, CFIndex size, CFOptionFlags hint, void *info) {
     GET_STACK_TRACE_HERE_FOR_MALLOC;
     return asan_malloc(size, &stack);
   } else {
-    if (__asan_mz_size(ptr)) {
+    if (asan_mz_size(ptr)) {
       GET_STACK_TRACE_HERE_FOR_MALLOC;
       return asan_realloc(ptr, size, &stack);
     } else {
@@ -279,11 +279,11 @@ void mi_log(malloc_zone_t *zone, void *address) {
 }
 
 void mi_force_lock(malloc_zone_t *zone) {
-  __asan_mz_force_lock();
+  asan_mz_force_lock();
 }
 
 void mi_force_unlock(malloc_zone_t *zone) {
-  __asan_mz_force_unlock();
+  asan_mz_force_unlock();
 }
 
 // This function is currently unused, and we build with -Werror.
diff --git a/compiler-rt/lib/asan/tests/asan_interface_test.cc b/compiler-rt/lib/asan/tests/asan_interface_test.cc
index c26ed92468b..521853ec756 100644
--- a/compiler-rt/lib/asan/tests/asan_interface_test.cc
+++ b/compiler-rt/lib/asan/tests/asan_interface_test.cc
@@ -27,7 +27,7 @@ TEST(AddressSanitizerInterface, GetEstimatedAllocatedSize) {
 }
 
 static const char* kGetAllocatedSizeErrorMsg =
-  "__asan_get_allocated_size failed";
+  "attempting to call __asan_get_allocated_size()";
 
 TEST(AddressSanitizerInterface, GetAllocatedSizeAndOwnershipTest) {
   const size_t kArraySize = 100;
diff --git a/compiler-rt/lib/asan/tests/asan_test.cc b/compiler-rt/lib/asan/tests/asan_test.cc
index 9c6144b92a2..363bb893e98 100644
--- a/compiler-rt/lib/asan/tests/asan_test.cc
+++ b/compiler-rt/lib/asan/tests/asan_test.cc
@@ -520,6 +520,25 @@ TEST(AddressSanitizer, ReallocTest) {
   }
 }
 
+#ifndef __APPLE__
+static const char *kMallocUsableSizeErrorMsg =
+  "AddressSanitizer attempting to call malloc_usable_size()";
+
+TEST(AddressSanitizer, MallocUsableSizeTest) {
+  const size_t kArraySize = 100;
+  char *array = Ident((char*)malloc(kArraySize));
+  int *int_ptr = Ident(new int);
+  EXPECT_EQ(0, malloc_usable_size(NULL));
+  EXPECT_EQ(kArraySize, malloc_usable_size(array));
+  EXPECT_EQ(sizeof(int), malloc_usable_size(int_ptr));
+  EXPECT_DEATH(malloc_usable_size((void*)0x123), kMallocUsableSizeErrorMsg);
+  EXPECT_DEATH(malloc_usable_size(array + kArraySize / 2),
+               kMallocUsableSizeErrorMsg);
+  free(array);
+  EXPECT_DEATH(malloc_usable_size(array), kMallocUsableSizeErrorMsg);
+}
+#endif
+
 void WrongFree() {
   int *x = (int*)malloc(100 * sizeof(int));
   // Use the allocated memory, otherwise Clang will optimize it out.