[sanitizer] Implement __sanitizer_get_allocated_size_fast

[Northbadge]

Summary

The primary motivation for this change is to allow FreeHooks to obtain
the allocated size of the pointer being freed in a fast, efficient manner.

[Northbadge]

Light cleanup

[vitalybuka]

Thanks, LGTM. Could you please update also compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp?

[Northbadge]

In D151360#4370274, @vitalybuka wrote:

Thanks, LGTM. Could you please update also compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp?

Thanks for the review! I don't see a compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp though?

[Northbadge]

Add missing semicolons

[Northbadge]

Synced to tip; Found the new test & updated it. Will let build bots run for a bit.

compiler-rt/include/sanitizer/allocator_interface.h

#endif
   /* If a pointer lies within an allocation, it will return the start address
      of the allocation. Otherwise, it returns nullptr. */
   const void *__sanitizer_get_allocated_begin(const void *p);
 
   /* Returns the number of bytes reserved for the pointer p.
      Requires (get_ownership(p) == true) or (p == 0). */
   size_t __sanitizer_get_allocated_size(const volatile void *p);
 
+  /* Returns the number of bytes reserved for the pointer p.
+     Requires __sanitizer_get_allocated_begin(p) == p. */
+  size_t __sanitizer_get_allocated_size_fast(const volatile void *p);
+
   /* Number of bytes, allocated and not yet freed by the application. */
   size_t __sanitizer_get_current_allocated_bytes(void);
 
   /* Number of bytes, mmaped by the allocator to fulfill allocation requests.
      Generally, for request of X bytes, allocator can reserve and add to free
      lists a large number of chunks of size X to use them for future requests.
      All these chunks count toward the heap size. Currently, allocator never
      releases memory to OS (instead, it just puts freed chunks to free

compiler-rt/lib/asan/asan_allocator.cpp

uptr AllocationSize(uptr p) {
     AsanChunk *m = GetAsanChunkByAddr(p);
     if (!m) return 0;
     if (atomic_load(&m->chunk_state, memory_order_acquire) != CHUNK_ALLOCATED)
       return 0;
     if (m->Beg() != p) return 0;
     return m->UsedSize();
   }
 
+  uptr AllocationSizeFast(uptr p) {
+    return reinterpret_cast<AsanChunk *>(p - kChunkHeaderSize)->UsedSize();
+  }
+
   AsanChunkView FindHeapChunkByAddress(uptr addr) {
     AsanChunk *m1 = GetAsanChunkByAddr(addr);
     sptr offset = 0;
     if (!m1 || AsanChunkView(m1).AddrIsAtLeft(addr, 1, &offset)) {
       // The address is in the chunk's left redzone, so maybe it is actually
       // a right buffer overflow from the other chunk before.
       // Search a bit before to see if there is another chunk.
       AsanChunk *m2 = nullptr;
uptr __sanitizer_get_allocated_size(const void *p) {
   // Die if p is not malloced or if it is already freed.
   if (allocated_size == 0) {
     GET_STACK_TRACE_FATAL_HERE;
     ReportSanitizerGetAllocatedSizeNotOwned(ptr, &stack);
   }
   return allocated_size;
 }
 
+uptr __sanitizer_get_allocated_size_fast(const void *p) {
+  DCHECK_EQ(p, __sanitizer_get_allocated_begin(p));
+  uptr ret = instance.AllocationSizeFast(reinterpret_cast<uptr>(p));
+  DCHECK_EQ(ret, __sanitizer_get_allocated_size(p));
+  return ret;
+}
+
 const void *__sanitizer_get_allocated_begin(const void *p) {
   return AllocationBegin(p);
 }
 
 void __sanitizer_purge_allocator() {
   GET_STACK_TRACE_MALLOC;
   instance.Purge(&stack);
 }
 
 int __asan_update_allocation_context(void* addr) {
   GET_STACK_TRACE_MALLOC;
   return instance.UpdateAllocationStack((uptr)addr, &stack);
 }

compiler-rt/lib/dfsan/dfsan_allocator.cpp

if (!p)
     return 0;
   const void *beg = allocator.GetBlockBegin(p);
   if (beg != p)
     return 0;
   Metadata *b = (Metadata *)allocator.GetMetaData(p);
   return b->requested_size;
 }
 
+static uptr AllocationSizeFast(const void *p) {
+  return reinterpret_cast<Metadata *>(allocator.GetMetaData(p))->requested_size;
+}
+
 void *dfsan_malloc(uptr size) {
   return SetErrnoOnNull(DFsanAllocate(size, sizeof(u64), false /*zeroise*/));
 }
 
 void *dfsan_calloc(uptr nmemb, uptr size) {
   return SetErrnoOnNull(DFsanCalloc(nmemb, size));
 }
 

 
 int __sanitizer_get_ownership(const void *p) { return AllocationSize(p) != 0; }
 
 const void *__sanitizer_get_allocated_begin(const void *p) {
   return AllocationBegin(p);
 }
 
 uptr __sanitizer_get_allocated_size(const void *p) { return AllocationSize(p); }
+
+uptr __sanitizer_get_allocated_size_fast(const void *p) {
+  DCHECK_EQ(p, __sanitizer_get_allocated_begin(p));
+  uptr ret = AllocationSizeFast(p);
+  DCHECK_EQ(ret, __sanitizer_get_allocated_size(p));
+  return ret;
+}

compiler-rt/lib/hwasan/hwasan_allocator.cpp

static uptr AllocationSize(const void *p) {
   if (!untagged_ptr) return 0;
   const void *beg = allocator.GetBlockBegin(untagged_ptr);
   if (!beg)
     return 0;
   Metadata *b = (Metadata *)allocator.GetMetaData(beg);
   return b->GetRequestedSize();
 }
 
+static uptr AllocationSizeFast(const void *p) {
+  const void *untagged_ptr = UntagPtr(p);
+  void *aligned_ptr = reinterpret_cast<void *>(
+      RoundDownTo(reinterpret_cast<uptr>(untagged_ptr), kShadowAlignment));
+  Metadata *meta =
+      reinterpret_cast<Metadata *>(allocator.GetMetaData(aligned_ptr));
+  return meta->GetRequestedSize();
+}
+
 void *hwasan_malloc(uptr size, StackTrace *stack) {
   return SetErrnoOnNull(HwasanAllocate(stack, size, sizeof(u64), false));
 }
 
 void *hwasan_calloc(uptr nmemb, uptr size, StackTrace *stack) {
   return SetErrnoOnNull(HwasanCalloc(stack, nmemb, size));
 }
 

 int __sanitizer_get_ownership(const void *p) { return AllocationSize(p) != 0; }
 
 const void *__sanitizer_get_allocated_begin(const void *p) {
   return AllocationBegin(p);
 }
 
 uptr __sanitizer_get_allocated_size(const void *p) { return AllocationSize(p); }
 
+uptr __sanitizer_get_allocated_size_fast(const void *p) {
+  DCHECK_EQ(p, __sanitizer_get_allocated_begin(p));
+  uptr ret = AllocationSizeFast(p);
+  DCHECK_EQ(ret, __sanitizer_get_allocated_size(p));
+  return ret;
+}
+
 void __sanitizer_purge_allocator() { allocator.ForceReleaseToOS(); }

compiler-rt/lib/lsan/lsan_allocator.cpp

 uptr GetMallocUsableSize(const void *p) {
   if (!p)
     return 0;
   ChunkMetadata *m = Metadata(p);
   if (!m) return 0;
   return m->requested_size;
 }
 
+uptr GetMallocUsableSizeFast(const void *p) {
+  return Metadata(p)->requested_size;
+}
+
 int lsan_posix_memalign(void **memptr, uptr alignment, uptr size,
                         const StackTrace &stack) {
   if (UNLIKELY(!CheckPosixMemalignAlignment(alignment))) {
     if (AllocatorMayReturnNull())
       return errno_EINVAL;
     ReportInvalidPosixMemalignAlignment(alignment, &stack);
   }
   void *ptr = Allocate(stack, size, alignment, kAlwaysClearMemory);

 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
 uptr __sanitizer_get_allocated_size(const void *p) {
   return GetMallocUsableSize(p);
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
+uptr __sanitizer_get_allocated_size_fast(const void *p) {
+  DCHECK_EQ(p, __sanitizer_get_allocated_begin(p));
+  uptr ret = GetMallocUsableSizeFast(p);
+  DCHECK_EQ(ret, __sanitizer_get_allocated_size(p));
+  return ret;
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE
 void __sanitizer_purge_allocator() { allocator.ForceReleaseToOS(); }
 
 } // extern "C"

compiler-rt/lib/memprof/memprof_allocator.cpp

uptr AllocationSize(uptr p) {
     MemprofChunk *m = GetMemprofChunkByAddr(p, user_requested_size);
     if (!m)
       return 0;
     if (m->Beg() != p)
       return 0;
     return user_requested_size;
   }
 
+  uptr AllocationSizeFast(uptr p) {
+    return reinterpret_cast<MemprofChunk *>(p - kChunkHeaderSize)->UsedSize();
+  }
+
   void Purge(BufferedStackTrace *stack) { allocator.ForceReleaseToOS(); }
 
   void PrintStats() { allocator.PrintStats(); }
 
   void ForceLock() SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
     allocator.ForceLock();
     fallback_mutex.Lock();
   }

 const void *__sanitizer_get_allocated_begin(const void *p) {
   return memprof_malloc_begin(p);
 }
 
 uptr __sanitizer_get_allocated_size(const void *p) {
   return memprof_malloc_usable_size(p, 0, 0);
 }
 
+uptr __sanitizer_get_allocated_size_fast(const void *p) {
+  DCHECK_EQ(p, __sanitizer_get_allocated_begin(p));
+  uptr ret = instance.AllocationSizeFast(reinterpret_cast<uptr>(p));
+  DCHECK_EQ(ret, __sanitizer_get_allocated_size(p));
+  return ret;
+}
+
 int __memprof_profile_dump() {
   instance.FinishAndWrite();
   // In the future we may want to return non-zero if there are any errors
   // detected during the dumping process.
   return 0;
 }

compiler-rt/lib/msan/msan_allocator.cpp

 static uptr AllocationSize(const void *p) {
   if (!p) return 0;
   const void *beg = allocator.GetBlockBegin(p);
   if (beg != p) return 0;
   Metadata *b = (Metadata *)allocator.GetMetaData(p);
   return b->requested_size;
 }
 
+static uptr AllocationSizeFast(const void *p) {
+  return reinterpret_cast<Metadata *>(allocator.GetMetaData(p))->requested_size;
+}
+
 void *msan_malloc(uptr size, StackTrace *stack) {
   return SetErrnoOnNull(MsanAllocate(stack, size, sizeof(u64), false));
 }
 
 void *msan_calloc(uptr nmemb, uptr size, StackTrace *stack) {
   return SetErrnoOnNull(MsanCalloc(stack, nmemb, size));
 }
 

 int __sanitizer_get_ownership(const void *p) { return AllocationSize(p) != 0; }
 
 const void *__sanitizer_get_allocated_begin(const void *p) {
   return AllocationBegin(p);
 }
 
 uptr __sanitizer_get_allocated_size(const void *p) { return AllocationSize(p); }
 
+uptr __sanitizer_get_allocated_size_fast(const void *p) {
+  DCHECK_EQ(p, __sanitizer_get_allocated_begin(p));
+  uptr ret = AllocationSizeFast(p);
+  DCHECK_EQ(ret, __sanitizer_get_allocated_size(p));
+  return ret;
+}
+
 void __sanitizer_purge_allocator() { allocator.ForceReleaseToOS(); }

compiler-rt/lib/sanitizer_common/sanitizer_allocator_interface.h

 extern "C" {
 SANITIZER_INTERFACE_ATTRIBUTE
 uptr __sanitizer_get_estimated_allocated_size(uptr size);
 SANITIZER_INTERFACE_ATTRIBUTE int __sanitizer_get_ownership(const void *p);
 SANITIZER_INTERFACE_ATTRIBUTE const void *__sanitizer_get_allocated_begin(
     const void *p);
 SANITIZER_INTERFACE_ATTRIBUTE uptr
 __sanitizer_get_allocated_size(const void *p);
+SANITIZER_INTERFACE_ATTRIBUTE uptr
+__sanitizer_get_allocated_size_fast(const void *p);
 SANITIZER_INTERFACE_ATTRIBUTE uptr __sanitizer_get_current_allocated_bytes();
 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 *));

compiler-rt/lib/sanitizer_common/sanitizer_common_interface.inc

 INTERFACE_WEAK_FUNCTION(__sanitizer_weak_hook_strstr)
 // Stacktrace interface.
 INTERFACE_FUNCTION(__sanitizer_get_module_and_offset_for_pc)
 INTERFACE_FUNCTION(__sanitizer_symbolize_global)
 INTERFACE_FUNCTION(__sanitizer_symbolize_pc)
 // Allocator interface.
 INTERFACE_FUNCTION(__sanitizer_get_allocated_begin)
 INTERFACE_FUNCTION(__sanitizer_get_allocated_size)
+INTERFACE_FUNCTION(__sanitizer_get_allocated_size_fast)
 INTERFACE_FUNCTION(__sanitizer_get_current_allocated_bytes)
 INTERFACE_FUNCTION(__sanitizer_get_estimated_allocated_size)
 INTERFACE_FUNCTION(__sanitizer_get_free_bytes)
 INTERFACE_FUNCTION(__sanitizer_get_heap_size)
 INTERFACE_FUNCTION(__sanitizer_get_ownership)
 INTERFACE_FUNCTION(__sanitizer_get_unmapped_bytes)
 INTERFACE_FUNCTION(__sanitizer_install_malloc_and_free_hooks)
 INTERFACE_FUNCTION(__sanitizer_purge_allocator)
 INTERFACE_FUNCTION(__sanitizer_print_memory_profile)
 INTERFACE_WEAK_FUNCTION(__sanitizer_free_hook)
 INTERFACE_WEAK_FUNCTION(__sanitizer_malloc_hook)

compiler-rt/lib/tsan/rtl/tsan_mman.cpp

uptr user_alloc_usable_size(const void *p) {
   MBlock *b = ctx->metamap.GetBlock((uptr)p);
   if (!b)
     return 0;  // Not a valid pointer.
   if (b->siz == 0)
     return 1;  // Zero-sized allocations are actually 1 byte.
   return b->siz;
 }
 
+uptr user_alloc_usable_size_fast(const void *p) {
+  MBlock *b = ctx->metamap.GetBlock((uptr)p);
+  if (b->siz == 0)
+    return 1;  // Zero-sized allocations are actually 1 byte.
+  return b->siz;
+}
+
 void invoke_malloc_hook(void *ptr, uptr size) {
   ThreadState *thr = cur_thread();
   if (ctx == 0 || !ctx->initialized || thr->ignore_interceptors)
     return;
   RunMallocHooks(ptr, size);
 }
 
 void invoke_free_hook(void *ptr) {

 const void *__sanitizer_get_allocated_begin(const void *p) {
   return user_alloc_begin(p);
 }
 
 uptr __sanitizer_get_allocated_size(const void *p) {
   return user_alloc_usable_size(p);
 }
 
+uptr __sanitizer_get_allocated_size_fast(const void *p) {
+  DCHECK_EQ(p, __sanitizer_get_allocated_begin(p));
+  uptr ret = user_alloc_usable_size_fast(p);
+  DCHECK_EQ(ret, __sanitizer_get_allocated_size(p));
+  return ret;
+}
+
 void __sanitizer_purge_allocator() {
   allocator()->ForceReleaseToOS();
 }
 
 void __tsan_on_thread_idle() {
   ThreadState *thr = cur_thread();
   allocator()->SwallowCache(&thr->proc()->alloc_cache);
   internal_allocator()->SwallowCache(&thr->proc()->internal_alloc_cache);
   ctx->metamap.OnProcIdle(thr->proc());
 }
 }  // extern "C"

compiler-rt/test/sanitizer_common/TestCases/Linux/malloc_usable_size.c

 
 int main() {
   assert(__sanitizer_get_allocated_size(NULL) == 0);
   assert(malloc_usable_size(NULL) == 0);
 
   int size = 1;
   p = malloc(size);
   assert(__sanitizer_get_allocated_size(p) == size);
+  assert(__sanitizer_get_allocated_size_fast(p) == size);
   assert(malloc_usable_size(p) == size);
   free(p);
 
   size = 1234567;
   p = malloc(size);
   assert(__sanitizer_get_allocated_size(p) == size);
+  assert(__sanitizer_get_allocated_size_fast(p) == size);
   assert(malloc_usable_size(p) == size);
   free(p);
   return 0;
 }

compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp

 
 void Test(int size) {
   auto allocated_bytes_before = __sanitizer_get_current_allocated_bytes();
   int *p = (int *)malloc(size);
   assert(__sanitizer_get_estimated_allocated_size(size) >= size);
   assert(__sanitizer_get_ownership(p));
   assert(!__sanitizer_get_ownership(&p));
   assert(__sanitizer_get_allocated_size(p) == size);
+  assert(__sanitizer_get_allocated_size_fast(p) == size);
   assert(__sanitizer_get_allocated_begin(p) == p);
   assert(__sanitizer_get_allocated_begin(p + 1) == p);
   assert(__sanitizer_get_current_allocated_bytes() >=
          size + allocated_bytes_before);
   assert(__sanitizer_get_current_allocated_bytes() <=
          2 * size + allocated_bytes_before);
   assert(__sanitizer_get_heap_size() >= size);
   free(p);

compiler-rt/test/sanitizer_common/TestCases/malloc_hook_get_allocated_size_fast.cpp

+// RUN: %clangxx -O2 %s -o %t && %run %t 2>&1
+
+// Malloc/free hooks are not supported on Windows.
+// XFAIL: target={{.*windows-msvc.*}}
+
+// Must not be implemented, no other reason to install interceptors.
+// XFAIL: ubsan
+
+// FIXME: Implement.
+// XFAIL: hwasan
+
+#include <assert.h>
+#include <sanitizer/allocator_interface.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+extern "C" {
+const volatile void *global_ptr;
+
+// 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 == sizeof(int)) {
+    global_ptr = ptr;
+    assert(__sanitizer_get_allocated_size_fast(ptr) == sizeof(int));
+  }
+}
+void __sanitizer_free_hook(const volatile void *ptr) {
+  if (__sanitizer_get_ownership(ptr) && ptr == global_ptr)
+    assert(__sanitizer_get_allocated_size_fast(ptr) == sizeof(int));
+}
+} // extern "C"
+
+volatile int *x;
+
+// Call this function with uninitialized arguments to poison
+// TLS shadow for function parameters before calling operator
+// new and, eventually, user-provided hook.
+__attribute__((noinline)) void allocate(int *unused1, int *unused2) {
+  x = reinterpret_cast<int *>(malloc(sizeof(int)));
+}
+
+int main() {
+  int *undef1, *undef2;
+  allocate(undef1, undef2);
+
+  // Check that malloc hook was called with correct argument.
+  if (global_ptr != (void *)x) {
+    _exit(1);
+  }
+
+  *x = -8;
+  free((void *)x);
+
+  return 0;
+}