[asan] Provide an interface to update an allocation stack trace.

[eugenis]

Summary

Sometimes an allocation stack trace is not very informative. Provide a
way to replace it with a stack trace of the user's choice.

[dvyukov]

FWIW for kernel sanitizers we want an ability to memorize few _additional_ stacks per heap object:
https://bugzilla.kernel.org/show_bug.cgi?id=198437
That's very useful for any kind of asynchronous processing environments that involves tasks, callbacks, thread pools, etc (read - almost all large C/C++ software today) because allocation/free stack may be meaningless and/or other stacks may be crucially important.

[kcc]

compiler-rt/include/sanitizer/asan_interface.h

320

Add some explanation of when it is going to be successful.
Also, what about threads?
This API is racey

[eugenis]

compiler-rt/include/sanitizer/asan_interface.h

320

Racey, as in someone could be deallocating this memory on another thread?
This API should be used in a situation when this is possible, of course.
It's exactly the same as malloc_usable_size().

[kcc]

compiler-rt/include/sanitizer/asan_interface.h

320

malloc_usable_size() only reads the data.
__asan_update_allocation writes the data, so if two threads try to simultaneously call this function, we'll have a race on alloc_context_id

BTW, the name "asan_update_allocation" doesn't seem to reflect the action.
Maybe "asan_update_allocation_context"?

[eugenis]

Use atomic to update the allocation stack.
Rename the function.

[eugenis]

In D69208#1716241, @dvyukov wrote:

FWIW for kernel sanitizers we want an ability to memorize few _additional_ stacks per heap object:
https://bugzilla.kernel.org/show_bug.cgi?id=198437
That's very useful for any kind of asynchronous processing environments that involves tasks, callbacks, thread pools, etc (read - almost all large C/C++ software today) because allocation/free stack may be meaningless and/or other stacks may be crucially important.

Agreed. We don't have space in the chunk header to fit any additional stack traces for free though. And the kernel has its own runtime implementation anyway.

compiler-rt/include/sanitizer/asan_interface.h

320

OK, I made the store atomic. It could still race with deallocation, but the same is true for malloc_usable_size.

Changed the name to asan_update_allocation_context

[eugenis]

Btw, we can implement this really nicely in MSan with origin tracking.
And if we really wanted, we could pull the chained origin stack depot code to sanitizer_common and use it in asan for the extra stack traces.
But this sounds like a lot of work for a fringe feature.

[kcc]

LGTM

compiler-rt/include/sanitizer/asan_interface.h

 
 /// Performs shadow memory cleanup of the current thread's stack before a
 /// function marked with the <c>[[noreturn]]</c> attribute is called.
 ///
 /// To avoid false positives on the stack, must be called before no-return
 /// functions like <c>_exit()</c> and <c>execl()</c>.
 void __asan_handle_no_return(void);
 
+/// Update allocation stack trace for the given allocation to the current stack
+/// trace. Returns 1 if successfull, 0 if not.
+int __asan_update_allocation_context(void* addr);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif
 
 #endif  // SANITIZER_ASAN_INTERFACE_H

compiler-rt/lib/asan/asan_allocator.cpp

AsanChunk *ChooseChunk(uptr addr, AsanChunk *left_chunk,
     sptr l_offset = 0, r_offset = 0;
     CHECK(AsanChunkView(left_chunk).AddrIsAtRight(addr, 1, &l_offset));
     CHECK(AsanChunkView(right_chunk).AddrIsAtLeft(addr, 1, &r_offset));
     if (l_offset < r_offset)
       return left_chunk;
     return right_chunk;
   }
 
+  bool UpdateAllocationStack(uptr addr, BufferedStackTrace *stack) {
+    AsanChunk *m = GetAsanChunkByAddr(addr);
+    if (!m) return false;
+    if (m->chunk_state != CHUNK_ALLOCATED) return false;
+    if (m->Beg() != addr) return false;
+    atomic_store((atomic_uint32_t *)&m->alloc_context_id, StackDepotPut(*stack),
+                 memory_order_relaxed);
+    return true;
+  }
+
   // -------------------- Allocation/Deallocation routines ---------------
   void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
                  AllocType alloc_type, bool can_fill) {
     if (UNLIKELY(!asan_inited))
       AsanInitFromRtl();
     if (RssLimitExceeded()) {
       if (AllocatorMayReturnNull())
         return nullptr;
uptr __sanitizer_get_allocated_size(const void *p) {
   return allocated_size;
 }
 
 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);
+}
+
 #if !SANITIZER_SUPPORTS_WEAK_HOOKS
 // Provide default (no-op) implementation of malloc hooks.
 SANITIZER_INTERFACE_WEAK_DEF(void, __sanitizer_malloc_hook,
                              void *ptr, uptr size) {
   (void)ptr;
   (void)size;
 }
 
 SANITIZER_INTERFACE_WEAK_DEF(void, __sanitizer_free_hook, void *ptr) {
   (void)ptr;
 }
 #endif

compiler-rt/lib/asan/asan_interface.inc

 INTERFACE_FUNCTION(__sanitizer_ptr_sub)
 INTERFACE_FUNCTION(__sanitizer_start_switch_fiber)
 INTERFACE_FUNCTION(__sanitizer_unaligned_load16)
 INTERFACE_FUNCTION(__sanitizer_unaligned_load32)
 INTERFACE_FUNCTION(__sanitizer_unaligned_load64)
 INTERFACE_FUNCTION(__sanitizer_unaligned_store16)
 INTERFACE_FUNCTION(__sanitizer_unaligned_store32)
 INTERFACE_FUNCTION(__sanitizer_unaligned_store64)
+INTERFACE_FUNCTION(__asan_update_allocation_context)
 INTERFACE_WEAK_FUNCTION(__asan_default_options)
 INTERFACE_WEAK_FUNCTION(__asan_default_suppressions)
 INTERFACE_WEAK_FUNCTION(__asan_on_error)

compiler-rt/lib/asan/asan_interface_internal.h

extern "C" {
   void __asan_alloca_poison(uptr addr, uptr size);
   SANITIZER_INTERFACE_ATTRIBUTE
   void __asan_allocas_unpoison(uptr top, uptr bottom);
 
   SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE
   const char* __asan_default_suppressions();
 
   SANITIZER_INTERFACE_ATTRIBUTE void __asan_handle_vfork(void *sp);
+
+  SANITIZER_INTERFACE_ATTRIBUTE int __asan_update_allocation_context(void* addr);
 }  // extern "C"
 
 #endif  // ASAN_INTERFACE_INTERNAL_H

compiler-rt/test/asan/TestCases/asan_update_allocation.cpp

+// RUN: %clangxx_asan -O0 -DSIZE=10 %s -o %t && not %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-%os --check-prefix=CHECK
+// RUN: %clangxx_asan -O0 -DSIZE=10000000 %s -o %t && not %run %t 2>&1 | FileCheck %s --check-prefix=CHECK-%os --check-prefix=CHECK
+// REQUIRES: stable-runtime
+
+#include <stdlib.h>
+#include <sanitizer/asan_interface.h>
+
+void UPDATE(void *p) {
+  __asan_update_allocation_context(p);
+}
+
+int main() {
+  char *x = (char*)malloc(SIZE * sizeof(char));
+  UPDATE(x);
+  free(x);
+  return x[5];
+  // CHECK: {{.*ERROR: AddressSanitizer: heap-use-after-free on address}}
+  // CHECK: UPDATE
+}