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

Author: bernhardu

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;
 

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);
+}