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

Author: bernhardu

compiler-rt/lib/asan/asan_allocator.cpp

@@ -525,6 +525,17 @@ struct Allocator {
     return true;
   }
 
+  bool UpdateDeallocationStack(uptr addr, BufferedStackTrace *stack) {
+    AsanChunk *m = GetAsanChunkByAddr(addr);
+    if (!m) return false;
+    if (atomic_load(&m->chunk_state, memory_order_acquire) != CHUNK_ALLOCATED)
+      return false;
+    if (m->Beg() != addr) return false;
+    AsanThread *t = GetCurrentThread();
+    m->SetFreeContext(t ? t->tid() : kMainTid, StackDepotPut(*stack));
+    return true;
+  }
+
   // -------------------- Allocation/Deallocation routines ---------------
   void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
                  AllocType alloc_type, bool can_fill) {
@@ -941,8 +952,8 @@ uptr AsanChunkView::AllocTid() const {
 }
 
 uptr AsanChunkView::FreeTid() const {
-  if (!IsQuarantined())
-    return kInvalidTid;
+  //if (!IsQuarantined())
+  //  return kInvalidTid;
   u32 tid = 0;
   u32 stack = 0;
   chunk_->GetFreeContext(tid, stack);
@@ -961,8 +972,8 @@ u32 AsanChunkView::GetAllocStackId() const {
 }
 
 u32 AsanChunkView::GetFreeStackId() const {
-  if (!IsQuarantined())
-    return 0;
+  //if (!IsQuarantined())
+  //  return 0;
   u32 tid = 0;
   u32 stack = 0;
   chunk_->GetFreeContext(tid, stack);
@@ -1119,6 +1130,10 @@ void asan_mz_force_unlock() SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
   instance.ForceUnlock();
 }
 
+int asan_update_deallocation_context(void* addr, BufferedStackTrace *stack) {
+  return instance.UpdateDeallocationStack((uptr)addr, stack);
+}
+
 }  // namespace __asan
 
 // --- Implementation of LSan-specific functions --- {{{1

compiler-rt/lib/asan/asan_allocator.h

@@ -292,6 +292,8 @@ uptr asan_mz_size(const void *ptr);
 void asan_mz_force_lock();
 void asan_mz_force_unlock();
 
+int asan_update_deallocation_context(void* addr, BufferedStackTrace *stack);
+
 void PrintInternalAllocatorStats();
 void AsanSoftRssLimitExceededCallback(bool exceeded);
 

compiler-rt/lib/asan/asan_malloc_win.cpp

@@ -17,6 +17,7 @@
 #include "asan_allocator.h"
 #include "asan_interceptors.h"
 #include "asan_internal.h"
+#include "asan_poisoning.h"
 #include "asan_stack.h"
 #include "interception/interception.h"
 #include <stddef.h>
@@ -296,7 +297,7 @@ void *SharedReAlloc(ReAllocFunction reallocFunc, SizeFunction heapSizeFunc,
     if (ownershipState == RTL ||
         (ownershipState == NEITHER && !only_asan_supported_flags)) {
       if (only_asan_supported_flags) {
-        // if this is a conversion to ASAN upported flags, transfer this
+        // if this is a conversion to ASAN supported flags, transfer this
         // allocation to the ASAN allocator
         void *replacement_alloc;
         if (dwFlags & HEAP_ZERO_MEMORY)
@@ -323,6 +324,27 @@ void *SharedReAlloc(ReAllocFunction reallocFunc, SizeFunction heapSizeFunc,
     }
 
     if (ownershipState == ASAN && !only_asan_supported_flags) {
+
+      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 lpMem;
+        } else if (dwBytes < 1 || dwBytes >= old_usable_size) {
+          // growing with HEAP_REALLOC_IN_PLACE_ONLY is not supported.
+          return nullptr;
+        } else {
+          // poison just the "freed" part
+          uptr aligned_ptr = RoundUpTo((uptr)((char *)lpMem + dwBytes), ASAN_SHADOW_GRANULARITY);
+          uptr aligned_size = RoundUpTo(dwBytes, ASAN_SHADOW_GRANULARITY);
+          PoisonShadow(aligned_ptr, aligned_size, kAsanHeapFreeMagic);
+          // unpoison the lower part, it could be poisoned by a previous realloc in place
+          __asan_unpoison_memory_region((char *)lpMem, dwBytes);
+          __asan::asan_update_deallocation_context(lpMem, &stack);
+          return lpMem;
+        }
+      }
+
       // Conversion to unsupported flags allocation,
       // transfer this allocation back to the original allocator.
       void *replacement_alloc = allocFunc(hHeap, dwFlags, dwBytes);
@@ -348,6 +370,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,68 @@
+// 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;
+  void *ret;
+  buffer = (char *)RtlAllocateHeap_ptr(GetProcessHeap(), 0, 23),
+
+  ret = RtlReAllocateHeap_ptr(GetProcessHeap(), HEAP_REALLOC_IN_PLACE_ONLY,
+                        buffer, 7);
+  if (!ret) {
+    puts("returned nullptr");
+  }
+  buffer[6] = 'a';
+  puts("Okay 6");
+  fflush(stdout);
+  // CHECK: Okay 6
+
+  ret = RtlReAllocateHeap_ptr(GetProcessHeap(), HEAP_REALLOC_IN_PLACE_ONLY,
+                        buffer, 15);
+  if (!ret) {
+    puts("returned nullptr");
+  }
+  buffer[14] = 'a';
+  puts("Okay 14");
+  fflush(stdout);
+  // CHECK: Okay 14
+
+  buffer[15] = 'a';
+  // CHECK: AddressSanitizer: heap-use-after-free on address [[ADDR:0x[0-9a-f]+]]
+  // CHECK: WRITE of size 1 at [[ADDR]] thread T0
+
+  RtlFreeHeap_ptr(GetProcessHeap(), 0, buffer);
+}