99 lines · 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%t7// RUN: %env_asan_opts=windows_hook_rtl_allocators=true %run %t 2>&1 | FileCheck %s8// UNSUPPORTED: asan-64-bits9 10using AllocateFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, SIZE_T);11using ReAllocateFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, PVOID, SIZE_T);12 13using FreeFunctionPtr = PVOID(__stdcall *)(PVOID, ULONG, PVOID);14 15int 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 rtl35 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 rtl44 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 owned56 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: AddressSanitizer98// CHECK: Success99