Index: lib/CodeGen/SafeStack.cpp =================================================================== --- lib/CodeGen/SafeStack.cpp +++ lib/CodeGen/SafeStack.cpp @@ -17,6 +17,7 @@ #include "llvm/ADT/Statistic.h" #include "llvm/ADT/Triple.h" +#include "llvm/Analysis/BranchProbabilityInfo.h" #include "llvm/Analysis/ScalarEvolution.h" #include "llvm/Analysis/ScalarEvolutionExpressions.h" #include "llvm/CodeGen/Passes.h" @@ -31,6 +32,7 @@ #include "llvm/IR/Instructions.h" #include "llvm/IR/IntrinsicInst.h" #include "llvm/IR/Intrinsics.h" +#include "llvm/IR/MDBuilder.h" #include "llvm/IR/Module.h" #include "llvm/Pass.h" #include "llvm/Support/CommandLine.h" @@ -40,6 +42,7 @@ #include "llvm/Support/raw_os_ostream.h" #include "llvm/Target/TargetLowering.h" #include "llvm/Target/TargetSubtargetInfo.h" +#include "llvm/Transforms/Utils/BasicBlockUtils.h" #include "llvm/Transforms/Utils/Local.h" #include "llvm/Transforms/Utils/ModuleUtils.h" @@ -122,6 +125,13 @@ /// \brief Build a value representing a pointer to the unsafe stack pointer. Value *getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F); + /// \brief Return the value of the stack canary. + Value *getStackGuard(IRBuilder<> &IRB, Function &F); + + /// \brief Load stack guard from the frame and check if it has changed. + void checkStackGuard(IRBuilder<> &IRB, Function &F, ReturnInst &RI, + AllocaInst *StackGuardSlot, Value *StackGuard); + /// \brief Find all static allocas, dynamic allocas, return instructions and /// stack restore points (exception unwind blocks and setjmp calls) in the /// given function and append them to the respective vectors. @@ -145,7 +155,8 @@ ArrayRef StaticAllocas, ArrayRef ByValArguments, ArrayRef Returns, - Instruction *BasePointer); + Instruction *BasePointer, + AllocaInst *StackGuardSlot); /// \brief Generate code to restore the stack after all stack restore points /// in \p StackRestorePoints. @@ -379,6 +390,16 @@ return UnsafeStackPtr; } +Value *SafeStack::getStackGuard(IRBuilder<> &IRB, Function &F) { + Value *StackGuardVar = nullptr; + if (TL) + StackGuardVar = TL->getStackCookieLocation(IRB); + if (!StackGuardVar) + StackGuardVar = + F.getParent()->getOrInsertGlobal("__stack_chk_guard", StackPtrTy); + return IRB.CreateLoad(StackGuardVar, "StackGuard"); +} + void SafeStack::findInsts(Function &F, SmallVectorImpl &StaticAllocas, SmallVectorImpl &DynamicAllocas, @@ -464,13 +485,33 @@ return DynamicTop; } +void SafeStack::checkStackGuard(IRBuilder<> &IRB, Function &F, ReturnInst &RI, + AllocaInst *StackGuardSlot, Value *StackGuard) { + Value *V = IRB.CreateLoad(StackGuardSlot); + Value *Cmp = IRB.CreateICmpNE(StackGuard, V); + + auto SuccessProb = BranchProbabilityInfo::getBranchProbStackProtector(true); + auto FailureProb = BranchProbabilityInfo::getBranchProbStackProtector(false); + MDNode *Weights = MDBuilder(F.getContext()) + .createBranchWeights(SuccessProb.getNumerator(), + FailureProb.getNumerator()); + Instruction *CheckTerm = + SplitBlockAndInsertIfThen(Cmp, &RI, + /* Unreachable */ true, Weights); + IRBuilder<> IRBFail(CheckTerm); + // FIXME: respect -fsanitize-trap / -ftrap-function here? + Constant *StackChkFail = F.getParent()->getOrInsertFunction( + "__stack_chk_fail", IRB.getVoidTy(), nullptr); + IRBFail.CreateCall(StackChkFail, {}); +} + /// We explicitly compute and set the unsafe stack layout for all unsafe /// static alloca instructions. We save the unsafe "base pointer" in the /// prologue into a local variable and restore it in the epilogue. Value *SafeStack::moveStaticAllocasToUnsafeStack( IRBuilder<> &IRB, Function &F, ArrayRef StaticAllocas, ArrayRef ByValArguments, ArrayRef Returns, - Instruction *BasePointer) { + Instruction *BasePointer, AllocaInst *StackGuardSlot) { if (StaticAllocas.empty() && ByValArguments.empty()) return BasePointer; @@ -506,6 +547,18 @@ int64_t StaticOffset = 0; // Current stack top. IRB.SetInsertPoint(BasePointer->getNextNode()); + if (StackGuardSlot) { + StaticOffset += getStaticAllocaAllocationSize(StackGuardSlot); + Value *Off = IRB.CreateGEP(BasePointer, // BasePointer is i8* + ConstantInt::get(Int32Ty, -StaticOffset)); + Value *NewAI = + IRB.CreateBitCast(Off, StackGuardSlot->getType(), "StackGuardSlot"); + + // Replace alloc with the new location. + StackGuardSlot->replaceAllUsesWith(NewAI); + StackGuardSlot->eraseFromParent(); + } + for (Argument *Arg : ByValArguments) { Type *Ty = Arg->getType()->getPointerElementType(); @@ -667,18 +720,6 @@ TL = TM ? TM->getSubtargetImpl(F)->getTargetLowering() : nullptr; SE = &getAnalysis().getSE(); - { - // Make sure the regular stack protector won't run on this function - // (safestack attribute takes precedence). - AttrBuilder B; - B.addAttribute(Attribute::StackProtect) - .addAttribute(Attribute::StackProtectReq) - .addAttribute(Attribute::StackProtectStrong); - F.removeAttributes( - AttributeSet::FunctionIndex, - AttributeSet::get(F.getContext(), AttributeSet::FunctionIndex, B)); - } - ++NumFunctions; SmallVector StaticAllocas; @@ -715,13 +756,29 @@ // Load the current stack pointer (we'll also use it as a base pointer). // FIXME: use a dedicated register for it ? Instruction *BasePointer = - IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr"); + IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr"); assert(BasePointer->getType() == StackPtrTy); - // The top of the unsafe stack after all unsafe static allocas are allocated. - Value *StaticTop = moveStaticAllocasToUnsafeStack(IRB, F, StaticAllocas, - ByValArguments, Returns, - BasePointer); + AllocaInst *StackGuardSlot = nullptr; + // FIXME: implement weaker forms of stack protector. + if (F.hasFnAttribute(Attribute::StackProtect) || + F.hasFnAttribute(Attribute::StackProtectStrong) || + F.hasFnAttribute(Attribute::StackProtectReq)) { + Value *StackGuard = getStackGuard(IRB, F); + StackGuardSlot = IRB.CreateAlloca(StackPtrTy, nullptr); + IRB.CreateStore(StackGuard, StackGuardSlot); + + for (ReturnInst *RI : Returns) { + IRBuilder<> IRBRet(RI); + checkStackGuard(IRBRet, F, *RI, StackGuardSlot, StackGuard); + } + } + + // The top of the unsafe stack after all unsafe static allocas are + // allocated. + Value *StaticTop = + moveStaticAllocasToUnsafeStack(IRB, F, StaticAllocas, ByValArguments, + Returns, BasePointer, StackGuardSlot); // Safe stack object that stores the current unsafe stack top. It is updated // as unsafe dynamic (non-constant-sized) allocas are allocated and freed. Index: lib/CodeGen/StackProtector.cpp =================================================================== --- lib/CodeGen/StackProtector.cpp +++ lib/CodeGen/StackProtector.cpp @@ -200,6 +200,9 @@ bool StackProtector::RequiresStackProtector() { bool Strong = false; bool NeedsProtector = false; + if (F->hasFnAttribute(Attribute::SafeStack)) + return false; + if (F->hasFnAttribute(Attribute::StackProtectReq)) { NeedsProtector = true; Strong = true; // Use the same heuristic as strong to determine SSPLayout Index: lib/IR/Attributes.cpp =================================================================== --- lib/IR/Attributes.cpp +++ lib/IR/Attributes.cpp @@ -1477,20 +1477,14 @@ AttributeSet::FunctionIndex, B); - if (Callee.hasFnAttribute(Attribute::SafeStack)) { - Caller.removeAttributes(AttributeSet::FunctionIndex, OldSSPAttr); - Caller.addFnAttr(Attribute::SafeStack); - } else if (Callee.hasFnAttribute(Attribute::StackProtectReq) && - !Caller.hasFnAttribute(Attribute::SafeStack)) { + if (Callee.hasFnAttribute(Attribute::StackProtectReq)) { Caller.removeAttributes(AttributeSet::FunctionIndex, OldSSPAttr); Caller.addFnAttr(Attribute::StackProtectReq); } else if (Callee.hasFnAttribute(Attribute::StackProtectStrong) && - !Caller.hasFnAttribute(Attribute::SafeStack) && !Caller.hasFnAttribute(Attribute::StackProtectReq)) { Caller.removeAttributes(AttributeSet::FunctionIndex, OldSSPAttr); Caller.addFnAttr(Attribute::StackProtectStrong); } else if (Callee.hasFnAttribute(Attribute::StackProtect) && - !Caller.hasFnAttribute(Attribute::SafeStack) && !Caller.hasFnAttribute(Attribute::StackProtectReq) && !Caller.hasFnAttribute(Attribute::StackProtectStrong)) Caller.addFnAttr(Attribute::StackProtect); Index: test/CodeGen/X86/safestack_ssp.ll =================================================================== --- /dev/null +++ test/CodeGen/X86/safestack_ssp.ll @@ -0,0 +1,26 @@ +; RUN: llc -mtriple=i386-linux < %s -o - | FileCheck --check-prefix=LINUX-I386 %s +; RUN: llc -mtriple=x86_64-linux < %s -o - | FileCheck --check-prefix=LINUX-X64 %s + +define void @_Z1fv() safestack sspreq { +entry: + %x = alloca i32, align 4 + %0 = bitcast i32* %x to i8* + call void @_Z7CapturePi(i32* nonnull %x) + ret void +} + +declare void @_Z7CapturePi(i32*) + +; LINUX-X64-DAG: movq __safestack_unsafe_stack_ptr@GOTTPOFF(%rip), %[[A:.*]] +; LINUX-X64-DAG: movq %fs:(%[[A]]), %[[B:.*]] +; LINUX-X64-DAG: movq %fs:40, %[[COOKIE:.*]] +; LINUX-X64-DAG: leaq -16(%[[B]]), %[[C:.*]] +; LINUX-X64-DAG: movq %[[C]], %fs:(%[[A]]) +; LINUX-X64-DAG: movq %[[COOKIE]], -8(%[[B]]) + +; LINUX-I386-DAG: movl __safestack_unsafe_stack_ptr@INDNTPOFF, %[[A:.*]] +; LINUX-I386-DAG: movl %gs:(%[[A]]), %[[B:.*]] +; LINUX-I386-DAG: movl %gs:20, %[[COOKIE:.*]] +; LINUX-I386-DAG: leal -16(%[[B]]), %[[C:.*]] +; LINUX-I386-DAG: movl %[[C]], %gs:(%[[A]]) +; LINUX-I386-DAG: movl %[[COOKIE]], -4(%[[B]]) Index: test/Transforms/SafeStack/AArch64/abi_ssp.ll =================================================================== --- /dev/null +++ test/Transforms/SafeStack/AArch64/abi_ssp.ll @@ -0,0 +1,22 @@ +; RUN: opt -safe-stack -S -mtriple=aarch64-linux-android < %s -o - | FileCheck --check-prefix=TLS %s + + +define void @foo() nounwind uwtable safestack sspreq { +entry: +; The first @llvm.aarch64.thread.pointer is for the unsafe stack pointer, skip it. +; TLS: call i8* @llvm.aarch64.thread.pointer() + +; TLS: %[[TP2:.*]] = call i8* @llvm.aarch64.thread.pointer() +; TLS: %[[B:.*]] = getelementptr i8, i8* %[[TP2]], i32 40 +; TLS: %[[C:.*]] = bitcast i8* %[[B]] to i8** +; TLS: %[[StackGuard:.*]] = load i8*, i8** %[[C]] +; TLS: store i8* %[[StackGuard]], i8** %[[StackGuardSlot:.*]] + %a = alloca i128, align 16 + call void @Capture(i128* %a) + +; TLS: %[[A:.*]] = load i8*, i8** %[[StackGuardSlot]] +; TLS: icmp ne i8* %[[StackGuard]], %[[A]] + ret void +} + +declare void @Capture(i128*) Index: test/Transforms/SafeStack/X86/abi_ssp.ll =================================================================== --- /dev/null +++ test/Transforms/SafeStack/X86/abi_ssp.ll @@ -0,0 +1,19 @@ +; RUN: opt -safe-stack -S -mtriple=i686-pc-linux-gnu < %s -o - | FileCheck --check-prefix=TLS --check-prefix=TLS32 %s +; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck --check-prefix=TLS --check-prefix=TLS64 %s +; RUN: opt -safe-stack -S -mtriple=i686-linux-android < %s -o - | FileCheck --check-prefix=TLS --check-prefix=TLS32 %s +; RUN: opt -safe-stack -S -mtriple=x86_64-linux-android < %s -o - | FileCheck --check-prefix=TLS --check-prefix=TLS64 %s + +define void @foo() safestack sspreq { +entry: +; TLS32: %[[StackGuard:.*]] = load i8*, i8* addrspace(256)* inttoptr (i32 20 to i8* addrspace(256)*) +; TLS64: %[[StackGuard:.*]] = load i8*, i8* addrspace(257)* inttoptr (i32 40 to i8* addrspace(257)*) +; TLS: store i8* %[[StackGuard]], i8** %[[StackGuardSlot:.*]] + %a = alloca i8, align 1 + call void @Capture(i8* %a) + +; TLS: %[[A:.*]] = load i8*, i8** %[[StackGuardSlot]] +; TLS: icmp ne i8* %[[StackGuard]], %[[A]] + ret void +} + +declare void @Capture(i8*) Index: test/Transforms/SafeStack/X86/ssp.ll =================================================================== --- /dev/null +++ test/Transforms/SafeStack/X86/ssp.ll @@ -0,0 +1,30 @@ +; RUN: opt -safe-stack -S -mtriple=x86_64-unknown < %s -o - | FileCheck %s + +define void @foo() safestack sspreq { +entry: +; CHECK: %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr +; CHECK: %[[USST:.*]] = getelementptr i8, i8* %[[USP]], i32 -16 +; CHECK: store i8* %[[USST]], i8** @__safestack_unsafe_stack_ptr + +; CHECK: %[[A:.*]] = getelementptr i8, i8* %[[USP]], i32 -8 +; CHECK: %[[StackGuardSlot:.*]] = bitcast i8* %[[A]] to i8** +; CHECK: %[[StackGuard:.*]] = load i8*, i8** @__stack_chk_guard +; CHECK: store i8* %[[StackGuard]], i8** %[[StackGuardSlot]] + %a = alloca i8, align 1 + +; CHECK: call void @Capture + call void @Capture(i8* %a) + +; CHECK: %[[B:.*]] = load i8*, i8** %[[StackGuardSlot]] +; CHECK: %[[COND:.*]] = icmp ne i8* %[[StackGuard]], %[[B]] +; CHECK: br i1 %[[COND]], {{.*}} !prof + +; CHECK: call void @__stack_chk_fail() +; CHECK-NEXT: unreachable + +; CHECK: store i8* %[[USP]], i8** @__safestack_unsafe_stack_ptr +; CHECK-NEXT: ret void + ret void +} + +declare void @Capture(i8*)