Unpoison stack memory in use-after-return + use-after-scope mode

[vitalybuka]

Summary

We still want to unpoison full stack even in use-after-return as it can be disabled at runtime.

PR27453

[eugenis]

Looks reasonable, but I'd like kcc to take a look.

[kcc]

LGTM

lib/Transforms/Instrumentation/AddressSanitizer.cpp

2183–2195 ↗

(On Diff #60241)

we have a bit of code duplication here. Up to you if you want to eliminate it or keep it as is.

[vitalybuka]

Dedup some code.

[kcc]

LGTM

llvm/trunk/lib/Transforms/Instrumentation/AddressSanitizer.cpp

for (int i = 0; i <= kMaxAsanStackMallocSizeClass; i++) {
     std::string Suffix = itostr(i);
     AsanStackMallocFunc[i] = checkSanitizerInterfaceFunction(
         M.getOrInsertFunction(kAsanStackMallocNameTemplate + Suffix, IntptrTy,
                               IntptrTy, nullptr));
     AsanStackFreeFunc[i] = checkSanitizerInterfaceFunction(
         M.getOrInsertFunction(kAsanStackFreeNameTemplate + Suffix,
                               IRB.getVoidTy(), IntptrTy, IntptrTy, nullptr));
   }
+  if (ASan.UseAfterScope) {
     AsanPoisonStackMemoryFunc = checkSanitizerInterfaceFunction(
         M.getOrInsertFunction(kAsanPoisonStackMemoryName, IRB.getVoidTy(),
                               IntptrTy, IntptrTy, nullptr));
     AsanUnpoisonStackMemoryFunc = checkSanitizerInterfaceFunction(
         M.getOrInsertFunction(kAsanUnpoisonStackMemoryName, IRB.getVoidTy(),
                               IntptrTy, IntptrTy, nullptr));
+  }
+
   AsanAllocaPoisonFunc = checkSanitizerInterfaceFunction(M.getOrInsertFunction(
       kAsanAllocaPoison, IRB.getVoidTy(), IntptrTy, IntptrTy, nullptr));
   AsanAllocasUnpoisonFunc =
       checkSanitizerInterfaceFunction(M.getOrInsertFunction(
           kAsanAllocasUnpoison, IRB.getVoidTy(), IntptrTy, IntptrTy, nullptr));
 }
 
 void FunctionStackPoisoner::poisonRedZones(ArrayRef<uint8_t> ShadowBytes,
Value *BasePlus2 = IRB.CreateIntToPtr(
                     ConstantInt::get(IntptrTy, 2 * ASan.LongSize / 8)),
       IntptrPtrTy);
   IRB.CreateStore(IRB.CreatePointerCast(&F, IntptrTy), BasePlus2);
 
   // Poison the stack redzones at the entry.
   Value *ShadowBase = ASan.memToShadow(LocalStackBase, IRB);
   poisonRedZones(L.ShadowBytes, IRB, ShadowBase, true);
 
+  auto UnpoisonStack = [&](IRBuilder<> &IRB) {
+    if (HavePoisonedAllocas) {
+      // If we poisoned some allocas in llvm.lifetime analysis,
+      // unpoison whole stack frame now.
+      poisonAlloca(LocalStackBase, LocalStackSize, IRB, false);
+    } else {
+      poisonRedZones(L.ShadowBytes, IRB, ShadowBase, false);
+    }
+  };
+
   // (Un)poison the stack before all ret instructions.
   for (auto Ret : RetVec) {
     IRBuilder<> IRBRet(Ret);
     // Mark the current frame as retired.
     IRBRet.CreateStore(ConstantInt::get(IntptrTy, kRetiredStackFrameMagic),
                        BasePlus0);
     if (DoStackMalloc) {
       assert(StackMallocIdx >= 0);
if (DoStackMalloc) {
       } else {
         // For larger frames call __asan_stack_free_*.
         IRBPoison.CreateCall(
             AsanStackFreeFunc[StackMallocIdx],
             {FakeStack, ConstantInt::get(IntptrTy, LocalStackSize)});
       }
 
       IRBuilder<> IRBElse(ElseTerm);
-      poisonRedZones(L.ShadowBytes, IRBElse, ShadowBase, false);
-    } else if (HavePoisonedAllocas) {
-      // If we poisoned some allocas in llvm.lifetime analysis,
-      // unpoison whole stack frame now.
-      poisonAlloca(LocalStackBase, LocalStackSize, IRBRet, false);
+      UnpoisonStack(IRBElse);
     } else {
-      poisonRedZones(L.ShadowBytes, IRBRet, ShadowBase, false);
+      UnpoisonStack(IRBRet);
     }
   }
 
   // We are done. Remove the old unused alloca instructions.
   for (auto AI : AllocaVec) AI->eraseFromParent();
 }
 
 void FunctionStackPoisoner::poisonAlloca(Value *V, uint64_t Size,

llvm/trunk/test/Instrumentation/AddressSanitizer/lifetime-uar-uas.ll

+; Test handling of llvm.lifetime intrinsics in UAR/UAS modes.
+; RUN: opt < %s -asan -asan-module -asan-use-after-return=0 -asan-use-after-scope=0 -S | FileCheck %s
+; RUN: opt < %s -asan -asan-module -asan-use-after-return=1 -asan-use-after-scope=0 -S | FileCheck %s
+; RUN: opt < %s -asan -asan-module -asan-use-after-return=0 -asan-use-after-scope=1 -S | FileCheck %s --check-prefix=CHECK-UAS
+; RUN: opt < %s -asan -asan-module -asan-use-after-return=1 -asan-use-after-scope=1 -S | FileCheck %s --check-prefix=CHECK-UAS
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+
+declare void @llvm.lifetime.start(i64, i8* nocapture) nounwind
+declare void @llvm.lifetime.end(i64, i8* nocapture) nounwind
+
+define i32 @basic_test() sanitize_address {
+  ; CHECK-LABEL: define i32 @basic_test()
+
+entry:
+  %retval = alloca i32, align 4
+  %c = alloca i8, align 1
+
+  call void @llvm.lifetime.start(i64 1, i8* %c)
+  ; Memory is unpoisoned at llvm.lifetime.start
+  ; CHECK-UAS: call void @__asan_unpoison_stack_memory(i64 %{{[^ ]+}}, i64 1)
+
+  store volatile i32 0, i32* %retval
+  store volatile i8 0, i8* %c, align 1
+
+  call void @llvm.lifetime.end(i64 1, i8* %c)
+  ; Memory is poisoned at llvm.lifetime.end
+  ; CHECK-UAS: call void @__asan_poison_stack_memory(i64 %{{[^ ]+}}, i64 1)
+
+  ; Unpoison memory at function exit in UAS mode.
+  ; CHECK-UAS: call void @__asan_unpoison_stack_memory(i64 %{{[^ ]+}}, i64 64)
+  ; CHECK-UAS: ret void
+
+  ret i32 0
+}
+
+; No poisoning/poisoning at all in plain mode.
+; CHECK-NOT: __asan_poison_stack_memory
+; CHECK-NOT: __asan_unpoison_stack_memory

llvm/trunk/test/Instrumentation/AddressSanitizer/lifetime-uar.ll

-; Test handling of llvm.lifetime intrinsics in UAR mode.
-; RUN: opt < %s -asan -asan-module -asan-use-after-return -asan-use-after-scope -S | FileCheck %s
-
-target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
-
-declare void @llvm.lifetime.start(i64, i8* nocapture) nounwind
-declare void @llvm.lifetime.end(i64, i8* nocapture) nounwind
-
-define i32 @basic_test() sanitize_address {
-  ; CHECK-LABEL: define i32 @basic_test()
-
-entry:
-  %retval = alloca i32, align 4
-  %c = alloca i8, align 1
-
-  call void @llvm.lifetime.start(i64 1, i8* %c)
-  ; Memory is unpoisoned at llvm.lifetime.start
-  ; CHECK: call void @__asan_unpoison_stack_memory(i64 %{{[^ ]+}}, i64 1)
-
-  store volatile i32 0, i32* %retval
-  store volatile i8 0, i8* %c, align 1
-
-  call void @llvm.lifetime.end(i64 1, i8* %c)
-  ; Memory is poisoned at llvm.lifetime.end
-  ; CHECK: call void @__asan_poison_stack_memory(i64 %{{[^ ]+}}, i64 1)
-
-  ; No need to unpoison memory at function exit in UAR mode.
-  ; CHECK-NOT: @__asan_unpoison_stack_memory
-  ; CHECK: ret void
-
-  ret i32 0
-}
-