Index: include/llvm/Target/TargetLowering.h =================================================================== --- include/llvm/Target/TargetLowering.h +++ include/llvm/Target/TargetLowering.h @@ -1070,8 +1070,8 @@ /// LOAD_STACK_GUARD, or customize @llvm.stackguard(). virtual Value *getIRStackGuard(IRBuilder<> &IRB) const; - /// Inserts necessary declarations for SSP purpose. Should be used only when - /// getIRStackGuard returns nullptr. + /// Inserts necessary declarations for SSP (stack protection) purpose + /// Should be used only when getIRStackGuard returns nullptr. virtual void insertSSPDeclarations(Module &M) const; /// Return the variable that's previously inserted by insertSSPDeclarations, @@ -1079,6 +1079,12 @@ /// getIRStackGuard returns nullptr. virtual Value *getSDagStackGuard(const Module &M) const; + /// If the target has a standard stack protection check function that + /// performs validation and error handling, returns the function. Otherwise, + /// returns nullptr. Must be previously inserted by insertSSPDeclarations. + /// Should be used only when getIRStackGuard returns nullptr. + virtual Value *getSSPStackGuardCheck(const Module &M) const; + /// If the target has a standard location for the unsafe stack pointer, /// returns the address of that location. Otherwise, returns nullptr. virtual Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const; Index: lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h =================================================================== --- lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h +++ lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h @@ -454,7 +454,14 @@ /// /// c. After we finish selecting the basic block, in FinishBasicBlock if /// the StackProtectorDescriptor attached to the SelectionDAGBuilder is - /// initialized, we first find a splice point in the parent basic block + /// initialized, we produce the validation code with one of these + /// techniques: + /// 1) with a call to a stack check function + /// 2) with inlined instrumentation + /// + /// 1) We insert a call to the check function before the terminator. + /// + /// 2) We first find a splice point in the parent basic block /// before the terminator and then splice the terminator of said basic /// block into the success basic block. Then we code-gen a new tail for /// the parent basic block consisting of the two loads, the comparison, @@ -473,15 +480,22 @@ return ParentMBB && SuccessMBB && FailureMBB; } + bool shouldEmitFunctionCheckStackProtector() const { + return ParentMBB && !SuccessMBB && !FailureMBB; + } + /// Initialize the stack protector descriptor structure for a new basic /// block. - void initialize(const BasicBlock *BB, MachineBasicBlock *MBB) { + void initialize(const BasicBlock *BB, MachineBasicBlock *MBB, + bool FunctionBasedInstrumentation) { // Make sure we are not initialized yet. assert(!shouldEmitStackProtector() && "Stack Protector Descriptor is " "already initialized!"); ParentMBB = MBB; - SuccessMBB = AddSuccessorMBB(BB, MBB, /* IsLikely */ true); - FailureMBB = AddSuccessorMBB(BB, MBB, /* IsLikely */ false, FailureMBB); + if (!FunctionBasedInstrumentation) { + SuccessMBB = AddSuccessorMBB(BB, MBB, /* IsLikely */ true); + FailureMBB = AddSuccessorMBB(BB, MBB, /* IsLikely */ false, FailureMBB); + } } /// Reset state that changes when we handle different basic blocks. Index: lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp =================================================================== --- lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp +++ lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp @@ -2040,6 +2040,39 @@ const Module &M = *ParentBB->getParent()->getFunction()->getParent(); unsigned Align = DL->getPrefTypeAlignment(Type::getInt8PtrTy(M.getContext())); + SDValue StackSlot = DAG.getLoad( + PtrTy, dl, DAG.getEntryNode(), StackSlotPtr, + MachinePointerInfo::getFixedStack(DAG.getMachineFunction(), FI), true, + false, false, Align); + + // Retrieve guard check function, nullptr if instrumentation is inlined. + if (const Value *GuardCheck = TLI.getSSPStackGuardCheck(M)) { + auto *Fn = cast(GuardCheck); + FunctionType *FnTy = Fn->getFunctionType(); + assert(FnTy->getNumParams() == 1 && "Invalid function signature"); + + // Emit a call to the stack protection check function. + TargetLowering::ArgListTy Args; + TargetLowering::ArgListEntry Entry; + Entry.Node = StackSlot; + Entry.Ty = FnTy->getParamType(0); + if (Fn->getCallingConv() == CallingConv::X86_FastCall) + Entry.isInReg = true; + Args.push_back(Entry); + + TargetLowering::CallLoweringInfo CLI(DAG); + CLI.setDebugLoc(getCurSDLoc()) + .setChain(DAG.getEntryNode()) + .setCallee(Fn->getCallingConv(), FnTy->getReturnType(), + getValue(GuardCheck), std::move(Args)); + + std::pair Result = TLI.LowerCallTo(CLI); + DAG.setRoot(Result.second); + return; + } + + llvm::outs() << "SelectionDAG representation : inline instrumentation\n"; + // If useLoadStackGuardNode returns true, generate LOAD_STACK_GUARD. // Otherwise, emit a volatile load to retrieve the stack guard value. SDValue Chain = DAG.getEntryNode(); @@ -2054,11 +2087,6 @@ true, false, false, Align); } - SDValue StackSlot = DAG.getLoad( - PtrTy, dl, DAG.getEntryNode(), StackSlotPtr, - MachinePointerInfo::getFixedStack(DAG.getMachineFunction(), FI), true, - false, false, Align); - // Perform the comparison via a subtract/getsetcc. EVT VT = Guard.getValueType(); SDValue Sub = DAG.getNode(ISD::SUB, dl, VT, Guard, StackSlot); Index: lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp =================================================================== --- lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp +++ lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp @@ -1463,8 +1463,12 @@ LowerArguments(Fn); } } - if (getAnalysis().shouldEmitSDCheck(*LLVMBB)) - SDB->SPDescriptor.initialize(LLVMBB, FuncInfo->MBBMap[LLVMBB]); + if (getAnalysis().shouldEmitSDCheck(*LLVMBB)) { + bool FunctionBasedInstrumentation = + TLI->getSSPStackGuardCheck(*Fn.getParent()); + SDB->SPDescriptor.initialize(LLVMBB, FuncInfo->MBBMap[LLVMBB], + FunctionBasedInstrumentation); + } if (Begin != BI) ++NumDAGBlocks; @@ -1552,7 +1556,7 @@ /// terminator, but additionally the copies that move the vregs into the /// physical registers. static MachineBasicBlock::iterator -FindSplitPointForStackProtector(MachineBasicBlock *BB, DebugLoc DL) { +FindSplitPointForStackProtector(MachineBasicBlock *BB) { MachineBasicBlock::iterator SplitPoint = BB->getFirstTerminator(); // if (SplitPoint == BB->begin()) @@ -1593,7 +1597,21 @@ } // Handle stack protector. - if (SDB->SPDescriptor.shouldEmitStackProtector()) { + if (SDB->SPDescriptor.shouldEmitFunctionCheckStackProtector()) { + MachineBasicBlock *ParentMBB = SDB->SPDescriptor.getParentMBB(); + + // Add load and check to the basicblock. + FuncInfo->MBB = ParentMBB; + FuncInfo->InsertPt = + FindSplitPointForStackProtector(ParentMBB); + SDB->visitSPDescriptorParent(SDB->SPDescriptor, ParentMBB); + CurDAG->setRoot(SDB->getRoot()); + SDB->clear(); + CodeGenAndEmitDAG(); + + // Clear the Per-BB State. + SDB->SPDescriptor.resetPerBBState(); + } else if (SDB->SPDescriptor.shouldEmitStackProtector()) { MachineBasicBlock *ParentMBB = SDB->SPDescriptor.getParentMBB(); MachineBasicBlock *SuccessMBB = SDB->SPDescriptor.getSuccessMBB(); @@ -1604,7 +1622,7 @@ // register allocation issues caused by us splitting the parent mbb. The // register allocator will clean up said virtual copies later on. MachineBasicBlock::iterator SplitPoint = - FindSplitPointForStackProtector(ParentMBB, SDB->getCurDebugLoc()); + FindSplitPointForStackProtector(ParentMBB); // Splice the terminator of ParentMBB into SuccessMBB. SuccessMBB->splice(SuccessMBB->end(), ParentMBB, Index: lib/CodeGen/StackProtector.cpp =================================================================== --- lib/CodeGen/StackProtector.cpp +++ lib/CodeGen/StackProtector.cpp @@ -336,12 +336,33 @@ if (!RI) continue; + // Generate prologue instrumentationm if not already generated. if (!HasPrologue) { HasPrologue = true; SupportsSelectionDAGSP &= CreatePrologue(F, M, RI, TLI, AI); } - if (!SupportsSelectionDAGSP) { + // SelectionDAG based code generation. Nothing else needs to be done here. + // The prologue instrumentation is postponed to SelectionDAGISel. + if (SupportsSelectionDAGSP) { + break; + } + + // Set HasIRCheck to true, so that SelectionDAG will not generate its own + // version. + HasIRCheck = true; + + // Generate epilogue instrumentation. The epilogue intrumentation can be + // function-based or inlined depending on which mecanism the target is + // providing. + if (Value* GuardCheck = TLI->getSSPStackGuardCheck(*M)) { + // Generate the function-based epilogue instrumentation. + // The target provides a guard check function, generate a call to it. + IRBuilder<> B(RI); + LoadInst *Cookie = B.CreateLoad(AI, true); + B.CreateCall(GuardCheck, {Cookie}); + } else { + // Generate the inlined epilogue instrumentation. // If we do not support SelectionDAG based tail calls, generate IR level // tail calls. // @@ -372,10 +393,6 @@ // fail BB generated by the stack protector pseudo instruction. BasicBlock *FailBB = CreateFailBB(); - // Set HasIRCheck to true, so that SelectionDAG will not generate its own - // version. - HasIRCheck = true; - // Split the basic block before the return instruction. BasicBlock *NewBB = BB->splitBasicBlock(RI->getIterator(), "SP_return"); @@ -403,7 +420,7 @@ BranchProbabilityInfo::getBranchProbStackProtector(false); MDNode *Weights = MDBuilder(F->getContext()) .createBranchWeights(SuccessProb.getNumerator(), - FailureProb.getNumerator()); + FailureProb.getNumerator()); B.CreateCondBr(Cmp, NewBB, FailBB, Weights); } } Index: lib/CodeGen/TargetLoweringBase.cpp =================================================================== --- lib/CodeGen/TargetLoweringBase.cpp +++ lib/CodeGen/TargetLoweringBase.cpp @@ -898,7 +898,7 @@ setOperationAction(ISD::ZERO_EXTEND_VECTOR_INREG, VT, Expand); } - // For most targets @llvm.get.dynamic.area.offest just returns 0. + // For most targets @llvm.get.dynamic.area.offset just returns 0. setOperationAction(ISD::GET_DYNAMIC_AREA_OFFSET, VT, Expand); } @@ -1835,3 +1835,7 @@ Value *TargetLoweringBase::getSDagStackGuard(const Module &M) const { return M.getGlobalVariable("__stack_chk_guard"); } + +Value *TargetLoweringBase::getSSPStackGuardCheck(const Module &M) const { + return nullptr; +} Index: lib/Target/X86/X86ISelLowering.h =================================================================== --- lib/Target/X86/X86ISelLowering.h +++ lib/Target/X86/X86ISelLowering.h @@ -962,7 +962,10 @@ /// returns the address of that location. Otherwise, returns nullptr. Value *getIRStackGuard(IRBuilder<> &IRB) const override; + bool useLoadStackGuardNode() const override; void insertSSPDeclarations(Module &M) const override; + Value *getSDagStackGuard(const Module &M) const override; + Value *getSSPStackGuardCheck(const Module &M) const override; /// Return true if the target stores SafeStack pointer at a fixed offset in /// some non-standard address space, and populates the address space and @@ -974,7 +977,6 @@ bool isNoopAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const override; - bool useLoadStackGuardNode() const override; /// \brief Customize the preferred legalization strategy for certain types. LegalizeTypeAction getPreferredVectorAction(EVT VT) const override; Index: lib/Target/X86/X86ISelLowering.cpp =================================================================== --- lib/Target/X86/X86ISelLowering.cpp +++ lib/Target/X86/X86ISelLowering.cpp @@ -1952,8 +1952,36 @@ } void X86TargetLowering::insertSSPDeclarations(Module &M) const { - if (!Subtarget.isTargetGlibc()) - TargetLowering::insertSSPDeclarations(M); + // MSVC CRT has a global variable holding security cookie. + if (Subtarget.getTargetTriple().isOSMSVCRT()) { + M.getOrInsertGlobal("__security_cookie", + Type::getInt8PtrTy(M.getContext())); + + auto* SecurityCheckCookie = cast( + M.getOrInsertFunction("__security_check_cookie", + Type::getVoidTy(M.getContext()), + Type::getInt8PtrTy(M.getContext()), nullptr)); + SecurityCheckCookie->setCallingConv(CallingConv::X86_FastCall); + return; + } + // glibc has a special slot for the stack guard. + if (Subtarget.isTargetGlibc()) + return; + TargetLowering::insertSSPDeclarations(M); +} + +Value *X86TargetLowering::getSDagStackGuard(const Module &M) const { + // MSVC CRT has a global variable holding security cookie. + if (Subtarget.getTargetTriple().isOSMSVCRT()) + return M.getGlobalVariable("__security_cookie"); + return TargetLowering::getSDagStackGuard(M); +} + +Value *X86TargetLowering::getSSPStackGuardCheck(const Module &M) const { + // MSVC CRT provides a function to validate security cookie. + if (Subtarget.getTargetTriple().isOSMSVCRT()) + return M.getFunction("__security_check_cookie"); + return TargetLowering::getSSPStackGuardCheck(M); } Value *X86TargetLowering::getSafeStackPointerLocation(IRBuilder<> &IRB) const {