[llvm] [Sanitizers] Don't inline unpoisoning of small stacks when inlining disabled (PR #75555)

[Mariusz Borsa via llvm-commits]

https://github.com/wrotki created https://github.com/llvm/llvm-project/pull/75555

When ASan.MaxInlinePoisoningSize == 0 , it means that no shadow memory operations should be made via inlined instrumentation code,
but only via calls to shadow setting functions. This change fixes one violation of this, which happened when the function allocas count
was small, i.e. less than 5 - in the code modifying the shadow just before ret instruction.
We now explicitly check ASan.MaxInlinePoisoningSize , and if it's 0 then we disallow inlining. It is required for the instrumentation
emitting code suitable for handling by ABI implementation.

rdar://119513720


>From 0f7a520ac855652ba212bd0f87ce6620f7b48c20 Mon Sep 17 00:00:00 2001
From: Mariusz Borsa <m_borsa at apple.com>
Date: Thu, 14 Dec 2023 18:08:28 -0800
Subject: [PATCH] [Sanitizers] Don't inline unpoisoning of small stacks when
 inlining disabled

When ASan.MaxInlinePoisoningSize == 0 , it means that no shadow memory operations should be made via inlined instrumentation code,
but only via calls to shadow setting functions. This change fixes one violation of this, which happened when the function allocas count
was small, i.e. less than 5 - in the code modifying the shadow just before ret instruction.
We now explicitly check ASan.MaxInlinePoisoningSize , and if it's 0 then we disallow inlining. It is required for the instrumentation
emitting code suitable for handling by ABI implementation.

rdar://119513720
---
 .../Instrumentation/AddressSanitizer.cpp      |  2 +-
 .../AddressSanitizer/calls-only-smallfn.ll    | 28 +++++++++++++++++++
 .../AddressSanitizer/calls-only.ll            |  4 +--
 3 files changed, 31 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Instrumentation/AddressSanitizer/calls-only-smallfn.ll

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index b175e6f93f3e8f..6468d07b4f4f45 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -3505,7 +3505,7 @@ void FunctionStackPoisoner::processStaticAllocas() {
       SplitBlockAndInsertIfThenElse(Cmp, Ret, &ThenTerm, &ElseTerm);
 
       IRBuilder<> IRBPoison(ThenTerm);
-      if (StackMallocIdx <= 4) {
+      if (ASan.MaxInlinePoisoningSize != 0 && StackMallocIdx <= 4) {
         int ClassSize = kMinStackMallocSize << StackMallocIdx;
         ShadowAfterReturn.resize(ClassSize / L.Granularity,
                                  kAsanStackUseAfterReturnMagic);
diff --git a/llvm/test/Instrumentation/AddressSanitizer/calls-only-smallfn.ll b/llvm/test/Instrumentation/AddressSanitizer/calls-only-smallfn.ll
new file mode 100644
index 00000000000000..3d67778049430f
--- /dev/null
+++ b/llvm/test/Instrumentation/AddressSanitizer/calls-only-smallfn.ll
@@ -0,0 +1,28 @@
+; RUN: opt < %s -passes=asan -asan-max-inline-poisoning-size=0   -asan-stack-dynamic-alloca=0 -S | FileCheck --check-prefix=OUTLINE %s
+; RUN: opt < %s -passes=asan -asan-max-inline-poisoning-size=999 -asan-stack-dynamic-alloca=0 -S | FileCheck --check-prefix=INLINE  %s
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-macosx13.0.0"
+
+; Function Attrs: noinline nounwind optnone sanitize_address ssp uwtable(sync)
+define void @foo() #0 {
+entry:
+  %array01 = alloca [1 x i8], align 1
+  %array02 = alloca [2 x i8], align 1
+; OUTLINE:  call void @__asan_set_shadow_f1(i64 %23, i64 4)
+; OUTLINE:  call void @__asan_set_shadow_01(i64 %24, i64 1)
+; OUTLINE:  call void @__asan_set_shadow_f2(i64 %25, i64 1)
+; OUTLINE:  call void @__asan_set_shadow_02(i64 %26, i64 1)
+; OUTLINE:  call void @__asan_set_shadow_f3(i64 %27, i64 1)
+; OUTLINE:  call void @__asan_stack_free_0(i64 %7, i64 64)
+; OUTLINE:  call void @__asan_set_shadow_00(i64 %55, i64 8)
+; INLINE:  store i64 -935919682371587599, ptr %24, align 1
+; INLINE:  store i64 -723401728380766731, ptr %52, align 1
+  %arrayidx = getelementptr inbounds [1 x i8], ptr %array01, i64 0, i64 1
+  store i8 1, ptr %arrayidx, align 1
+  %arrayidx1 = getelementptr inbounds [2 x i8], ptr %array02, i64 0, i64 2
+  store i8 2, ptr %arrayidx1, align 1
+  ret void
+}
+attributes #0 = { noinline nounwind optnone sanitize_address ssp uwtable(sync) "frame-pointer"="non-leaf" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="apple-m1" "target-features"="+aes,+crc,+crypto,+dotprod,+fp-armv8,+fp16fml,+fullfp16,+lse,+neon,+ras,+rcpc,+rdm,+sha2,+sha3,+sm4,+v8.1a,+v8.2a,+v8.3a,+v8.4a,+v8.5a,+v8a,+zcm,+zcz" }
+
diff --git a/llvm/test/Instrumentation/AddressSanitizer/calls-only.ll b/llvm/test/Instrumentation/AddressSanitizer/calls-only.ll
index 2cf0070cf08623..fa491105e017d7 100644
--- a/llvm/test/Instrumentation/AddressSanitizer/calls-only.ll
+++ b/llvm/test/Instrumentation/AddressSanitizer/calls-only.ll
@@ -29,8 +29,8 @@ entry:
 ; OUTLINE:  call void @__asan_set_shadow_f2(i64 %45, i64 3)
 ; OUTLINE:  call void @__asan_set_shadow_07(i64 %46, i64 1)
 ; OUTLINE:  call void @__asan_set_shadow_f3(i64 %47, i64 3)
-; OUTLINE:  call void @__asan_set_shadow_f5(i64 %134, i64 32)
-; OUTLINE:  call void @__asan_set_shadow_00(i64 %140, i64 24)
+; OUTLINE:  call void @__asan_stack_free_2(i64 %7, i64 192)
+; OUTLINE:  call void @__asan_set_shadow_00(i64 %135, i64 24)
 ; INLINE:  store i64 -1007977276409515535, ptr %34, align 1
 ; INLINE:  store i64 -940423264817843709, ptr %36, align 1
 ; INLINE:  store i64 -868083087686045178, ptr %38, align 1