// User-defined global heap prototypes extern void mem_global_init(); extern void mem_global_term(); extern void* mem_global_allocate(size_t size, size_t align); extern void mem_global_deallocate(void* ptr); extern size_t mem_global_get_size(void* ptr); // Actual code #define BREAK() __debugbreak() #include <stdio.h> #include <stdlib.h> #define WIN32_LEAN_AND_MEAN #include <windows.h> static inline void patch_with_jump(void* dest, void* address) { // get offset for relative jmp unsigned int offset = (unsigned int)((char*)address - (char*)dest - 5); // unprotect memory unsigned long old_protect; VirtualProtect(dest, 5, PAGE_READWRITE, &old_protect); // write jmp *(unsigned char*)dest = 0xe9; *(unsigned int*)((char*)dest + 1) = offset; // protect memory VirtualProtect(dest, 5, old_protect, &old_protect); } static void break_stub() { BREAK(); } #define PATCH(target, stub) void target(); patch_with_jump((void*)(target), (void*)(stub)) #define PATCH_STUB(target) PATCH(target, target ## _stub) #define PATCH_BREAK(target) PATCH(target, break_stub) #ifdef _DEBUG #define PATCH_DEBUG(target, stub) PATCH(target, stub) #define PATCH_BREAK_DEBUG(target) PATCH_BREAK(target) #else #define PATCH_DEBUG(target, stub) (void)0 #define PATCH_BREAK_DEBUG(target) (void)0 #endif #define PATCH_BREAK_ALL(target) PATCH_BREAK(target); PATCH_BREAK_DEBUG(target ## _base); PATCH_BREAK_DEBUG(target ## _dbg) extern "C" { const size_t malloc_alignment = 8; // This is a copy of __freeCrtMemory (we can't use it directly since it's not in CRT in release) static void freeCrtMemory() { // code removed to avoid EULA issues (I hate licenses) } static void cleanup_crt_leaks() { freeCrtMemory(); // turn off buffering for stdout/stderr and release buffers extern void* _stdbuf[2]; setvbuf(stdout, 0, _IONBF, 0); free(_stdbuf[0]); _stdbuf[0] = NULL; setvbuf(stderr, 0, _IONBF, 0); free(_stdbuf[1]); _stdbuf[1] = NULL; // free ptmcbinfo (allocated in _setmbcp) extern pthreadmbcinfo __ptmbcinfo; free(__ptmbcinfo); __ptmbcinfo = NULL; // free ioinfo (allocated in _ioinit) extern void* __pioinfo[]; free(__pioinfo[0]); __pioinfo[0] = NULL; // free ptd (allocated in _mtinit) void* _getptd_noexit(); free(_getptd_noexit()); // free TLS indices (allocated in _mtinit) and critical sections (allocated in _mtinitlocknum) void _mtterm(); _mtterm(); } static int heap_init_stub(int) { // call user init function mem_global_init(); // pretend that we initialized CRT heap (needed for _mtinitlocknum) extern HANDLE _crtheap; _crtheap = (HANDLE)-1; return 1; } static void heap_term_stub() { // cleanup memory that CRT writers are too lazy to free themselves (bastards!) cleanup_crt_leaks(); // call user init function mem_global_term(); } static void* malloc_stub(size_t size) { return mem_global_allocate(size, malloc_alignment); } static void* calloc_stub(size_t count, size_t size) { void* ptr = mem_global_allocate(count * size, malloc_alignment); if (ptr) memset(ptr, 0, mem_global_get_size(ptr)); return ptr; } static void* realloc_stub(void* ptr, size_t size) { if (!ptr) { return mem_global_allocate(size, malloc_alignment); } if (!size) { mem_global_deallocate(ptr); return 0; } size_t old_size = mem_global_get_size(ptr); void* result = mem_global_allocate(size, malloc_alignment); if (result) { memcpy(result, ptr, min(size, old_size)); mem_global_deallocate(ptr); } return result; } // Note: this function has a bug - since it calculates size to clear using // mem_global_get_size, and it can return inaccurate values (i.e. more than // allocated size), it's possible that not everything is cleared (i.e. malloc(0), // followed by recalloc(100) will not clear first 12 bytes for current implementation). // However, there is no problem if only calloc/_recalloc were used for the block in question, // and this is the case in CRT, so I won't worry about it. static void* recalloc_stub(void* ptr, size_t count, size_t size) { if (!ptr) { return calloc_stub(count, size); } if (!size) { mem_global_deallocate(ptr); return 0; } size_t old_size = mem_global_get_size(ptr); void* result = calloc_stub(count, size); if (result) { memcpy(result, ptr, min(size, old_size)); mem_global_deallocate(ptr); } return result; } static void free_stub(void* ptr) { mem_global_deallocate(ptr); } static size_t msize_stub(void* ptr) { return mem_global_get_size(ptr); } static void patch_memory_management_functions() { // heap init/term // note: we patch __crtCorExitProcess because _heap_term is called only in DLL CRT PATCH(_heap_init, heap_init_stub); PATCH(__crtCorExitProcess, heap_term_stub); // malloc PATCH(malloc, malloc_stub); PATCH(_malloc_crt, malloc_stub); PATCH_DEBUG(_malloc_base, malloc_stub); PATCH_DEBUG(_malloc_dbg, malloc_stub); // calloc PATCH(calloc, calloc_stub); PATCH(_calloc_crt, calloc_stub); PATCH_DEBUG(_calloc_base, calloc_stub); PATCH_DEBUG(_calloc_dbg, calloc_stub); // realloc PATCH(realloc, realloc_stub); PATCH(_realloc_crt, realloc_stub); PATCH_DEBUG(_realloc_base, realloc_stub); PATCH_DEBUG(_realloc_dbg, realloc_stub); // recalloc PATCH(_recalloc, recalloc_stub); PATCH(_recalloc_crt, recalloc_stub); PATCH_DEBUG(_recalloc_base, recalloc_stub); PATCH_DEBUG(_recalloc_dbg, recalloc_stub); // free PATCH(free, free_stub); PATCH_DEBUG(_free_nolock, free_stub); PATCH_DEBUG(_free_dbg_nolock, free_stub); PATCH_DEBUG(_free_base, free_stub); PATCH_DEBUG(_free_dbg, free_stub); // msize PATCH(_msize, msize_stub); PATCH_DEBUG(_msize_base, msize_stub); PATCH_DEBUG(_msize_dbg, msize_stub); // aligned functions PATCH_BREAK_ALL(_aligned_malloc); PATCH_BREAK_ALL(_aligned_offset_malloc); PATCH_BREAK_ALL(_aligned_realloc); PATCH_BREAK_ALL(_aligned_recalloc); PATCH_BREAK_ALL(_aligned_offset_realloc); PATCH_BREAK_ALL(_aligned_msize); PATCH_BREAK_ALL(_aligned_offset_recalloc); PATCH_BREAK_ALL(_aligned_free); PATCH_BREAK_ALL(_aligned_malloc); // supporting functions PATCH_BREAK_DEBUG(_nh_malloc); PATCH_BREAK_DEBUG(_nh_malloc_dbg); PATCH_BREAK_ALL(_heap_alloc); PATCH_BREAK_ALL(_expand); } int entrypoint() { int mainCRTStartup(); patch_memory_management_functions(); return mainCRTStartup(); } };