Skip to content

Instantly share code, notes, and snippets.

@kevinw
Last active June 26, 2021 04:04
Show Gist options
  • Save kevinw/87aa506af8e9f7bb966bc10eedbff21c to your computer and use it in GitHub Desktop.
Save kevinw/87aa506af8e9f7bb966bc10eedbff21c to your computer and use it in GitHub Desktop.
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!)");
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment