Merge pull request #5122 from Lperlind/asan-allocators

Add asan support for various allocators and stack unpoisoning

4 files changed, 34 insertions, 11 deletions

diff --git a/src/llvm_backend.hpp b/src/llvm_backend.hpp
index 6177fcf6e..de6841ed8 100644
--- a/src/llvm_backend.hpp
+++ b/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> asan_stack_locals;
 };
 
 
diff --git a/src/llvm_backend_general.cpp b/src/llvm_backend_general.cpp
index 421720c4c..dad5d4dd5 100644
--- a/src/llvm_backend_general.cpp
+++ b/src/llvm_backend_general.cpp
@@ -3070,6 +3070,13 @@ 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);
+
+		// NOTE(lucas): In LLVM 20 and below we do not have the option to have asan cleanup poisoned stack
+		// locals ourselves. So we need to manually track and unpoison these locals on proc return.
+		// LLVM 21 adds the 'use-after-scope' asan option which does this for us.
+		if (build_context.sanitizer_flags & SanitizerFlag_Address && !p->entity->Procedure.no_sanitize_address) {
+			array_add(&p->asan_stack_locals, val);
+		}
 	}
 
 	if (zero_init) {
diff --git a/src/llvm_backend_proc.cpp b/src/llvm_backend_proc.cpp
index 14157455e..1f023037d 100644
--- a/src/llvm_backend_proc.cpp
+++ b/src/llvm_backend_proc.cpp
@@ -115,12 +115,13 @@ gb_internal lbProcedure *lb_create_procedure(lbModule *m, Entity *entity, bool i
 	p->is_entry_point = false;
 
 	gbAllocator a = heap_allocator();
-	p->children.allocator      = a;
-	p->defer_stmts.allocator   = a;
-	p->blocks.allocator        = a;
-	p->branch_blocks.allocator = a;
-	p->context_stack.allocator = a;
-	p->scope_stack.allocator   = a;
+	p->children.allocator          = a;
+	p->defer_stmts.allocator       = a;
+	p->blocks.allocator            = a;
+	p->branch_blocks.allocator     = a;
+	p->context_stack.allocator     = a;
+	p->scope_stack.allocator       = a;
+	p->asan_stack_locals.allocator = a;
 	// map_init(&p->selector_values,  0);
 	// map_init(&p->selector_addr,    0);
 	// map_init(&p->tuple_fix_map,    0);
@@ -385,11 +386,12 @@ gb_internal lbProcedure *lb_create_dummy_procedure(lbModule *m, String link_name
 	p->is_entry_point = false;
 
 	gbAllocator a = permanent_allocator();
-	p->children.allocator      = a;
-	p->defer_stmts.allocator   = a;
-	p->blocks.allocator        = a;
-	p->branch_blocks.allocator = a;
-	p->context_stack.allocator = a;
+	p->children.allocator          = a;
+	p->defer_stmts.allocator       = a;
+	p->blocks.allocator            = a;
+	p->branch_blocks.allocator     = a;
+	p->context_stack.allocator     = a;
+	p->asan_stack_locals.allocator = a;
 	map_init(&p->tuple_fix_map, 0);
 
 
diff --git a/src/llvm_backend_stmt.cpp b/src/llvm_backend_stmt.cpp
index 96a5d0db1..89737a454 100644
--- a/src/llvm_backend_stmt.cpp
+++ b/src/llvm_backend_stmt.cpp
@@ -2917,6 +2917,18 @@ gb_internal void lb_emit_defer_stmts(lbProcedure *p, lbDeferExitKind kind, lbBlo
 	}
 	defer (p->branch_location_pos = prev_token_pos);
 
+	// TODO(lucas): In LLVM 21 use the 'use-after-scope' asan option which does this for us.
+	if (kind == lbDeferExit_Return) {
+		for_array(i, p->asan_stack_locals) {
+			lbValue local = p->asan_stack_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) {