Index: lib/CodeGen/SafeStack.cpp =================================================================== --- lib/CodeGen/SafeStack.cpp +++ lib/CodeGen/SafeStack.cpp @@ -61,6 +61,11 @@ "Non-thread-local storage"), clEnumValEnd)); +static cl::opt USPInit("safe-stack-usp-init", + cl::Hidden, cl::init(false), + cl::desc("Use a single-threaded variable as the unsafe stack pointer " + "while the program is being initialized.")); + namespace llvm { STATISTIC(NumFunctions, "Total number of functions"); @@ -112,7 +117,13 @@ Type *Int32Ty; Type *Int8Ty; + /// Pointer to default unsafe stack pointer Value *UnsafeStackPtr = nullptr; + /// Pointer to single-threaded unsafe stack pointer to be used during program + /// initialization + GlobalVariable *UnsafeStackPtrInit = nullptr; + /// Load of init unsafe stack pointer at beginning of function + Value *UnsafeStackPtrInitLoad = nullptr; /// Unsafe stack alignment. Each stack frame must ensure that the stack is /// aligned to this value. We need to re-align the unsafe stack if the @@ -122,8 +133,26 @@ /// might expect to appear on the stack on most common targets. enum { StackAlignment = 16 }; + /// \brief Load the unsafe stack pointer. Assumes that + /// getOrCreateUnsafeStackPtr was invoked at the beginning of the function + /// being processed. + /// + /// \param Top Set to true if this load is at the beginning/top of the + /// function. + Instruction *loadUnsafeStackPtr(IRBuilder<> &IRB, bool Top = false); + /// \brief Build a value representing a pointer to the unsafe stack pointer. - Value *getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F); + /// If USPInit is set, also build a load instruction from the single-threaded + /// unsafe stack pointer to be used during program initialization and save it + /// in UnsafeStackPtrInitLoad. The loaded value is initially used to select + /// between that init pointer and the default unsafe stack pointer, depending + /// on whether the init pointer is non-null. That loaded value is also used + /// to determine where to later store the updated unsafe stack pointer. + Instruction *getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F); + + /// \brief Store the updated value of the unsafe stack pointer. + void storeUnsafeStackPtr(IRBuilder<> &IRB, Value *UpdatedUSP, + Value *NameSource = nullptr); /// \brief Return the value of the stack canary. Value *getStackGuard(IRBuilder<> &IRB, Function &F); @@ -171,8 +200,7 @@ /// \brief Replace all allocas in \p DynamicAllocas with code to allocate /// space dynamically on the unsafe stack and store the dynamic unsafe stack /// top to \p DynamicTop if non-null. - void moveDynamicAllocasToUnsafeStack(Function &F, Value *UnsafeStackPtr, - AllocaInst *DynamicTop, + void moveDynamicAllocasToUnsafeStack(Function &F, AllocaInst *DynamicTop, ArrayRef DynamicAllocas); bool IsSafeStackAlloca(const Value *AllocaPtr, uint64_t AllocaSize); @@ -359,18 +387,70 @@ return true; } -Value *SafeStack::getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F) { +Instruction *SafeStack::loadUnsafeStackPtr(IRBuilder<> &IRB, bool Top) { + if (!USPInit) + return IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr"); + + auto IsInit = IRB.CreateIsNotNull(UnsafeStackPtrInitLoad); + + auto PHI = IRB.CreatePHI(StackPtrTy, 2); + TerminatorInst *ThenTerm, *ElseTerm; + SplitBlockAndInsertIfThenElse(IsInit, PHI, &ThenTerm, &ElseTerm); + IRBuilder<> InitIRB(ThenTerm); + IRBuilder<> DefaultIRB(ElseTerm); + auto Tail = ThenTerm->getSuccessor(0); + IRB.SetInsertPoint(Tail, Tail->getFirstInsertionPt()); + + auto InitLoad = Top? + // There is no need to load this twice at the top of the function. + UnsafeStackPtrInitLoad : + InitIRB.CreateLoad(UnsafeStackPtrInit, false, "unsafe_stack_ptr_init"); + auto DefaultLoad = DefaultIRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr"); + + PHI->addIncoming(InitLoad, ThenTerm->getParent()); + PHI->addIncoming(DefaultLoad, ElseTerm->getParent()); + + return PHI; +} + +Instruction *SafeStack::getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F) { + Module &M = *F.getParent(); + + if (USPInit) { + const char *UnsafeStackPtrInitName = "__safestack_unsafe_stack_ptr_init"; + UnsafeStackPtrInit = + dyn_cast_or_null(M.getNamedValue(UnsafeStackPtrInitName)); + if (!UnsafeStackPtrInit) { + // The global variable is not defined yet, define it ourselves. + UnsafeStackPtrInit = new GlobalVariable( + M, StackPtrTy, false, GlobalValue::ExternalLinkage, nullptr, + UnsafeStackPtrInitName); + } else { + // The variable exists, check its type and attributes. + if (UnsafeStackPtrInit->getValueType() != StackPtrTy) + report_fatal_error(Twine(UnsafeStackPtrInitName) + + " must have void* type"); + if (UnsafeStackPtrInit->isThreadLocal()) + report_fatal_error(Twine(UnsafeStackPtrInitName) + + " must not be thread-local"); + } + + UnsafeStackPtrInitLoad = + IRB.CreateLoad(UnsafeStackPtrInit, false, "unsafe_stack_ptr_init"); + } + // Check if there is a target-specific location for the unsafe stack pointer. - if (TL) - if (Value *V = TL->getSafeStackPointerLocation(IRB)) - return V; + if (TL) { + UnsafeStackPtr = TL->getSafeStackPointerLocation(IRB); + if (UnsafeStackPtr) + return loadUnsafeStackPtr(IRB, true); + } // Otherwise, assume the target links with compiler-rt, which provides a // thread-local variable with a magic name. - Module &M = *F.getParent(); - const char *UnsafeStackPtrVar = "__safestack_unsafe_stack_ptr"; - auto UnsafeStackPtr = - dyn_cast_or_null(M.getNamedValue(UnsafeStackPtrVar)); + const char *UnsafeStackPtrName = "__safestack_unsafe_stack_ptr"; + UnsafeStackPtr = + dyn_cast_or_null(M.getNamedValue(UnsafeStackPtrName)); bool UseTLS = USPStorage == ThreadLocalUSP; @@ -383,16 +463,47 @@ // variable living anywhere other than in the main executable. UnsafeStackPtr = new GlobalVariable( M, StackPtrTy, false, GlobalValue::ExternalLinkage, nullptr, - UnsafeStackPtrVar, nullptr, TLSModel); + UnsafeStackPtrName, nullptr, TLSModel); } else { + auto UnsafeStackPtrVar = dyn_cast(UnsafeStackPtr); + // The variable exists, check its type and attributes. - if (UnsafeStackPtr->getValueType() != StackPtrTy) - report_fatal_error(Twine(UnsafeStackPtrVar) + " must have void* type"); - if (UseTLS != UnsafeStackPtr->isThreadLocal()) - report_fatal_error(Twine(UnsafeStackPtrVar) + " must " + + if (UnsafeStackPtrVar->getValueType() != StackPtrTy) + report_fatal_error(Twine(UnsafeStackPtrName) + " must have void* type"); + if (UseTLS != UnsafeStackPtrVar->isThreadLocal()) + report_fatal_error(Twine(UnsafeStackPtrName) + " must " + (UseTLS ? "" : "not ") + "be thread-local"); } - return UnsafeStackPtr; + + return loadUnsafeStackPtr(IRB, true); +} + +void SafeStack::storeUnsafeStackPtr(IRBuilder<> &IRB, Value *UpdatedUSP, + Value *NameSource) { + if (!USPInit) { + auto DefaultStore = IRB.CreateStore(UpdatedUSP, UnsafeStackPtr); + if (NameSource != nullptr) + DefaultStore->takeName(NameSource); + return; + } + + auto IsInit = IRB.CreateIsNotNull(UnsafeStackPtrInitLoad); + TerminatorInst *ThenTerm, *ElseTerm; + SplitBlockAndInsertIfThenElse(IsInit, (Instruction *)IRB.GetInsertPoint(), + &ThenTerm, &ElseTerm); + IRBuilder<> InitIRB(ThenTerm); + IRBuilder<> DefaultIRB(ElseTerm); + auto Tail = ThenTerm->getSuccessor(0); + IRB.SetInsertPoint(Tail, Tail->getFirstInsertionPt()); + + auto InitStore = InitIRB.CreateStore(UpdatedUSP, UnsafeStackPtrInit); + auto DefaultStore = DefaultIRB.CreateStore(UpdatedUSP, UnsafeStackPtr); + + if (NameSource == nullptr) + return; + + InitStore->takeName(NameSource); + DefaultStore->takeName(NameSource); } Value *SafeStack::getStackGuard(IRBuilder<> &IRB, Function &F) { @@ -484,7 +595,7 @@ IRB.SetInsertPoint(I->getNextNode()); Value *CurrentTop = DynamicTop ? IRB.CreateLoad(DynamicTop) : StaticTop; - IRB.CreateStore(CurrentTop, UnsafeStackPtr); + storeUnsafeStackPtr(IRB, CurrentTop); } return DynamicTop; @@ -634,12 +745,12 @@ Value *StaticTop = IRB.CreateGEP(BasePointer, ConstantInt::get(Int32Ty, -StaticOffset), "unsafe_stack_static_top"); - IRB.CreateStore(StaticTop, UnsafeStackPtr); + storeUnsafeStackPtr(IRB, StaticTop); return StaticTop; } void SafeStack::moveDynamicAllocasToUnsafeStack( - Function &F, Value *UnsafeStackPtr, AllocaInst *DynamicTop, + Function &F, AllocaInst *DynamicTop, ArrayRef DynamicAllocas) { DIBuilder DIB(*F.getParent()); @@ -655,7 +766,7 @@ uint64_t TySize = DL->getTypeAllocSize(Ty); Value *Size = IRB.CreateMul(ArraySize, ConstantInt::get(IntPtrTy, TySize)); - Value *SP = IRB.CreatePtrToInt(IRB.CreateLoad(UnsafeStackPtr), IntPtrTy); + Value *SP = IRB.CreatePtrToInt(loadUnsafeStackPtr(IRB), IntPtrTy); SP = IRB.CreateSub(SP, Size); // Align the SP value to satisfy the AllocaInst, type and stack alignments. @@ -669,7 +780,7 @@ StackPtrTy); // Save the stack pointer. - IRB.CreateStore(NewTop, UnsafeStackPtr); + storeUnsafeStackPtr(IRB, NewTop); if (DynamicTop) IRB.CreateStore(NewTop, DynamicTop); @@ -683,6 +794,7 @@ } if (!DynamicAllocas.empty()) { + std::vector StackSaves, StackRestores; // Now go through the instructions again, replacing stacksave/stackrestore. for (inst_iterator It = inst_begin(&F), Ie = inst_end(&F); It != Ie;) { Instruction *I = &*(It++); @@ -690,19 +802,25 @@ if (!II) continue; - if (II->getIntrinsicID() == Intrinsic::stacksave) { - IRBuilder<> IRB(II); - Instruction *LI = IRB.CreateLoad(UnsafeStackPtr); - LI->takeName(II); - II->replaceAllUsesWith(LI); - II->eraseFromParent(); - } else if (II->getIntrinsicID() == Intrinsic::stackrestore) { - IRBuilder<> IRB(II); - Instruction *SI = IRB.CreateStore(II->getArgOperand(0), UnsafeStackPtr); - SI->takeName(II); - assert(II->use_empty()); - II->eraseFromParent(); - } + if (II->getIntrinsicID() == Intrinsic::stacksave) + StackSaves.push_back(II); + else if (II->getIntrinsicID() == Intrinsic::stackrestore) + StackRestores.push_back(II); + } + + for (auto II : StackSaves) { + IRBuilder<> IRB(II); + Instruction *LI = loadUnsafeStackPtr(IRB); + LI->takeName(II); + II->replaceAllUsesWith(LI); + II->eraseFromParent(); + } + + for (auto II : StackRestores) { + IRBuilder<> IRB(II); + storeUnsafeStackPtr(IRB, II->getArgOperand(0), II); + assert(II->use_empty()); + II->eraseFromParent(); } } } @@ -756,12 +874,10 @@ ++NumUnsafeStackRestorePointsFunctions; IRBuilder<> IRB(&F.front(), F.begin()->getFirstInsertionPt()); - UnsafeStackPtr = getOrCreateUnsafeStackPtr(IRB, F); // 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"); + Instruction *BasePointer = getOrCreateUnsafeStackPtr(IRB, F); assert(BasePointer->getType() == StackPtrTy); AllocaInst *StackGuardSlot = nullptr; @@ -795,13 +911,12 @@ IRB, F, StackRestorePoints, StaticTop, !DynamicAllocas.empty()); // Handle dynamic allocas. - moveDynamicAllocasToUnsafeStack(F, UnsafeStackPtr, DynamicTop, - DynamicAllocas); + moveDynamicAllocasToUnsafeStack(F, DynamicTop, DynamicAllocas); // Restore the unsafe stack pointer before each return. for (ReturnInst *RI : Returns) { IRB.SetInsertPoint(RI); - IRB.CreateStore(BasePointer, UnsafeStackPtr); + storeUnsafeStackPtr(IRB, BasePointer); } DEBUG(dbgs() << "[SafeStack] safestack applied\n"); Index: test/Transforms/SafeStack/array.ll =================================================================== --- test/Transforms/SafeStack/array.ll +++ test/Transforms/SafeStack/array.ll @@ -1,12 +1,16 @@ ; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s ; RUN: opt -safe-stack -safe-stack-usp-storage=single-thread -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck -check-prefix=SINGLE-THREAD %s +; RUN: opt -safe-stack -safe-stack-usp-init -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck -check-prefix=INIT %s ; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s ; RUN: opt -safe-stack -safe-stack-usp-storage=single-thread -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck -check-prefix=SINGLE-THREAD %s +; RUN: opt -safe-stack -safe-stack-usp-init -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck -check-prefix=INIT %s ; array [4 x i8] ; Requires protector. ; CHECK: @__safestack_unsafe_stack_ptr = external thread_local(initialexec) global i8* +; INIT: @__safestack_unsafe_stack_ptr_init = external global i8* +; INIT: @__safestack_unsafe_stack_ptr = external thread_local(initialexec) global i8* ; SINGLE-THREAD: @__safestack_unsafe_stack_ptr = external global i8* define void @foo(i8* %a) nounwind uwtable safestack { @@ -58,6 +62,15 @@ ; CHECK-LABEL: define i8 @StaticArrayFixedUnsafe( ; CHECK: __safestack_unsafe_stack_ptr ; CHECK: ret i8 + + ; INIT-LABEL: define i8 @StaticArrayFixedUnsafe( + ; INIT: %unsafe_stack_ptr_init = load i8*, i8** @__safestack_unsafe_stack_ptr_init + ; INIT-NEXT: [[REG0:%[0-9]+]] = icmp ne i8* %unsafe_stack_ptr_init, null + ; INIT-NEXT: br i1 [[REG0]] + ; INIT: %unsafe_stack_ptr = load i8*, i8** @__safestack_unsafe_stack_ptr + ; INIT: store i8* [[REG1:.*]], i8** @__safestack_unsafe_stack_ptr_init + ; INIT: store i8* [[REG1]], i8** @__safestack_unsafe_stack_ptr + %buf = alloca i8, i32 4, align 1 %gep = getelementptr inbounds i8, i8* %buf, i32 5 %x = load i8, i8* %gep, align 1