1 // RUN: %clang_cl_asan -Od %s -Fe%t /MD
2 // RUN: %env_asan_opts=windows_hook_rtl_allocators=true not %run %t 2>&1 | FileCheck %s
3 // UNSUPPORTED: asan-64-bits
4 // REQUIRES: asan-rtl-heap-interception
5 
6 #include <assert.h>
7 #include <stdio.h>
8 #include <windows.h>
9 
10 using AllocateFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, SIZE_T);
11 using ReAllocateFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, PVOID, SIZE_T);
12 using FreeFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, PVOID);
13 
14 int main() {
15   HMODULE NtDllHandle = GetModuleHandle("ntdll.dll");
16   if (!NtDllHandle) {
17     puts("Couldn't load ntdll??");
18     return -1;
19   }
20 
21   auto RtlAllocateHeap_ptr = (AllocateFunctionPtr)GetProcAddress(NtDllHandle, "RtlAllocateHeap");
22   if (RtlAllocateHeap_ptr == 0) {
23     puts("Couldn't find RtlAllocateHeap");
24     return -1;
25   }
26 
27   auto RtlReAllocateHeap_ptr = (ReAllocateFunctionPtr)GetProcAddress(NtDllHandle, "RtlReAllocateHeap");
28   if (RtlReAllocateHeap_ptr == 0) {
29     puts("Couldn't find RtlReAllocateHeap");
30     return -1;
31   }
32 
33   char *buffer;
34   SIZE_T buffer_size = 32;
35   SIZE_T new_buffer_size = buffer_size * 2;
36 
37   buffer = (char *)RtlAllocateHeap_ptr(GetProcessHeap(), HEAP_ZERO_MEMORY, buffer_size);
38   assert(buffer != nullptr);
39   // Check that the buffer is zeroed.
40   for (SIZE_T i = 0; i < buffer_size; ++i) {
41     assert(buffer[i] == 0);
42   }
43   memset(buffer, 0xcc, buffer_size);
44 
45   // Zero the newly allocated memory.
46   buffer = (char *)RtlReAllocateHeap_ptr(GetProcessHeap(), HEAP_ZERO_MEMORY, buffer, new_buffer_size);
47   assert(buffer != nullptr);
48   // Check that the first part of the buffer still has the old contents.
49   for (SIZE_T i = 0; i < buffer_size; ++i) {
50     assert(buffer[i] == (char)0xcc);
51   }
52   // Check that the new part of the buffer is zeroed.
53   for (SIZE_T i = buffer_size; i < new_buffer_size; ++i) {
54     assert(buffer[i] == 0x0);
55   }
56 
57   // Shrink the buffer back down.
58   buffer = (char *)RtlReAllocateHeap_ptr(GetProcessHeap(), HEAP_ZERO_MEMORY, buffer, buffer_size);
59   assert(buffer != nullptr);
60   // Check that the first part of the buffer still has the old contents.
61   for (SIZE_T i = 0; i < buffer_size; ++i) {
62     assert(buffer[i] == (char)0xcc);
63   }
64 
65   buffer[buffer_size + 1] = 'a';
66   // CHECK: AddressSanitizer: heap-buffer-overflow on address [[ADDR:0x[0-9a-f]+]]
67   // CHECK: WRITE of size 1 at [[ADDR]] thread T0
68 }
69