[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 882cf0418f0ae1c9039e0e939d4ff7a1eb39fe9c 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.

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      | 74 +++++++++++++++++++
 2 files changed, 90 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 3278f07219876..31b944ff7d5cb 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 0000000000000..2248399d6ff18
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/Windows/rtlallocateheap_realloc_in_place.cpp
@@ -0,0 +1,74 @@
+// 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);
+}