Overriding CRT heap functions to use a custom allocator.

crtheap.cpp

// 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();
    }
};