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

[via llvm-commits]

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

>From bfcb81b9607eba0ef9971be1e4a0665d5b6936c4 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.

Currently it prints only a "use-after-poison" message,
not a proper "heap-buffer-overflow" with a hint to a reallocation.
---
 compiler-rt/lib/asan/asan_malloc_win.cpp      | 17 +++++-
 .../rtlallocateheap_realloc_in_place.cpp      | 61 +++++++++++++++++++
 2 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 compiler-rt/test/asan/TestCases/Windows/rtlallocateheap_realloc_in_place.cpp

diff --git a/compiler-rt/lib/asan/asan_malloc_win.cpp b/compiler-rt/lib/asan/asan_malloc_win.cpp
index 3278f07219876..4a8eb0acb174e 100644
--- a/compiler-rt/lib/asan/asan_malloc_win.cpp
+++ b/compiler-rt/lib/asan/asan_malloc_win.cpp
@@ -323,12 +323,24 @@ void *SharedReAlloc(ReAllocFunction reallocFunc, SizeFunction heapSizeFunc,
     }
 
     if (ownershipState == ASAN && !only_asan_supported_flags) {
+      size_t old_usable_size = asan_malloc_usable_size(lpMem, pc, bp);
+      if (dwFlags & HEAP_REALLOC_IN_PLACE_ONLY) {
+        if (dwBytes >= 1 && dwBytes <= old_usable_size) {
+          if (dwBytes < old_usable_size) {
+            __asan_poison_memory_region((char *)lpMem + dwBytes,
+                                        old_usable_size - dwBytes);
+          }
+          __asan_unpoison_memory_region((char *)lpMem, dwBytes);
+          return lpMem;
+        } else {
+          return nullptr;
+        }
+      }
+
       // Conversion to unsupported flags allocation,
       // transfer this allocation back to the original allocator.
       void *replacement_alloc = allocFunc(hHeap, dwFlags, dwBytes);
-      size_t old_usable_size = 0;
       if (replacement_alloc) {
-        old_usable_size = asan_malloc_usable_size(lpMem, pc, bp);
         REAL(memcpy)(replacement_alloc, lpMem,
                      Min<size_t>(dwBytes, old_usable_size));
         asan_free(lpMem, &stack, FROM_MALLOC);
@@ -348,6 +360,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..ffe62b7c30723
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/Windows/rtlallocateheap_realloc_in_place.cpp
@@ -0,0 +1,61 @@
+// 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;
+  buffer = (char *)RtlAllocateHeap_ptr(GetProcessHeap(), 0, 48),
+
+  RtlReAllocateHeap_ptr(GetProcessHeap(), HEAP_REALLOC_IN_PLACE_ONLY, buffer,
+                        16);
+  buffer[15] = 'a';
+  puts("Okay 15");
+  fflush(stdout);
+  // CHECK: Okay 15
+
+  RtlReAllocateHeap_ptr(GetProcessHeap(), HEAP_REALLOC_IN_PLACE_ONLY, buffer,
+                        32);
+  buffer[31] = 'a';
+  puts("Okay 31");
+  fflush(stdout);
+  // CHECK: Okay 31
+
+  buffer[32] = 'a';
+  // CHECK: AddressSanitizer: use-after-poison on address [[ADDR:0x[0-9a-f]+]]
+  // CHECK: WRITE of size 1 at [[ADDR]] thread T0
+
+  RtlFreeHeap_ptr(GetProcessHeap(), 0, buffer);
+}