[compiler-rt] [win/asan] Improve SharedReAlloc with HEAP_REALLOC_IN_PLACE_ONLY. (PR #132558)

via llvm-commits llvm-commits at lists.llvm.org
Wed May 28 11:18:27 PDT 2025


https://github.com/bernhardu updated https://github.com/llvm/llvm-project/pull/132558

>From bf73a58fd3c438106ad9dde227c82640068be0af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bernhard=20=C3=9Cbelacker?= <bernhardu at mailbox.org>
Date: Tue, 18 Mar 2025 18:44:47 +0100
Subject: [PATCH] [win/asan] Improve SharedReAlloc with
 HEAP_REALLOC_IN_PLACE_ONLY.

This patch allows reallocations in place if the size is
below or equal to the initial allocated size.
---
 compiler-rt/lib/asan/asan_allocator.cpp       | 23 +++++--
 compiler-rt/lib/asan/asan_allocator.h         |  2 +
 compiler-rt/lib/asan/asan_descriptions.cpp    |  4 +-
 compiler-rt/lib/asan/asan_descriptions.h      |  1 +
 compiler-rt/lib/asan/asan_malloc_win.cpp      | 25 ++++++-
 .../rtlallocateheap_realloc_in_place.cpp      | 68 +++++++++++++++++++
 6 files changed, 117 insertions(+), 6 deletions(-)
 create mode 100644 compiler-rt/test/asan/TestCases/Windows/rtlallocateheap_realloc_in_place.cpp

diff --git a/compiler-rt/lib/asan/asan_allocator.cpp b/compiler-rt/lib/asan/asan_allocator.cpp
index 3a55c2af65653..4cb1767291996 100644
--- a/compiler-rt/lib/asan/asan_allocator.cpp
+++ b/compiler-rt/lib/asan/asan_allocator.cpp
@@ -525,6 +525,17 @@ struct Allocator {
     return true;
   }
 
+  bool UpdateDeallocationStack(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;
+    AsanThread *t = GetCurrentThread();
+    m->SetFreeContext(t ? t->tid() : kMainTid, StackDepotPut(*stack));
+    return true;
+  }
+
   // -------------------- Allocation/Deallocation routines ---------------
   void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
                  AllocType alloc_type, bool can_fill) {
@@ -941,8 +952,8 @@ uptr AsanChunkView::AllocTid() const {
 }
 
 uptr AsanChunkView::FreeTid() const {
-  if (!IsQuarantined())
-    return kInvalidTid;
+  //if (!IsQuarantined())
+  //  return kInvalidTid;
   u32 tid = 0;
   u32 stack = 0;
   chunk_->GetFreeContext(tid, stack);
@@ -961,8 +972,8 @@ u32 AsanChunkView::GetAllocStackId() const {
 }
 
 u32 AsanChunkView::GetFreeStackId() const {
-  if (!IsQuarantined())
-    return 0;
+  //if (!IsQuarantined())
+  //  return 0;
   u32 tid = 0;
   u32 stack = 0;
   chunk_->GetFreeContext(tid, stack);
@@ -1119,6 +1130,10 @@ void asan_mz_force_unlock() SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
   instance.ForceUnlock();
 }
 
+int asan_update_deallocation_context(void* addr, BufferedStackTrace *stack) {
+  return instance.UpdateDeallocationStack((uptr)addr, stack);
+}
+
 }  // namespace __asan
 
 // --- Implementation of LSan-specific functions --- {{{1
diff --git a/compiler-rt/lib/asan/asan_allocator.h b/compiler-rt/lib/asan/asan_allocator.h
index db8dc3bebfc62..7236543e3f39e 100644
--- a/compiler-rt/lib/asan/asan_allocator.h
+++ b/compiler-rt/lib/asan/asan_allocator.h
@@ -292,6 +292,8 @@ uptr asan_mz_size(const void *ptr);
 void asan_mz_force_lock();
 void asan_mz_force_unlock();
 
+int asan_update_deallocation_context(void* addr, BufferedStackTrace *stack);
+
 void PrintInternalAllocatorStats();
 void AsanSoftRssLimitExceededCallback(bool exceeded);
 
diff --git a/compiler-rt/lib/asan/asan_descriptions.cpp b/compiler-rt/lib/asan/asan_descriptions.cpp
index c9f3e4d682d95..7d24b9d6c27e9 100644
--- a/compiler-rt/lib/asan/asan_descriptions.cpp
+++ b/compiler-rt/lib/asan/asan_descriptions.cpp
@@ -123,6 +123,7 @@ static void GetAccessToHeapChunkInformation(ChunkAccess *descr,
   descr->chunk_size = chunk.UsedSize();
   descr->user_requested_alignment = chunk.UserRequestedAlignment();
   descr->alloc_type = chunk.GetAllocType();
+  descr->is_carantined = chunk.IsQuarantined();
 }
 
 static void PrintHeapChunkAccess(uptr addr, const ChunkAccess &descr) {
@@ -422,7 +423,8 @@ void HeapAddressDescription::Print() const {
   AsanThreadContext *free_thread = nullptr;
   if (free_tid != kInvalidTid) {
     free_thread = GetThreadContextByTidLocked(free_tid);
-    Printf("%sfreed by thread %s here:%s\n", d.Allocation(),
+    Printf("%s%sfreed by thread %s here:%s\n", d.Allocation(),
+           chunk_access.is_carantined ? "": "partially ",
            AsanThreadIdAndName(free_thread).c_str(), d.Default());
     StackTrace free_stack = GetStackTraceFromId(free_stack_id);
     free_stack.Print();
diff --git a/compiler-rt/lib/asan/asan_descriptions.h b/compiler-rt/lib/asan/asan_descriptions.h
index a614f47d461bb..54e110c537b1b 100644
--- a/compiler-rt/lib/asan/asan_descriptions.h
+++ b/compiler-rt/lib/asan/asan_descriptions.h
@@ -113,6 +113,7 @@ struct ChunkAccess {
   u32 user_requested_alignment : 12;
   u32 access_type : 2;
   u32 alloc_type : 2;
+  u32 is_carantined : 1;
 };
 
 struct HeapAddressDescription {
diff --git a/compiler-rt/lib/asan/asan_malloc_win.cpp b/compiler-rt/lib/asan/asan_malloc_win.cpp
index 3278f07219876..dbef1943a09c2 100644
--- a/compiler-rt/lib/asan/asan_malloc_win.cpp
+++ b/compiler-rt/lib/asan/asan_malloc_win.cpp
@@ -17,6 +17,7 @@
 #include "asan_allocator.h"
 #include "asan_interceptors.h"
 #include "asan_internal.h"
+#include "asan_poisoning.h"
 #include "asan_stack.h"
 #include "interception/interception.h"
 #include <stddef.h>
@@ -296,7 +297,7 @@ void *SharedReAlloc(ReAllocFunction reallocFunc, SizeFunction heapSizeFunc,
     if (ownershipState == RTL ||
         (ownershipState == NEITHER && !only_asan_supported_flags)) {
       if (only_asan_supported_flags) {
-        // if this is a conversion to ASAN upported flags, transfer this
+        // if this is a conversion to ASAN supported flags, transfer this
         // allocation to the ASAN allocator
         void *replacement_alloc;
         if (dwFlags & HEAP_ZERO_MEMORY)
@@ -323,6 +324,27 @@ void *SharedReAlloc(ReAllocFunction reallocFunc, SizeFunction heapSizeFunc,
     }
 
     if (ownershipState == ASAN && !only_asan_supported_flags) {
+
+      if (dwFlags & HEAP_REALLOC_IN_PLACE_ONLY) {
+        size_t old_usable_size = asan_malloc_usable_size(lpMem, pc, bp);
+        if (dwBytes == old_usable_size) {
+          // nothing to change
+          return lpMem;
+        } else if (dwBytes < 1 || dwBytes >= old_usable_size) {
+          // growing with HEAP_REALLOC_IN_PLACE_ONLY is not supported.
+          return nullptr;
+        } else {
+          // poison just the "freed" part
+          uptr aligned_ptr = RoundUpTo((uptr)((char *)lpMem + dwBytes), ASAN_SHADOW_GRANULARITY);
+          uptr aligned_size = RoundUpTo(dwBytes, ASAN_SHADOW_GRANULARITY);
+          PoisonShadow(aligned_ptr, aligned_size, kAsanHeapFreeMagic);
+          // unpoison the lower part, it could be poisoned by a previous realloc in place
+          __asan_unpoison_memory_region((char *)lpMem, dwBytes);
+          __asan::asan_update_deallocation_context(lpMem, &stack);
+          return lpMem;
+        }
+      }
+
       // Conversion to unsupported flags allocation,
       // transfer this allocation back to the original allocator.
       void *replacement_alloc = allocFunc(hHeap, dwFlags, dwBytes);
@@ -348,6 +370,7 @@ void *SharedReAlloc(ReAllocFunction reallocFunc, SizeFunction heapSizeFunc,
   }
   // asan_realloc will never reallocate in place, so for now this flag is
   // unsupported until we figure out a way to fake this.
+  // Small exception when shrinking or staying below the inital size, see above.
   if (dwFlags & HEAP_REALLOC_IN_PLACE_ONLY)
     return nullptr;
 
diff --git a/compiler-rt/test/asan/TestCases/Windows/rtlallocateheap_realloc_in_place.cpp b/compiler-rt/test/asan/TestCases/Windows/rtlallocateheap_realloc_in_place.cpp
new file mode 100644
index 0000000000000..5773bca4bd3e5
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/Windows/rtlallocateheap_realloc_in_place.cpp
@@ -0,0 +1,68 @@
+// RUN: %clang_cl_asan %Od %s %Fe%t %MD
+// RUN: %env_asan_opts=windows_hook_rtl_allocators=true:halt_on_error=false not %run %t 2>&1 | FileCheck %s
+
+#include <stdio.h>
+#include <windows.h>
+
+using AllocateFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, SIZE_T);
+using ReAllocateFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, PVOID, SIZE_T);
+using FreeFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, PVOID);
+
+int main() {
+  HMODULE NtDllHandle = GetModuleHandle("ntdll.dll");
+  if (!NtDllHandle) {
+    puts("Couldn't load ntdll??");
+    return -1;
+  }
+
+  auto RtlAllocateHeap_ptr =
+      (AllocateFunctionPtr)GetProcAddress(NtDllHandle, "RtlAllocateHeap");
+  if (RtlAllocateHeap_ptr == 0) {
+    puts("Couldn't RtlAllocateHeap");
+    return -1;
+  }
+
+  auto RtlReAllocateHeap_ptr =
+      (ReAllocateFunctionPtr)GetProcAddress(NtDllHandle, "RtlReAllocateHeap");
+  if (RtlReAllocateHeap_ptr == 0) {
+    puts("Couldn't find RtlReAllocateHeap");
+    return -1;
+  }
+
+  auto RtlFreeHeap_ptr =
+      (FreeFunctionPtr)GetProcAddress(NtDllHandle, "RtlFreeHeap");
+  if (RtlFreeHeap_ptr == 0) {
+    puts("Couldn't RtlFreeHeap");
+    return -1;
+  }
+
+  char *buffer;
+  void *ret;
+  buffer = (char *)RtlAllocateHeap_ptr(GetProcessHeap(), 0, 23),
+
+  ret = RtlReAllocateHeap_ptr(GetProcessHeap(), HEAP_REALLOC_IN_PLACE_ONLY,
+                        buffer, 7);
+  if (!ret) {
+    puts("returned nullptr");
+  }
+  buffer[6] = 'a';
+  puts("Okay 6");
+  fflush(stdout);
+  // CHECK: Okay 6
+
+  ret = RtlReAllocateHeap_ptr(GetProcessHeap(), HEAP_REALLOC_IN_PLACE_ONLY,
+                        buffer, 15);
+  if (!ret) {
+    puts("returned nullptr");
+  }
+  buffer[14] = 'a';
+  puts("Okay 14");
+  fflush(stdout);
+  // CHECK: Okay 14
+
+  buffer[15] = 'a';
+  // CHECK: AddressSanitizer: heap-use-after-free on address [[ADDR:0x[0-9a-f]+]]
+  // CHECK: WRITE of size 1 at [[ADDR]] thread T0
+
+  RtlFreeHeap_ptr(GetProcessHeap(), 0, buffer);
+}



More information about the llvm-commits mailing list