a guard page allocator for odin

guard_allocator.odin

package debug_alloc

import "core:mem"
import "core:os"
import "core:sys/win32"

guard_allocator := mem.Allocator {
    procedure = guard_allocator_proc,
    data = nil,
};

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 {

    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);
                when #defined(TEST_GUARD_ALLOC) do fmt.println("virtual_alloc(", block, ",", valloc_size, ")");

                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))^;
                when #defined(TEST_GUARD_ALLOC) do fmt.println("virtual_free(", orig_alloc_address, ",", orig_alloc_size, ")");
                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 #defined(TEST_GUARD_ALLOC) {
    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!)");
    }
}