Allocator which maps each allocation to a separate page, and locks that page immediately on free

strict_allocator.odin

package util

import "base:runtime"
import "base:intrinsics"

import "core:mem"
import "core:mem/virtual"
import "core:os"
import "core:fmt"

// Ensures that freeing or resizing renders the original pointer invalid, such that UAFs are guarenteed to crash.
// This is achieved by making a seperate allocation for each and every allocation request.
Strict_Allocator :: struct {
    pointers: map[rawptr]^Block, // pointer -> size
    frees: map[rawptr]runtime.Source_Code_Location,
}

strict_allocator :: proc(state: ^Strict_Allocator) -> runtime.Allocator {
	return {
		procedure = strict_allocator_proc,
		data = state,
	}
}

strict_allocator_init :: proc(state: ^Strict_Allocator, ally: mem.Allocator) {
	state^ = {}
	state.pointers.allocator = ally
	state.frees.allocator = ally
}

strict_allocator_destroy :: proc(state: ^Strict_Allocator) {
	free_all(strict_allocator(state))
	delete(state.pointers)
	delete(state.frees)
}

page_size := i64(os.get_page_size())

strict_allocator_proc :: proc(allocator_data: rawptr, mode: runtime.Allocator_Mode,
	size, alignment: int,
	old_memory: rawptr, old_size: int, loc := #caller_location) -> (memory: []byte, err: runtime.Allocator_Error) {
	state := (^Strict_Allocator)(allocator_data)


	// NOTE: On freeing, don't unmap the old memory pages; this means that it won't be remapped again.
	// This will help to catch invalid accesses on long-dead memory, as accessing them will immediately crash.
	switch mode {
	case .Alloc, .Alloc_Non_Zeroed:
		block := strict_allocator_alloc(state, size, loc) or_return
		return block_user_data(block), nil
	case .Free:
		strict_allocator_free(state, old_memory, old_size, loc)
		return nil, nil
	case .Free_All:
		for _, block in state.pointers {
			if block == nil { continue }
			block_destroy(block)
		}
		clear(&state.frees)
		clear(&state.pointers)
	case .Resize, .Resize_Non_Zeroed:
		new_block := strict_allocator_alloc(state, size, loc) or_return
		if old_memory != nil {
			old_block, belongs_to_us := state.pointers[old_memory]
			if belongs_to_us {
				copy(block_user_data(new_block), block_user_data(old_block))
			}
			strict_allocator_free(state, old_memory, old_size, loc)
		}
		return block_user_data(new_block), nil
	case .Query_Features, .Query_Info:
		return nil, .Mode_Not_Implemented
	}

	return nil, nil
}

strict_allocator_alloc :: proc(state: ^Strict_Allocator, n: int, loc: runtime.Source_Code_Location) -> (block: ^Block, err: mem.Allocator_Error) {
	block = block_init(n, loc) or_return
	state.pointers[block.user_data] = block
	return block, nil
}

strict_allocator_free :: proc(state: ^Strict_Allocator, user_data: rawptr, user_size: int, loc: runtime.Source_Code_Location) {
	block, belongs_to_us := state.pointers[user_data]
	if !belongs_to_us {
		fmt.panicf("memory being freed does not belong to this allocator (%p)", user_data, loc = loc)
	}

	freed_loc, has_been_freed := state.frees[user_data]
	if has_been_freed {
		fmt.panicf("memory has already been freed (%p) at %v", user_data, freed_loc, loc = loc)
	}

	assert(block.user_data == user_data)
	assert(block._magic1 == BLOCK_START_MAGIC)
	assert(block._magic2 == BLOCK_END_MAGIC)

	user_data_offset := mem.ptr_sub((^byte)(user_data), (^byte)(block.block_data))
	user_data_size := block.rounded_size - user_data_offset
	fmt.assertf(user_data_size == user_size, "attempt to free with wrong size: %v, but should be %v (%p)", user_size, user_data_size, user_data)

	delete_key(&state.pointers, user_data)
	state.frees[user_data] = loc
	block_destroy(block)
	// NOTE: don't release, only decommit so that this page isn't handed out again.
}

Block :: struct {
	_magic1:        u64,
	block_data:     rawptr,
	rounded_size:   int,
	user_data:      rawptr,
	user_size:      int,
	allocated_at:   runtime.Source_Code_Location,
	_magic2:        u64,
}
BLOCK_START_MAGIC :: 0xaaaabbbbccccdddd
BLOCK_END_MAGIC   :: 0xeeeeffffaaaabbbb

block_user_data :: proc(b: ^Block) -> []byte {
	p := ([^]byte)(b.user_data)
	return p[:b.user_size]
}

block_init :: proc(size: int, loc: runtime.Source_Code_Location) -> (block: ^Block, err: mem.Allocator_Error) {
	rounded_size := align_forward_i64(i64(size) + size_of(Block), page_size)
	memory := virtual.reserve_and_commit(uint(rounded_size)) or_return
	user_part := memory[len(memory) - size:]
	block = mem.ptr_offset((^Block)(raw_data(user_part)), -1)
	block^ = {
		_magic1 = BLOCK_START_MAGIC,
		_magic2 = BLOCK_END_MAGIC,
		block_data = raw_data(memory),
		rounded_size = int(rounded_size),
		user_data = raw_data(user_part),
		user_size = size,
		allocated_at = loc,
	}
	return block, nil
}

block_destroy :: proc(b: ^Block) {
	virtual.decommit(b.block_data, cast(uint) b.rounded_size)
}