Skip to content

Instantly share code, notes, and snippets.

@kevinw
Last active June 26, 2021 04:04

Revisions

  1. kevinw revised this gist Nov 28, 2020. 1 changed file with 2 additions and 0 deletions.
    2 changes: 2 additions & 0 deletions guard_allocator.odin
    Original file line number Diff line number Diff line change
    @@ -1,3 +1,5 @@
    // public domain, use at your own risk

    package debug_alloc

    import "core:mem"
  2. kevinw renamed this gist Nov 28, 2020. 1 changed file with 0 additions and 0 deletions.
    File renamed without changes.
  3. kevinw created this gist Nov 28, 2020.
    133 changes: 133 additions & 0 deletions gistfile1.txt
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,133 @@
    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!)");
    }
    }