guard_allocator.odin

guard_allocator.odin

// public domain, use at your own risk

package debug_alloc

import "core:mem"
import "core:os"
import "core:fmt"
import "core:sys/win32"
import "../stacktrace"

Guard_Allocator_Data :: struct {
    verbose: bool,
    stack_trace_pages: int,
    alloc_count: int,
}

guard_allocator :: proc(verbose := false, stack_trace_pages := 0) -> mem.Allocator {
    if context.allocator.procedure == guard_allocator_proc do
        panic("cyclic initialization of the guard allocator with itself");

    alloc_data := new(Guard_Allocator_Data);
    alloc_data.verbose = verbose;
    alloc_data.stack_trace_pages = stack_trace_pages;
    alloc_data.alloc_count = 0;

    return {
        procedure = guard_allocator_proc,
        data = alloc_data,
    };
};

PAGE_GUARD :: 0x100;

guard_allocator_proc :: proc(
    allocator_data: rawptr,
    mode: mem.Allocator_Mode,
    size, alignment: int,
    old_memory: rawptr,
    old_size: int,
    flags: u64,
    location := #caller_location)
-> rawptr {

    guard_alloc_data := cast(^Guard_Allocator_Data)allocator_data;

    when ODIN_OS == "windows" {

        check_result :: proc(res: win32.Bool) {
            zero : win32.Bool;
            assert(res != zero, "virtual_protect failed");
        }

        using win32;
        switch mode {
            case .Alloc:
                page_size := os.get_page_size();

                // Here we save room for an (addr, size) pair allocated before
                // the user's block of memory, so we can free it later.
                actual_size := (size_of(uintptr)*2) + size;

                // we ask VirtualAlloc for 2 pages before and after our user block
                num_pages_needed := ((actual_size - 1) / page_size) + 1;
                valloc_size := (2 + num_pages_needed) * page_size;
                block := cast(^byte)virtual_alloc(nil, cast(uint)valloc_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
                if guard_alloc_data.verbose {
                    fmt.printf("alloc %d bytes (%d total calls)\n", valloc_size, guard_alloc_data.alloc_count);
                    guard_alloc_data.alloc_count += 1;
                }
                if guard_alloc_data.stack_trace_pages > 0 && num_pages_needed > guard_alloc_data.stack_trace_pages {
                    stacktrace.print_stack_trace();
                }

                assert(block != nil, "virtual alloc failed");

                before_page := block;
                after_page  := mem.ptr_offset(before_page, cast(int)page_size * (1 + num_pages_needed));
                
                old_protect: u32;
                new_protect :u32 = PAGE_GUARD | PAGE_READONLY;
                check_result(virtual_protect(before_page, cast(uint)page_size, new_protect, &old_protect));
                check_result(virtual_protect(after_page,  cast(uint)page_size, new_protect, &old_protect));

                ptr := mem.ptr_offset(after_page, -size);

                // We store the original virtual_alloc result and size just before the returned block memory.
                (cast(^int)mem.ptr_offset(ptr, -size_of(int)*2))^ = cast(int)cast(uintptr)block;
                (cast(^int)mem.ptr_offset(ptr, -size_of(int)*1))^ = valloc_size;

                return ptr;
            case .Free:
                orig_alloc_size := (cast(^int)mem.ptr_offset(cast(^byte)old_memory, -size_of(int)))^;
                orig_alloc_address := (cast(^rawptr)cast(uintptr)mem.ptr_offset(cast(^byte)old_memory, -size_of(int)*2))^;
                if guard_alloc_data.verbose {
                    fmt.printf("free %d bytes (%d allocs left)\n", orig_alloc_size, guard_alloc_data.alloc_count);
                    guard_alloc_data.alloc_count -= 1;
                }
                virtual_free(orig_alloc_address, cast(uint)orig_alloc_size, MEM_RELEASE);

        case .Free_All:
            assert(false, "unimplemented");

        case .Resize:
            return mem.default_resize_align(old_memory, old_size, size, alignment, context.allocator, location);

        }
    } else {
        #panic("guard_allocator is unimplemented for this OS");
        return nil;
    }

    return nil;
}

when #config(TEST_GUARD_ALLOC, false) {
    import "core:fmt"
    main :: proc() {
        context.allocator = guard_allocator;

        N :: 5000;
        rptr := mem.alloc(N);
        ptr := cast(^byte)rptr;

        ptr_to_valid_memory := mem.ptr_offset(ptr, N-1);
        os.write_string(os.stdout, "writing to valid memory...");
        ptr_to_valid_memory^ = 0;
        os.write_string(os.stdout, "...it worked!\n");

        ptr_to_invalid_memory := mem.ptr_offset(ptr, N);
        os.write_string(os.stdout, "about to crash...");
        ptr_to_invalid_memory^ = 0;
        mem.free(rptr);
        os.write_string(os.stdout, "...success! (you shouldn't see this!)");
    }
}