diff --git a/include/llvm/Transforms/Utils/BasicBlockUtils.h b/include/llvm/Transforms/Utils/BasicBlockUtils.h index 6bdf53facb6..cab0d5f29c2 100644 --- a/include/llvm/Transforms/Utils/BasicBlockUtils.h +++ b/include/llvm/Transforms/Utils/BasicBlockUtils.h @@ -205,6 +205,26 @@ TerminatorInst *SplitBlockAndInsertIfThen(Value *Cond, Instruction *SplitBefore, bool Unreachable, MDNode *BranchWeights = 0); + +/// SplitBlockAndInsertIfThenElse is similar to SplitBlockAndInsertIfThen, +/// but also creates the ElseBlock. +/// Before: +/// Head +/// SplitBefore +/// Tail +/// After: +/// Head +/// if (Cond) +/// ThenBlock +/// else +/// ElseBlock +/// SplitBefore +/// Tail +void SplitBlockAndInsertIfThenElse(Value *Cond, Instruction *SplitBefore, + TerminatorInst **ThenTerm, + TerminatorInst **ElseTerm, + MDNode *BranchWeights = 0); + /// /// GetIfCondition - Check whether BB is the merge point of a if-region. /// If so, return the boolean condition that determines which entry into diff --git a/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/lib/Transforms/Instrumentation/AddressSanitizer.cpp index 758669f1257..17c3c58e9d6 100644 --- a/lib/Transforms/Instrumentation/AddressSanitizer.cpp +++ b/lib/Transforms/Instrumentation/AddressSanitizer.cpp @@ -1508,27 +1508,31 @@ void FunctionStackPoisoner::poisonStack() { Value *ShadowBase = ASan.memToShadow(LocalStackBase, IRB); poisonRedZones(L.ShadowBytes, IRB, ShadowBase, true); - // Unpoison the stack before all ret instructions. + // (Un)poison the stack before all ret instructions. for (size_t i = 0, n = RetVec.size(); i < n; i++) { Instruction *Ret = RetVec[i]; IRBuilder<> IRBRet(Ret); // Mark the current frame as retired. IRBRet.CreateStore(ConstantInt::get(IntptrTy, kRetiredStackFrameMagic), BasePlus0); - // Unpoison the stack. - poisonRedZones(L.ShadowBytes, IRBRet, ShadowBase, false); if (DoStackMalloc) { assert(StackMallocIdx >= 0); - // In use-after-return mode, mark the whole stack frame unaddressable. + // if LocalStackBase != OrigStackBase: + // // In use-after-return mode, poison the whole stack frame. + // if StackMallocIdx <= 4 + // // For small sizes inline the whole thing: + // memset(ShadowBase, kAsanStackAfterReturnMagic, ShadowSize); + // **SavedFlagPtr(LocalStackBase) = 0 + // else + // __asan_stack_free_N(LocalStackBase, OrigStackBase) + // else + // + Value *Cmp = IRBRet.CreateICmpNE(LocalStackBase, OrigStackBase); + TerminatorInst *ThenTerm, *ElseTerm; + SplitBlockAndInsertIfThenElse(Cmp, Ret, &ThenTerm, &ElseTerm); + + IRBuilder<> IRBPoison(ThenTerm); if (StackMallocIdx <= 4) { - // For small sizes inline the whole thing: - // if LocalStackBase != OrigStackBase: - // memset(ShadowBase, kAsanStackAfterReturnMagic, ShadowSize); - // **SavedFlagPtr(LocalStackBase) = 0 - // FIXME: if LocalStackBase != OrigStackBase don't call poisonRedZones. - Value *Cmp = IRBRet.CreateICmpNE(LocalStackBase, OrigStackBase); - TerminatorInst *PoisonTerm = SplitBlockAndInsertIfThen(Cmp, Ret, false); - IRBuilder<> IRBPoison(PoisonTerm); int ClassSize = kMinStackMallocSize << StackMallocIdx; SetShadowToStackAfterReturnInlined(IRBPoison, ShadowBase, ClassSize >> Mapping.Scale); @@ -1542,15 +1546,20 @@ void FunctionStackPoisoner::poisonStack() { IRBPoison.CreateIntToPtr(SavedFlagPtr, IRBPoison.getInt8PtrTy())); } else { // For larger frames call __asan_stack_free_*. - IRBRet.CreateCall3(AsanStackFreeFunc[StackMallocIdx], LocalStackBase, - ConstantInt::get(IntptrTy, LocalStackSize), - OrigStackBase); + IRBPoison.CreateCall3(AsanStackFreeFunc[StackMallocIdx], LocalStackBase, + ConstantInt::get(IntptrTy, LocalStackSize), + OrigStackBase); } + + 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. assert(LocalStackBase == OrigStackBase); poisonAlloca(LocalStackBase, LocalStackSize, IRBRet, false); + } else { + poisonRedZones(L.ShadowBytes, IRBRet, ShadowBase, false); } } diff --git a/lib/Transforms/Utils/BasicBlockUtils.cpp b/lib/Transforms/Utils/BasicBlockUtils.cpp index 17bd115ee95..214a3aa538e 100644 --- a/lib/Transforms/Utils/BasicBlockUtils.cpp +++ b/lib/Transforms/Utils/BasicBlockUtils.cpp @@ -670,6 +670,39 @@ TerminatorInst *llvm::SplitBlockAndInsertIfThen(Value *Cond, return CheckTerm; } +/// SplitBlockAndInsertIfThenElse is similar to SplitBlockAndInsertIfThen, +/// but also creates the ElseBlock. +/// Before: +/// Head +/// SplitBefore +/// Tail +/// After: +/// Head +/// if (Cond) +/// ThenBlock +/// else +/// ElseBlock +/// SplitBefore +/// Tail +void llvm::SplitBlockAndInsertIfThenElse(Value *Cond, Instruction *SplitBefore, + TerminatorInst **ThenTerm, + TerminatorInst **ElseTerm, + MDNode *BranchWeights) { + BasicBlock *Head = SplitBefore->getParent(); + BasicBlock *Tail = Head->splitBasicBlock(SplitBefore); + TerminatorInst *HeadOldTerm = Head->getTerminator(); + LLVMContext &C = Head->getContext(); + BasicBlock *ThenBlock = BasicBlock::Create(C, "", Head->getParent(), Tail); + BasicBlock *ElseBlock = BasicBlock::Create(C, "", Head->getParent(), Tail); + *ThenTerm = BranchInst::Create(Tail, ThenBlock); + *ElseTerm = BranchInst::Create(Tail, ElseBlock); + BranchInst *HeadNewTerm = + BranchInst::Create(/*ifTrue*/ThenBlock, /*ifFalse*/ElseBlock, Cond); + HeadNewTerm->setMetadata(LLVMContext::MD_prof, BranchWeights); + ReplaceInstWithInst(HeadOldTerm, HeadNewTerm); +} + + /// GetIfCondition - Given a basic block (BB) with two predecessors, /// check to see if the merge at this block is due /// to an "if condition". If so, return the boolean condition that determines diff --git a/test/Instrumentation/AddressSanitizer/stack-poisoning.ll b/test/Instrumentation/AddressSanitizer/stack-poisoning.ll new file mode 100644 index 00000000000..2d69879925b --- /dev/null +++ b/test/Instrumentation/AddressSanitizer/stack-poisoning.ll @@ -0,0 +1,43 @@ +; RUN: opt < %s -asan -asan-use-after-return -S | FileCheck --check-prefix=CHECK-UAR %s +; RUN: opt < %s -asan -S | FileCheck --check-prefix=CHECK-PLAIN %s +target datalayout = "e-i64:64-f80:128-s:64-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare void @Foo(i8*) + +define void @Bar() uwtable sanitize_address { +entry: +; CHECK-PLAIN-LABEL: Bar +; CHECK-PLAIN-NOT: label +; CHECK-PLAIN: ret void + +; CHECK-UAR-LABEL: Bar +; CHECK-UAR: load i32* @__asan_option_detect_stack_use_after_return +; CHECK-UAR: label +; CHECK-UAR: call i64 @__asan_stack_malloc_1 +; CHECK-UAR: label +; CHECK-UAR: call void @Foo +; If LocalStackBase != OrigStackBase +; CHECK-UAR: label +; Then Block: poison the entire frame. + ; CHECK-UAR: store i64 -723401728380766731 + ; CHECK-UAR: store i64 -723401728380766731 + ; CHECK-UAR: store i8 0 + ; CHECK-UAR-NOT: store + ; CHECK-UAR: label +; Else Block: no UAR frame. Only unpoison the redzones. + ; CHECK-UAR: store i64 0 + ; CHECK-UAR: store i32 0 + ; CHECK-UAR-NOT: store + ; CHECK-UAR: label +; Done, no more stores. +; CHECK-UAR-NOT: store +; CHECK-UAR: ret void + + %x = alloca [20 x i8], align 16 + %arraydecay = getelementptr inbounds [20 x i8]* %x, i64 0, i64 0 + call void @Foo(i8* %arraydecay) + ret void +} + +