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

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

PR27453

Reviewers: eugenis, kcc

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D21202

llvm-svn: 272334

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -1869,12 +1869,15 @@ void FunctionStackPoisoner::initializeCallbacks(Module &M) {
         M.getOrInsertFunction(kAsanStackFreeNameTemplate + Suffix,
                               IRB.getVoidTy(), IntptrTy, IntptrTy, nullptr));
   }
-  AsanPoisonStackMemoryFunc = checkSanitizerInterfaceFunction(
-      M.getOrInsertFunction(kAsanPoisonStackMemoryName, IRB.getVoidTy(),
-                            IntptrTy, IntptrTy, nullptr));
-  AsanUnpoisonStackMemoryFunc = checkSanitizerInterfaceFunction(
-      M.getOrInsertFunction(kAsanUnpoisonStackMemoryName, 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 =
@@ -2133,6 +2136,16 @@ void FunctionStackPoisoner::poisonStack() {
   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);
@@ -2177,13 +2190,9 @@ void FunctionStackPoisoner::poisonStack() {
       }
 
       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);
     }
   }
 

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

@@ -0,0 +1,39 @@
+; 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