From d349daa135648e7e18da5be0c38ffbdbc7a208ca Mon Sep 17 00:00:00 2001
From: bernhardu <bernhardu@mailbox.org>
Date: Thu, 4 Sep 2025 01:34:55 +0200
Subject: [PATCH] [win/asan] Improve SharedReAlloc with
 HEAP_REALLOC_IN_PLACE_ONLY. (#132558)

Currently with HEAP_REALLOC_IN_PLACE_ONLY a new allocation
gets returned with the content copied from the original pointer,
which gets freed.

But applications may rely on HEAP_REALLOC_IN_PLACE_ONLY returning
the same pointer as they give as input to e.g. RtlReAllocateHeap.
If e.g. growing is not possible it fails without
modifying the input pointer.

Downside of this patch is, it won't detect accesses to the area
getting "free" by a shrinking reallocation.
---
 compiler-rt/lib/asan/asan_malloc_win.cpp      | 16 +++++
 .../rtlallocateheap_realloc_in_place.cpp      | 71 +++++++++++++++++++
 2 files changed, 87 insertions(+)
 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 8d98da940800..ea6f7dfaa08c 100644
--- a/compiler-rt/lib/asan/asan_malloc_win.cpp
+++ b/compiler-rt/lib/asan/asan_malloc_win.cpp
@@ -322,6 +322,22 @@ void *SharedReAlloc(ReAllocFunction reallocFunc, SizeFunction heapSizeFunc,
       }
     }
 
+    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 the current pointer.
+        return lpMem;
+      } else if (dwBytes >= old_usable_size) {
+        // Growing with HEAP_REALLOC_IN_PLACE_ONLY is not supported.
+        return nullptr;
+      } else {
+        // Shrinking with HEAP_REALLOC_IN_PLACE_ONLY is not yet supported.
+        // For now return the current pointer and
+        // leave the allocation size as it is.
+        return lpMem;
+      }
+    }
+
     if (ownershipState == ASAN && !only_asan_supported_flags) {
       // Conversion to unsupported flags allocation,
       // transfer this allocation back to the original allocator.
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 000000000000..35fa9ce57429
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/Windows/rtlallocateheap_realloc_in_place.cpp
@@ -0,0 +1,71 @@
+// RUN: %clang_cl_asan %Od %s %Fe%t %MD
+// RUN: %env_asan_opts=windows_hook_rtl_allocators=true 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) {
+    fputs("Couldn't load ntdll??\n", stderr);
+    return -1;
+  }
+
+  auto RtlAllocateHeap_ptr =
+      (AllocateFunctionPtr)GetProcAddress(NtDllHandle, "RtlAllocateHeap");
+  if (RtlAllocateHeap_ptr == 0) {
+    fputs("Couldn't RtlAllocateHeap\n", stderr);
+    return -1;
+  }
+
+  auto RtlReAllocateHeap_ptr =
+      (ReAllocateFunctionPtr)GetProcAddress(NtDllHandle, "RtlReAllocateHeap");
+  if (RtlReAllocateHeap_ptr == 0) {
+    fputs("Couldn't find RtlReAllocateHeap\n", stderr);
+    return -1;
+  }
+
+  auto RtlFreeHeap_ptr =
+      (FreeFunctionPtr)GetProcAddress(NtDllHandle, "RtlFreeHeap");
+  if (RtlFreeHeap_ptr == 0) {
+    fputs("Couldn't RtlFreeHeap\n", stderr);
+    return -1;
+  }
+
+  char *ptr1;
+  char *ptr2;
+  ptr2 = ptr1 = (char *)RtlAllocateHeap_ptr(GetProcessHeap(), 0, 15);
+  if (ptr1)
+    fputs("Okay alloc\n", stderr);
+  // CHECK: Okay alloc
+
+  // TODO: Growing is currently not supported
+  ptr2 = (char *)RtlReAllocateHeap_ptr(GetProcessHeap(),
+                                       HEAP_REALLOC_IN_PLACE_ONLY, ptr1, 23);
+  if (ptr2 == NULL)
+    fputs("Okay grow failed\n", stderr);
+  // CHECK: Okay grow failed
+
+  // TODO: Shrinking is currently not supported
+  ptr2 = (char *)RtlReAllocateHeap_ptr(GetProcessHeap(),
+                                       HEAP_REALLOC_IN_PLACE_ONLY, ptr1, 7);
+  if (ptr2 == ptr1)
+    fputs("Okay shrinking return the original pointer\n", stderr);
+  // CHECK: Okay shrinking return the original pointer
+
+  ptr1[7] = 'a';
+  fputs("Okay 7\n", stderr);
+  // CHECK: Okay 7
+
+  // TODO: Writing behind the shrinked part is currently not detected.
+  //       Therefore test writing behind the original allocation for now.
+  ptr1[16] = 'a';
+  // CHECK: AddressSanitizer: heap-buffer-overflow on address [[ADDR:0x[0-9a-f]+]]
+  // CHECK: WRITE of size 1 at [[ADDR]] thread T0
+
+  RtlFreeHeap_ptr(GetProcessHeap(), 0, ptr1);
+}