Unpoison stack on procedure return

Author: Lperlind

base/intrinsics/intrinsics.odin

@@ -368,3 +368,6 @@ valgrind_client_request :: proc(default: uintptr, request: uintptr, a0, a1, a2,
 // Internal compiler use only
 
 __entry_point :: proc() ---
+
+
+__asan_unpoison_memory_region :: proc "system" (address: rawptr, size: uint) ---

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)
+}
+

core/mem/virtual/arena.odin

@@ -82,7 +82,7 @@ arena_init_buffer :: proc(arena: ^Arena, buffer: []byte) -> (err: Allocator_Erro
 
 	arena.kind = .Buffer
 
-	// sanitizer.address_poison(buffer[:])
+	sanitizer.address_poison(buffer[:])
 
 	block_base := raw_data(buffer)
 	block := (^Memory_Block)(block_base)
@@ -182,7 +182,7 @@ arena_static_reset_to :: proc(arena: ^Arena, pos: uint, loc := #caller_location)
 			mem.zero_slice(arena.curr_block.base[arena.curr_block.used:][:prev_pos-pos])
 		}
 		arena.total_used = arena.curr_block.used
-		// sanitizer.address_poison(arena.curr_block.base[:arena.curr_block.committed])
+		sanitizer.address_poison(arena.curr_block.base[:arena.curr_block.committed])
 		return true
 	} else if pos == 0 {
 		arena.total_used = 0
@@ -349,7 +349,7 @@ arena_allocator_proc :: proc(allocator_data: rawptr, mode: mem.Allocator_Mode,
 			if size < old_size {
 				// shrink data in-place
 				data = old_data[:size]
-				// sanitizer.address_poison(old_data[size:old_size])
+				sanitizer.address_poison(old_data[size:old_size])
 				return
 			}
 
@@ -374,7 +374,7 @@ arena_allocator_proc :: proc(allocator_data: rawptr, mode: mem.Allocator_Mode,
 			return
 		}
 		copy(new_memory, old_data[:old_size])
-		// sanitizer.address_poison(old_data[:old_size])
+		sanitizer.address_poison(old_data[:old_size])
 		return new_memory, nil
 	case .Query_Features:
 		set := (^mem.Allocator_Mode_Set)(old_memory)

core/testing/runner.odin

@@ -29,7 +29,7 @@ import "core:thread"
 import "core:time"
 
 // Specify how many threads to use when running tests.
-TEST_THREADS          : int    : #config(ODIN_TEST_THREADS, 1)
+TEST_THREADS          : int    : #config(ODIN_TEST_THREADS, 0)
 // Track the memory used by each test.
 TRACKING_MEMORY       : bool   : #config(ODIN_TEST_TRACK_MEMORY, true)
 // Always report how much memory is used, even when there are no leaks or bad frees.

src/checker_builtin_procs.hpp

@@ -344,6 +344,8 @@ BuiltinProc__type_end,
 
 	BuiltinProc_valgrind_client_request,
 
+	BuiltinProc_asan_unpoison_memory_region,
+
 	BuiltinProc_COUNT,
 };
 gb_global BuiltinProc builtin_procs[BuiltinProc_COUNT] = {
@@ -687,4 +689,6 @@ gb_global BuiltinProc builtin_procs[BuiltinProc_COUNT] = {
 	{STR_LIT("wasm_memory_atomic_notify32"), 2, false, Expr_Expr, BuiltinProcPkg_intrinsics},
 
 	{STR_LIT("valgrind_client_request"), 7, false, Expr_Expr, BuiltinProcPkg_intrinsics},
+
+	{STR_LIT("__asan_unpoison_memory_region"), 2, false, Expr_Expr, BuiltinProcPkg_intrinsics},
 };

src/llvm_backend.hpp

@@ -383,6 +383,8 @@ struct lbProcedure {
 	PtrMap<Ast *, lbValue> selector_values;
 	PtrMap<Ast *, lbAddr>  selector_addr;
 	PtrMap<LLVMValueRef, lbTupleFix> tuple_fix_map;
+
+	Array<lbValue> locals;
 };
 
 

src/llvm_backend_general.cpp

@@ -3070,6 +3070,10 @@ gb_internal lbAddr lb_add_local(lbProcedure *p, Type *type, Entity *e, bool zero
 	if (e != nullptr) {
 		lb_add_entity(p->module, e, val);
 		lb_add_debug_local_variable(p, ptr, type, e->token);
+
+		if (build_context.sanitizer_flags & SanitizerFlag_Address && !p->entity->Procedure.no_sanitize_address) {
+			array_add(&p->locals, val);
+		}
 	}
 
 	if (zero_init) {

src/llvm_backend_proc.cpp

@@ -121,6 +121,7 @@ gb_internal lbProcedure *lb_create_procedure(lbModule *m, Entity *entity, bool i
 	p->branch_blocks.allocator = a;
 	p->context_stack.allocator = a;
 	p->scope_stack.allocator   = a;
+	p->locals.allocator        = a;
 	// map_init(&p->selector_values,  0);
 	// map_init(&p->selector_addr,    0);
 	// map_init(&p->tuple_fix_map,    0);
@@ -390,6 +391,7 @@ gb_internal lbProcedure *lb_create_dummy_procedure(lbModule *m, String link_name
 	p->blocks.allocator        = a;
 	p->branch_blocks.allocator = a;
 	p->context_stack.allocator = a;
+	p->locals.allocator        = a;
 	map_init(&p->tuple_fix_map, 0);
 
 

src/llvm_backend_stmt.cpp

@@ -2917,6 +2917,17 @@ gb_internal void lb_emit_defer_stmts(lbProcedure *p, lbDeferExitKind kind, lbBlo
 	}
 	defer (p->branch_location_pos = prev_token_pos);
 
+	if (kind == lbDeferExit_Return) {
+		for_array(i, p->locals) {
+			lbValue local = p->locals[i];
+
+			auto args = array_make<lbValue>(temporary_allocator(), 2);
+			args[0] = lb_emit_conv(p, local, t_rawptr);
+			args[1] = lb_const_int(p->module, t_int, type_size_of(local.type->Pointer.elem));
+			lb_emit_runtime_call(p, "__asan_unpoison_memory_region", args);
+		}
+	}
+
 	isize count = p->defer_stmts.count;
 	isize i = count;
 	while (i --> 0) {