[ASan] Disable dynamic alloca and UAR detection in presence of returns_twice calls.

[samsonov]

Summary

returns_twice (most importantly, setjmp) functions are
optimization-hostile: if local variable is promoted to register, and is
changed between setjmp() and longjmp() calls, this update will be
undone. This is the reason why "man setjmp" advises to mark all these
locals as "volatile".

This can not be enough for ASan, though: when it replaces static alloca
with dynamic one, optionally called if UAR mode is enabled, it adds a
whole lot of SSA values, and computations of local variable addresses,
that can involve virtual registers, and cause unexpected behavior, when
these registers are restored from buffer saved in setjmp.

To fix this, just disable dynamic alloca and UAR tricks whenever we see
a returns_twice call in the function.

[rnk]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

685

It is possible to invoke setjmp, so you should change this to visitCallSite(CallSite CS) to handle both.

[samsonov]

• Use "visitCallSite" to handle both call and invoke.

[samsonov]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

685

Done

[kcc]

LGTM

lib/Transforms/Instrumentation/AddressSanitizer.cpp

struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   SmallVector<IntrinsicInst *, 1> StackRestoreVec;
   AllocaInst *DynamicAllocaLayout = nullptr;
   IntrinsicInst *LocalEscapeCall = nullptr;
 
   // Maps Value to an AllocaInst from which the Value is originated.
   typedef DenseMap<Value *, AllocaInst *> AllocaForValueMapTy;
   AllocaForValueMapTy AllocaForValue;
 
-  bool HasNonEmptyInlineAsm;
+  bool HasNonEmptyInlineAsm = false;
+  bool HasReturnsTwiceCall = false;
   std::unique_ptr<CallInst> EmptyInlineAsm;
 
   FunctionStackPoisoner(Function &F, AddressSanitizer &ASan)
       : F(F),
         ASan(ASan),
         DIB(*F.getParent(), /*AllowUnresolved*/ false),
         C(ASan.C),
         IntptrTy(ASan.IntptrTy),
         IntptrPtrTy(PointerType::get(IntptrTy, 0)),
         Mapping(ASan.Mapping),
         StackAlignment(1 << Mapping.Scale),
-        HasNonEmptyInlineAsm(false),
         EmptyInlineAsm(CallInst::Create(ASan.EmptyAsm)) {}
 
   bool runOnFunction() {
     if (!ClStack) return false;
     // Collect alloca, ret, lifetime instructions etc.
     for (BasicBlock *BB : depth_first(&F.getEntryBlock())) visit(*BB);
 
     if (AllocaVec.empty() && DynamicAllocaVec.empty()) return false;
void visitIntrinsicInst(IntrinsicInst &II) {
     // Find alloca instruction that corresponds to llvm.lifetime argument.
     AllocaInst *AI = findAllocaForValue(II.getArgOperand(1));
     if (!AI) return;
     bool DoPoison = (ID == Intrinsic::lifetime_end);
     AllocaPoisonCall APC = {&II, AI, SizeValue, DoPoison};
     AllocaPoisonCallVec.push_back(APC);
   }
 
   void visitCallInst(CallInst &CI) {
     HasNonEmptyInlineAsm |=
         CI.isInlineAsm() && !CI.isIdenticalTo(EmptyInlineAsm.get());
+    HasReturnsTwiceCall |= CI.canReturnTwice();
   }
 
   // ---------------------- Helpers.
   void initializeCallbacks(Module &M);
 
   bool doesDominateAllExits(const Instruction *I) const {
     for (auto Ret : RetVec) {
       if (!ASan.getDominatorTree().dominates(I, Ret)) return false;
void FunctionStackPoisoner::poisonStack() {
   // i.e. 32 bytes on 64-bit platforms and 16 bytes in 32-bit platforms.
   size_t MinHeaderSize = ASan.LongSize / 2;
   ASanStackFrameLayout L;
   ComputeASanStackFrameLayout(SVD, 1UL << Mapping.Scale, MinHeaderSize, &L);
   DEBUG(dbgs() << L.DescriptionString << " --- " << L.FrameSize << "\n");
   uint64_t LocalStackSize = L.FrameSize;
   bool DoStackMalloc = ClUseAfterReturn && !ASan.CompileKernel &&
                        LocalStackSize <= kMaxStackMallocSize;
-  // Don't do dynamic alloca or stack malloc in presence of inline asm:
-  // too often it makes assumptions on which registers are available.
-  bool DoDynamicAlloca = ClDynamicAllocaStack && !HasNonEmptyInlineAsm;
-  DoStackMalloc &= !HasNonEmptyInlineAsm;
+  bool DoDynamicAlloca = ClDynamicAllocaStack;
+  // Don't do dynamic alloca or stack malloc if:
+  // 1) There is inline asm: too often it makes assumptions on which registers
+  //    are available.
+  // 2) There is a returns_twice call (typically setjmp), which is
+  //    optimization-hostile, and doesn't play well with introduced indirect
+  //    register-relative calculation of local variable addresses.
+  DoDynamicAlloca &= !HasNonEmptyInlineAsm && !HasReturnsTwiceCall;
+  DoStackMalloc &= !HasNonEmptyInlineAsm && !HasReturnsTwiceCall;
 
   Value *StaticAlloca =
       DoDynamicAlloca ? nullptr : createAllocaForLayout(IRB, L, false);
 
   Value *FakeStack;
   Value *LocalStackBase;
 
   if (DoStackMalloc) {

test/Instrumentation/AddressSanitizer/stack_dynamic_alloca.ll

 ; CHECK: ret void
 
   %XXX = alloca [20 x i8], align 1
   %arr.ptr = bitcast [20 x i8]* %XXX to i8*
   store volatile i8 0, i8* %arr.ptr
   call void asm sideeffect "mov %%rbx, %%rcx", "~{dirflag},~{fpsr},~{flags}"() nounwind
   ret void
 }
+
+; Test that dynamic alloca is not used when setjmp is present.
+%struct.__jmp_buf_tag = type { [8 x i64], i32, %struct.__sigset_t }
+%struct.__sigset_t = type { [16 x i64] }
+@_ZL3buf = internal global [1 x %struct.__jmp_buf_tag] zeroinitializer, align 16
+
+define void @Func3() uwtable sanitize_address {
+; CHECK-LABEL: define void @Func3
+; CHECK-NOT: __asan_option_detect_stack_use_after_return
+; CHECK-NOT: __asan_stack_malloc
+; CHECK: call void @__asan_handle_no_return
+; CHECK: call void @longjmp
+; CHECK: ret void
+entry:
+  %a = alloca i32, align 4
+  %call = call i32 @_setjmp(%struct.__jmp_buf_tag* getelementptr inbounds ([1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* @_ZL3buf, i32 0, i32 0)) nounwind returns_twice
+  %cmp = icmp eq i32 0, %call
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @longjmp(%struct.__jmp_buf_tag* getelementptr inbounds ([1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* @_ZL3buf, i32 0, i32 0), i32 1) noreturn nounwind
+  unreachable
+
+if.end:                                           ; preds = %entry
+  call void @_Z10escape_ptrPi(i32* %a)
+  ret void
+}
+
+declare i32 @_setjmp(%struct.__jmp_buf_tag*) nounwind returns_twice
+declare void @longjmp(%struct.__jmp_buf_tag*, i32) noreturn nounwind
+declare void @_Z10escape_ptrPi(i32*)