Add asan support for various allocators

3 files changed, 27 insertions, 2 deletions

diff --git a/base/runtime/default_temp_allocator_arena.odin b/base/runtime/default_temp_allocator_arena.odin
index 5f25dac95..74994344a 100644
--- a/base/runtime/default_temp_allocator_arena.odin
+++ b/base/runtime/default_temp_allocator_arena.odin
@@ -1,6 +1,7 @@
 package runtime
 
 import "base:intrinsics"
+import "base:sanitizer"
 
 DEFAULT_ARENA_GROWING_MINIMUM_BLOCK_SIZE :: uint(DEFAULT_TEMP_ALLOCATOR_BACKING_SIZE)
 
@@ -43,6 +44,8 @@ memory_block_alloc :: proc(allocator: Allocator, capacity: uint, alignment: uint
 	block.base = ([^]byte)(uintptr(block) + base_offset)
 	block.capacity = uint(end - uintptr(block.base))
 
+	sanitizer.address_poison(block.base, block.capacity)
+
 	// Should be zeroed
 	assert(block.used == 0)
 	assert(block.prev == nil)
@@ -52,6 +55,7 @@ memory_block_alloc :: proc(allocator: Allocator, capacity: uint, alignment: uint
 memory_block_dealloc :: proc(block_to_free: ^Memory_Block, loc := #caller_location) {
 	if block_to_free != nil {
 		allocator := block_to_free.allocator
+		sanitizer.address_unpoison(block_to_free.base, block_to_free.capacity)
 		mem_free(block_to_free, allocator, loc)
 	}
 }
@@ -83,6 +87,7 @@ alloc_from_memory_block :: proc(block: ^Memory_Block, min_size, alignment: uint)
 		return
 	}
 	data = block.base[block.used+alignment_offset:][:min_size]
+	sanitizer.address_unpoison(block.base[block.used:block.used+size])
 	block.used += size
 	return
 }
@@ -162,6 +167,7 @@ arena_free_all :: proc(arena: ^Arena, loc := #caller_location) {
 	if arena.curr_block != nil {
 		intrinsics.mem_zero(arena.curr_block.base, arena.curr_block.used)
 		arena.curr_block.used = 0
+		sanitizer.address_poison(arena.curr_block.base, arena.curr_block.capacity)
 	}
 	arena.total_used = 0
 }
@@ -226,6 +232,7 @@ arena_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 					// grow data in-place, adjusting next allocation
 					block.used = uint(new_end)
 					data = block.base[start:new_end]
+					sanitizer.address_unpoison(data)
 					return
 				}
 			}
@@ -299,6 +306,7 @@ arena_temp_end :: proc(temp: Arena_Temp, loc := #caller_location) {
 			assert(block.used >= temp.used, "out of order use of arena_temp_end", loc)
 			amount_to_zero := block.used-temp.used
 			intrinsics.mem_zero(block.base[temp.used:], amount_to_zero)
+			sanitizer.address_poison(block.base[temp.used:block.capacity])
 			block.used = temp.used
 			arena.total_used -= amount_to_zero
 		}
diff --git a/base/runtime/heap_allocator_windows.odin b/base/runtime/heap_allocator_windows.odin
index e07df7559..04a6f149b 100644
--- a/base/runtime/heap_allocator_windows.odin
+++ b/base/runtime/heap_allocator_windows.odin
@@ -1,5 +1,7 @@
 package runtime
 
+import "../sanitizer"
+
 foreign import kernel32 "system:Kernel32.lib"
 
 @(private="file")
@@ -16,7 +18,10 @@ foreign kernel32 {
 
 _heap_alloc :: proc "contextless" (size: int, zero_memory := true) -> rawptr {
 	HEAP_ZERO_MEMORY :: 0x00000008
-	return HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY if zero_memory else 0, uint(size))
+	ptr := HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY if zero_memory else 0, uint(size))
+	// NOTE(lucas): asan not guarunteed to unpoison win32 heap out of the box, do it ourselves
+	sanitizer.address_unpoison(ptr, size)
+	return ptr
 }
 _heap_resize :: proc "contextless" (ptr: rawptr, new_size: int) -> rawptr {
 	if new_size == 0 {
@@ -28,7 +33,10 @@ _heap_resize :: proc "contextless" (ptr: rawptr, new_size: int) -> rawptr {
 	}
 
 	HEAP_ZERO_MEMORY :: 0x00000008
-	return HeapReAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, ptr, uint(new_size))
+	new_ptr := HeapReAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, ptr, uint(new_size))
+	// NOTE(lucas): asan not guarunteed to unpoison win32 heap out of the box, do it ourselves
+	sanitizer.address_unpoison(new_ptr, new_size)
+	return new_ptr
 }
 _heap_free :: proc "contextless" (ptr: rawptr) {
 	if ptr == nil {
diff --git a/base/runtime/internal.odin b/base/runtime/internal.odin
index 59811a525..7c8a8294b 100644
--- a/base/runtime/internal.odin
+++ b/base/runtime/internal.odin
@@ -1106,3 +1106,12 @@ __read_bits :: proc "contextless" (dst, src: [^]byte, offset: uintptr, size: uin
 		dst[j>>3]  |= the_bit<<(j&7)
 	}
 }
+
+@(no_sanitize_address)
+__asan_unpoison_memory_region :: #force_inline proc "contextless" (address: rawptr, size: uint) {
+	foreign {
+		__asan_unpoison_memory_region :: proc(address: rawptr, size: uint) ---
+	}
+	__asan_unpoison_memory_region(address, size)
+}
+