compiler-rt/test/asan/TestCases/Windows/heapalloc_rtl_transfer.cpp

   1 #include "sanitizer\allocator_interface.h"
   2 #include <cassert>
   3 #include <stdio.h>
   4 #include <windows.h>
   5
   6 // RUN: %clang_cl_asan %s -o%t
   7 // RUN: %env_asan_opts=windows_hook_rtl_allocators=true %run %t 2>&1 | FileCheck %s
   8 // UNSUPPORTED: asan-64-bits
   9
  10 using AllocateFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, SIZE_T);
  11 using ReAllocateFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, PVOID, SIZE_T);
  12
  13 using FreeFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, PVOID);
  14
  15 int main() {
  16   HMODULE NtDllHandle = GetModuleHandle("ntdll.dll");
  17   if (!NtDllHandle) {
  18     puts("Couldn't load ntdll??");
  19     return -1;
  20   }
  21
  22   auto RtlAllocateHeap_ptr = (AllocateFunctionPtr)GetProcAddress(NtDllHandle, "RtlAllocateHeap");
  23   if (RtlAllocateHeap_ptr == 0) {
  24     puts("Couldn't find RtlAllocateHeap");
  25     return -1;
  26   }
  27
  28   auto RtlReAllocateHeap_ptr = (ReAllocateFunctionPtr)GetProcAddress(NtDllHandle, "RtlReAllocateHeap");
  29   if (RtlReAllocateHeap_ptr == 0) {
  30     puts("Couldn't find RtlReAllocateHeap");
  31     return -1;
  32   }
  33
  34   //owned by rtl
  35   void *alloc = RtlAllocateHeap_ptr(GetProcessHeap(),
  36                                     HEAP_GENERATE_EXCEPTIONS | HEAP_ZERO_MEMORY, 100);
  37   assert(alloc);
  38   for (int i = 0; i < 100; i++) {
  39     assert(((char *)alloc)[i] == 0);
  40     ((char *)alloc)[i] = '\xcc';
  41   }
  42
  43   // still owned by rtl
  44   alloc = RtlReAllocateHeap_ptr(GetProcessHeap(),
  45                                 HEAP_GENERATE_EXCEPTIONS | HEAP_ZERO_MEMORY, alloc, 500);
  46   assert(alloc && !__sanitizer_get_ownership(alloc) && HeapValidate(GetProcessHeap(), 0, alloc));
  47   for (int i = 0; i < 100; i++) {
  48     assert(((char *)alloc)[i] == '\xcc');
  49   }
  50   for (int i = 100; i < 500; i++) {
  51     assert(((char *)alloc)[i] == 0);
  52     ((char *)alloc)[i] = '\xcc';
  53   }
  54
  55   //convert to asan owned
  56   void *realloc = RtlReAllocateHeap_ptr(GetProcessHeap(),
  57                                         HEAP_ZERO_MEMORY, alloc, 600);
  58   alloc = nullptr;
  59   assert(realloc && __sanitizer_get_ownership(realloc));
  60
  61   for (int i = 0; i < 500; i++) {
  62     assert(((char *)realloc)[i] == '\xcc');
  63   }
  64   for (int i = 500; i < 600; i++) {
  65     assert(((char *)realloc)[i] == 0);
  66     ((char *)realloc)[i] = '\xcc';
  67   }
  68   realloc = RtlReAllocateHeap_ptr(GetProcessHeap(),
  69                                   HEAP_ZERO_MEMORY, realloc, 2048);
  70   assert(realloc && __sanitizer_get_ownership(realloc));
  71
  72   for (int i = 0; i < 600; i++) {
  73     assert(((char *)realloc)[i] == '\xcc');
  74   }
  75   for (int i = 600; i < 2048; i++) {
  76     assert(((char *)realloc)[i] == 0);
  77     ((char *)realloc)[i] = '\xcc';
  78   }
  79   //convert back to rtl owned;
  80   alloc = RtlReAllocateHeap_ptr(GetProcessHeap(),
  81                                 HEAP_ZERO_MEMORY | HEAP_GENERATE_EXCEPTIONS, realloc, 100);
  82   assert(alloc && !__sanitizer_get_ownership(alloc) && HeapValidate(GetProcessHeap(), 0, alloc));
  83   for (int i = 0; i < 100; i++) {
  84     assert(((char *)alloc)[i] == '\xcc');
  85     ((char *)alloc)[i] = 0;
  86   }
  87
  88   auto usable_size = HeapSize(GetProcessHeap(), 0, alloc);
  89   for (int i = 100; i < usable_size; i++) {
  90     assert(((char *)alloc)[i] == 0);
  91   }
  92
  93   printf("Success\n");
  94 }
  95
  96 // CHECK-NOT: Assertion failed:
  97 // CHECK-NOT: AddressSanitizer
  98 // CHECK: Success