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