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

2 files changed, 87 insertions, 0 deletions

diff --git a/compiler-rt/lib/asan/asan_malloc_win.cpp b/compiler-rt/lib/asan/asan_malloc_win.cpp
index 8d98da9..ea6f7df 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 0000000..35fa9ce
--- /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);
+}