[Asan] Fix __asan_update_allocation_context

[vitalybuka]

Summary

Update both thread and stack.
Update thread and stack as atomic operation.
Keep all 32bit of TID as now we have enough bits.

Depends on D87135.

[morehouse]

compiler-rt/lib/asan/asan_allocator.cpp

89–90

Nit: Since we're adding member functions and private members, let's change struct to class.

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

7

Looks like CHECK-%os is unused, let's remove it.

30

s/=/-

30

Why do we expect the threaded test to not give the allocation thread?

[vitalybuka]

update

[vitalybuka]

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

30

I tried to check that this is not main thread to avoid assuming thread numbering
Probably unnecessary.

[morehouse]

LGTM

compiler-rt/lib/asan/asan_allocator.h

 namespace __asan {
 
 enum AllocType {
   FROM_MALLOC = 1,  // Memory block came from malloc, calloc, realloc, etc.
   FROM_NEW = 2,     // Memory block came from operator new.
   FROM_NEW_BR = 3   // Memory block came from operator new [ ]
 };
 
-struct AsanChunk;
+class AsanChunk;
 
 struct AllocatorOptions {
   u32 quarantine_size_mb;
   u32 thread_local_quarantine_size_kb;
   u16 min_redzone;
   u16 max_redzone;
   u8 may_return_null;
   u8 alloc_dealloc_mismatch;

compiler-rt/lib/asan/asan_allocator.cpp

static u32 RZSize2Log(u32 rz_size) {
   CHECK(IsPowerOfTwo(rz_size));
   u32 res = Log2(rz_size) - 4;
   CHECK_EQ(rz_size, RZLog2Size(res));
   return res;
 }
 
 static AsanAllocator &get_allocator();
 
+static void AtomicContextStore(volatile atomic_uint64_t *atomic_context,
+                               u32 tid, u32 stack) {
+  u64 context = tid;
+  context <<= 32;
+  context += stack;
+  atomic_store(atomic_context, context, memory_order_relaxed);
+}
+
+static void AtomicContextLoad(const volatile atomic_uint64_t *atomic_context,
+                              u32 &tid, u32 &stack) {
+  u64 context = atomic_load(atomic_context, memory_order_relaxed);
+  stack = context;
+  context >>= 32;
+  tid = context;
+}
+
 // The memory chunk allocated from the underlying allocator looks like this:
 // L L L L L L H H U U U U U U R R
 //   L -- left redzone words (0 or more bytes)
 //   H -- ChunkHeader (16 bytes), which is also a part of the left redzone.
 //   U -- user memory.
 //   R -- right redzone (0 or more bytes)
 // ChunkBase consists of ChunkHeader and other bytes that overlap with user
 // memory.
 
 // If the left redzone is greater than the ChunkHeader size we store a magic
 // value in the first uptr word of the memory block and store the address of
 // ChunkBase in the next uptr.
 // M B L L L L L L L L L  H H U U U U U U
 //   |                    ^
 //   ---------------------|
 //   M -- magic value kAllocBegMagic
 //   B -- address of ChunkHeader pointing to the first 'H'
 static const uptr kAllocBegMagic = 0xCC6E96B9;
 
-struct ChunkHeader {
+class ChunkHeader {
+ public:
   atomic_uint8_t chunk_state;
   u8 from_memalign : 1;
   u8 alloc_type : 2;
   u8 rz_log : 3;
   u8 lsan_tag : 2;
+
   // This field is used for small sizes. For large sizes it is equal to
   // SizeClassMap::kMaxSize and the actual size is stored in the
   // SecondaryAllocator's metadata.
   u32 user_requested_size : 29;
   // align < 8 -> 0
   // else      -> log2(min(align, 512)) - 2
   u32 user_requested_alignment_log : 3;
-  u32 alloc_tid;
-  atomic_uint32_t alloc_context_id;
+
+ private:
+  atomic_uint64_t alloc_context_id;
+
+ public:
+  void SetAllocContext(u32 tid, u32 stack) {
+    AtomicContextStore(&alloc_context_id, tid, stack);
+  }
+
+  void GetAllocContext(u32 &tid, u32 &stack) const {
+    AtomicContextLoad(&alloc_context_id, tid, stack);
+  }
 };
 
-struct ChunkBase : ChunkHeader {
-  // Header2, intersects with user memory.
-  u32 free_context_id;
-  u32 free_tid;
+class ChunkBase : public ChunkHeader {
+  atomic_uint64_t free_context_id;
+
+ public:
+  void SetFreeContext(u32 tid, u32 stack) {
+    AtomicContextStore(&free_context_id, tid, stack);
+  }
+
+  void GetFreeContext(u32 &tid, u32 &stack) const {
+    AtomicContextLoad(&free_context_id, tid, stack);
+  }
 };
 
 static const uptr kChunkHeaderSize = sizeof(ChunkHeader);
 static const uptr kChunkHeader2Size = sizeof(ChunkBase) - kChunkHeaderSize;
 COMPILER_CHECK(kChunkHeaderSize == 16);
 COMPILER_CHECK(kChunkHeader2Size <= 16);
 
 enum {
   // Either just allocated by underlying allocator, but AsanChunk is not yet
   // ready, or almost returned to undelying allocator and AsanChunk is already
   // meaningless.
   CHUNK_INVALID = 0,
   // The chunk is allocated and not yet freed.
   CHUNK_ALLOCATED = 2,
   // The chunk was freed and put into quarantine zone.
   CHUNK_QUARANTINE = 3,
 };
 
-struct AsanChunk: ChunkBase {
+class AsanChunk : public ChunkBase {
+ public:
   uptr Beg() { return reinterpret_cast<uptr>(this) + kChunkHeaderSize; }
   uptr UsedSize(bool locked_version = false) {
     if (user_requested_size != SizeClassMap::kMaxSize)
       return user_requested_size;
     return *reinterpret_cast<uptr *>(
         get_allocator().GetMetaData(AllocBeg(locked_version)));
   }
   void *AllocBeg(bool locked_version = false) {
struct QuarantineCallback {
 
   void Recycle(AsanChunk *m) {
     u8 old_chunk_state = CHUNK_QUARANTINE;
     if (!atomic_compare_exchange_strong(&m->chunk_state, &old_chunk_state,
                                         CHUNK_INVALID, memory_order_acquire)) {
       CHECK_EQ(old_chunk_state, CHUNK_QUARANTINE);
     }
 
-    CHECK_NE(m->alloc_tid, kInvalidTid);
-    CHECK_NE(m->free_tid, kInvalidTid);
     PoisonShadow(m->Beg(),
                  RoundUpTo(m->UsedSize(), SHADOW_GRANULARITY),
                  kAsanHeapLeftRedzoneMagic);
     void *p = reinterpret_cast<void *>(m->AllocBeg());
     if (p != m) {
       uptr *alloc_magic = reinterpret_cast<uptr *>(p);
       CHECK_EQ(alloc_magic[0], kAllocBegMagic);
       // Clear the magic value, as allocator internals may overwrite the
struct Allocator {
   }
 
   bool UpdateAllocationStack(uptr addr, BufferedStackTrace *stack) {
     AsanChunk *m = GetAsanChunkByAddr(addr);
     if (!m) return false;
     if (atomic_load(&m->chunk_state, memory_order_acquire) != CHUNK_ALLOCATED)
       return false;
     if (m->Beg() != addr) return false;
-    atomic_store(&m->alloc_context_id, StackDepotPut(*stack),
-                 memory_order_relaxed);
+    AsanThread *t = GetCurrentThread();
+    m->SetAllocContext(t ? t->tid() : 0, StackDepotPut(*stack));
     return true;
   }
 
   // -------------------- Allocation/Deallocation routines ---------------
   void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
                  AllocType alloc_type, bool can_fill) {
     if (UNLIKELY(!asan_inited))
       AsanInitFromRtl();
void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
     if (!IsAligned(user_beg, alignment))
       user_beg = RoundUpTo(user_beg, alignment);
     uptr user_end = user_beg + size;
     CHECK_LE(user_end, alloc_end);
     uptr chunk_beg = user_beg - kChunkHeaderSize;
     AsanChunk *m = reinterpret_cast<AsanChunk *>(chunk_beg);
     m->alloc_type = alloc_type;
     m->rz_log = rz_log;
-    u32 alloc_tid = t ? t->tid() : 0;
-    m->alloc_tid = alloc_tid;
-    CHECK_EQ(alloc_tid, m->alloc_tid);  // Does alloc_tid fit into the bitfield?
     m->from_memalign = user_beg != beg_plus_redzone;
     if (alloc_beg != chunk_beg) {
       CHECK_LE(alloc_beg + 2 * sizeof(uptr), chunk_beg);
       reinterpret_cast<uptr *>(alloc_beg)[0] = kAllocBegMagic;
       reinterpret_cast<uptr *>(alloc_beg)[1] = chunk_beg;
     }
     if (using_primary_allocator) {
       CHECK(size);
       m->user_requested_size = size;
       CHECK(allocator.FromPrimary(allocated));
     } else {
       CHECK(!allocator.FromPrimary(allocated));
       m->user_requested_size = SizeClassMap::kMaxSize;
       uptr *meta = reinterpret_cast<uptr *>(allocator.GetMetaData(allocated));
       meta[0] = size;
       meta[1] = chunk_beg;
     }
     m->user_requested_alignment_log = user_requested_alignment_log;
 
-    atomic_store(&m->alloc_context_id, StackDepotPut(*stack),
-                 memory_order_relaxed);
+    m->SetAllocContext(t ? t->tid() : 0, StackDepotPut(*stack));
 
     uptr size_rounded_down_to_granularity =
         RoundDownTo(size, SHADOW_GRANULARITY);
     // Unpoison the bulk of the memory region.
     if (size_rounded_down_to_granularity)
       PoisonShadow(user_beg, size_rounded_down_to_granularity, 0);
     // Deal with the end of the region if size is not aligned to granularity.
     if (size != size_rounded_down_to_granularity && CanPoisonMemory()) {
if (!atomic_compare_exchange_strong(&m->chunk_state, &old_chunk_state,
                                         CHUNK_QUARANTINE,
                                         memory_order_acquire)) {
       ReportInvalidFree(ptr, old_chunk_state, stack);
       // It's not safe to push a chunk in quarantine on invalid free.
       return false;
     }
     CHECK_EQ(CHUNK_ALLOCATED, old_chunk_state);
     // It was a user data.
-    m->free_tid = kInvalidTid;
-    m->free_context_id = 0;
+    m->SetFreeContext(kInvalidTid, 0);
     return true;
   }
 
   // Expects the chunk to already be marked as quarantined by using
   // AtomicallySetQuarantineFlagIfAllocated.
   void QuarantineChunk(AsanChunk *m, void *ptr, BufferedStackTrace *stack) {
     CHECK_EQ(atomic_load(&m->chunk_state, memory_order_relaxed),
              CHUNK_QUARANTINE);
     AsanThread *t = GetCurrentThread();
-    m->free_tid = t ? t->tid() : 0;
-    m->free_context_id = StackDepotPut(*stack);
+    m->SetFreeContext(t ? t->tid() : 0, StackDepotPut(*stack));
 
     Flags &fl = *flags();
     if (fl.max_free_fill_size > 0) {
       // We have to skip the chunk header, it contains free_context_id.
       uptr scribble_start = (uptr)m + kChunkHeaderSize + kChunkHeader2Size;
       if (m->UsedSize() >= kChunkHeader2Size) {  // Skip Header2 in user area.
         uptr size_to_fill = m->UsedSize() - kChunkHeader2Size;
         size_to_fill = Min(size_to_fill, (uptr)fl.max_free_fill_size);
return chunk_ && atomic_load(&chunk_->chunk_state, memory_order_relaxed) ==
                        CHUNK_QUARANTINE;
 }
 uptr AsanChunkView::Beg() const { return chunk_->Beg(); }
 uptr AsanChunkView::End() const { return Beg() + UsedSize(); }
 uptr AsanChunkView::UsedSize() const { return chunk_->UsedSize(); }
 u32 AsanChunkView::UserRequestedAlignment() const {
   return Allocator::ComputeUserAlignment(chunk_->user_requested_alignment_log);
 }
-uptr AsanChunkView::AllocTid() const { return chunk_->alloc_tid; }
+
+uptr AsanChunkView::AllocTid() const {
+  u32 tid = 0;
+  u32 stack = 0;
+  chunk_->GetAllocContext(tid, stack);
+  return tid;
+}
+
 uptr AsanChunkView::FreeTid() const {
-  return IsQuarantined() ? chunk_->free_tid : kInvalidTid;
+  if (!IsQuarantined())
+    return kInvalidTid;
+  u32 tid = 0;
+  u32 stack = 0;
+  chunk_->GetFreeContext(tid, stack);
+  return tid;
 }
+
 AllocType AsanChunkView::GetAllocType() const {
   return (AllocType)chunk_->alloc_type;
 }
 
 static StackTrace GetStackTraceFromId(u32 id) {
   CHECK(id);
   StackTrace res = StackDepotGet(id);
   CHECK(res.trace);
   return res;
 }
 
 u32 AsanChunkView::GetAllocStackId() const {
-  return atomic_load(&chunk_->alloc_context_id, memory_order_relaxed);
+  u32 tid = 0;
+  u32 stack = 0;
+  chunk_->GetAllocContext(tid, stack);
+  return stack;
 }
+
 u32 AsanChunkView::GetFreeStackId() const {
-  return IsQuarantined() ? chunk_->free_context_id : 0;
+  if (!IsQuarantined())
+    return 0;
+  u32 tid = 0;
+  u32 stack = 0;
+  chunk_->GetFreeContext(tid, stack);
+  return stack;
 }
 
 StackTrace AsanChunkView::GetAllocStack() const {
   return GetStackTraceFromId(GetAllocStackId());
 }
 
 StackTrace AsanChunkView::GetFreeStack() const {
   return GetStackTraceFromId(GetFreeStackId());

 
 uptr LsanMetadata::requested_size() const {
   __asan::AsanChunk *m = reinterpret_cast<__asan::AsanChunk *>(metadata_);
   return m->UsedSize(/*locked_version=*/true);
 }
 
 u32 LsanMetadata::stack_trace_id() const {
   __asan::AsanChunk *m = reinterpret_cast<__asan::AsanChunk *>(metadata_);
-  return atomic_load(&m->alloc_context_id, memory_order_relaxed);
+  u32 tid = 0;
+  u32 stack = 0;
+  m->GetAllocContext(tid, stack);
+  return stack;
 }
 
 void ForEachChunk(ForEachChunkCallback callback, void *arg) {
   __asan::get_allocator().ForEachChunk(callback, arg);
 }
 
 IgnoreObjectResult IgnoreObjectLocked(const void *p) {
   uptr addr = reinterpret_cast<uptr>(p);

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
+// RUN: %clangxx_asan -O0 %s -o %t
+
+// RUN: not %run %t 10 0 2>&1 | FileCheck %s --check-prefix=CHECK-%os --check-prefixes=CHECK,T0
+// RUN: not %run %t 10000000 0 2>&1 | FileCheck %s --check-prefixes=CHECK,T0
+
+// RUN: not %run %t 10 1 2>&1 | FileCheck %s --check-prefix=CHECK-%os --check-prefixes=CHECK,T1
+// RUN: not %run %t 10000000 1 2>&1 | FileCheck %s --check-prefixes=CHECK,T1
+
 // REQUIRES: stable-runtime
 
-#include <stdlib.h>
 #include <sanitizer/asan_interface.h>
+#include <stdlib.h>
+#include <thread>
 
 void UPDATE(void *p) {
   __asan_update_allocation_context(p);
 }
 
-int main() {
-  char *x = (char*)malloc(SIZE * sizeof(char));
+int main(int argc, char *argv[]) {
+  char *x = (char *)malloc(atoi(argv[1]) * sizeof(char));
+  if (atoi(argv[2]))
+    std::thread([&]() { UPDATE(x); }).join();
+  else
     UPDATE(x);
   free(x);
   return x[5];
   // CHECK: {{.*ERROR: AddressSanitizer: heap-use-after-free on address}}
+  // CHECK: READ of size 1 at {{.*}} thread T0
+  // T0: allocated by thread T0 here
+  // T1: allocated by thread T1 here
   // CHECK: UPDATE
 }