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
main()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