Debug allocator that catches leaks, use-after-free, double-free, buffer overrun, invalid/mismatched free. Windows x64. Simple but slow (no page reuse).

dead_simple_debug_allocator.c

// libs: Kernel32.lib User32.lib DbgHelp.lib msvcrt.lib (unless you have your own printf)
// includes: Windows.h (LEAN_AND_MEAN is ok) WinUser.h DbgHelp.h stdio.h

// call SymInitialize() at startup and debug_check_leaks at shutdown

#ifndef NDEBUG

static void print_stack_trace(void **stack, u64 stack_count, const char *indent) {
    struct {
        SYMBOL_INFO si;
        char name[256];
    } symbol = {{sizeof(symbol.si)}};
    symbol.si.MaxNameLen = COUNTOF(symbol.name);
	
    IMAGEHLP_LINE64 line = {sizeof(line)};
	
    for (int i = 0; i < stack_count; i += 1) {
        unsigned long long dummy1 = 0;
        if (!SymFromAddr(GetCurrentProcess(), (u64)stack[i], &dummy1, &symbol.si)) {
            printf("%s(unknown: 0x%0.16p)\n", indent, stack[i]);
            continue;
        }
        printf("%s%s", indent, symbol.si.Name);
        
        unsigned long dummy2 = 0;
        if (SymGetLineFromAddr64(GetCurrentProcess(), (u64)stack[i], &dummy2, &line)) {
            printf(" (%s:%d)", line.FileName, line.LineNumber);
        }
        printf("\n");
    }
}

static void debug_free_stack_trace(const char *error, void *base) {
    void *stack[512];
    printf("%s  Stack trace of free:\n", error);
    print_stack_trace(stack, RtlCaptureStackBackTrace(2, 512, stack, 0), "    ");
    if (base) {
        printf("\n  Stack trace of allocation:\n");
        print_stack_trace((void **)base + 4, ((u64 *)base)[3], "    ");
    }
    printf("\n");
}

static void *debug_last_allocation;
static void *debug_alloc(u64 n) { // @Todo: mutex // @Todo: different page size
    u64 rounded = (n + 4095) & ~4095ull;
    u64 *base = (u64 *)VirtualAlloc(0, rounded + 8192, MEM_RESERVE, PAGE_NOACCESS);
    base = (u64 *)VirtualAlloc(base, rounded + 4096, MEM_COMMIT, PAGE_READWRITE);
    base[0] = n;
    base[1] = (u64)debug_last_allocation;
    base[3] = RtlCaptureStackBackTrace(1, 512 - 4, (void **)base + 4, 0);
    debug_last_allocation = base;
    return (char *)base + 4096 + (rounded - n);
}

static int debug_free(void *block, u64 n) {
    MEMORY_BASIC_INFORMATION mbi;
    VirtualQuery(block, &mbi, sizeof(mbi));
    if (mbi.State == MEM_FREE) {
        debug_free_stack_trace("Error: Invalid free!\n", 0);
        return 0;
    }
    if (mbi.State == MEM_RESERVE) {
        debug_free_stack_trace("Error: Double-free!\n", mbi.AllocationBase);
        return 0;
    }
    bool print = false;
    u64 real_n = ((u64 *)mbi.AllocationBase)[0];
    u64 rounded = (real_n + 4095) & ~4095ull;
    if (block != (char *)mbi.AllocationBase + 4096 + (rounded - real_n)) {
        printf("Warning: Freeing wrong base pointer.\n");
        print = true;
    }
    if (n != real_n) {
        printf("Warning: Freeing wrong block size.\n");
        print = true;
    }
    if (print) {
        debug_free_stack_trace("", mbi.AllocationBase);
    }
    if (!rounded) { // 0-size block
        return 1; // success, can't decommit 0 pages.
    }
    return VirtualFree((char *)mbi.AllocationBase + 4096, rounded, MEM_DECOMMIT);
}

static void *debug_mark_leaked(void *block) {
    MEMORY_BASIC_INFORMATION mbi;
    VirtualQuery(block, &mbi, sizeof(mbi));
    if (mbi.State & (MEM_COMMIT | MEM_RESERVE)) {
        ((u64 *)mbi.AllocationBase)[2] = 1;
    }
    return block;
}

static void debug_check_leaks(void) {
    int num_leaks = 0;
    for (void *base = debug_last_allocation; base; base = ((void **)base)[1]) {
        MEMORY_BASIC_INFORMATION mbi;
        VirtualQuery(base, &mbi, sizeof(mbi));
        
        assert(mbi.State == MEM_COMMIT);
        if (mbi.RegionSize > 4096) { // Committed area extends beyond header page
            u64 n = ((u64 *)base)[0];
            u64 flag = ((u64 *)base)[2];
            if (flag == 0) {
                u64 rounded = (n + 4095) & ~4095;
                void *user_facing = (char *)base + 4096 + (rounded - n);
                printf("Leaked %llu bytes at 0x%0.16p!\n  Stack trace:\n", n, user_facing);
                print_stack_trace((void **)base + 4, ((u64 *)base)[3], "    ");
                num_leaks += 1;
            }
        }
    }
    if (num_leaks) {
        char buf[512];
        snprintf(buf, sizeof(buf), "%d memory leaks!", num_leaks);
        MessageBoxA(0, buf, "Error", 0);
    }
}

#define ALLOC(size) debug_alloc(size)
#define LEAK(pointer) ((void)debug_mark_leaked(pointer))
#define FREE(pointer, size) (void)(debug_free(pointer, size) || (__debugbreak(), 0))

#else

#define ALLOC(size) release_alloc(size)
#define LEAK(pointer) ((void)(pointer))
#define FREE(pointer, size) release_free(pointer, size)

#endif

// Public domain; no warranty implied.

// Basically all credit goes to Niklas Gray (https://twitter.com/niklasfrykholm)
// for this post http://ourmachinery.com/post/virtual-memory-tricks/