Index: include/llvm/Target/TargetLowering.h =================================================================== --- include/llvm/Target/TargetLowering.h +++ include/llvm/Target/TargetLowering.h @@ -1079,6 +1079,14 @@ /// getIRStackGuard returns nullptr. virtual Value *getSDagStackGuard(const Module &M) const; + /// Return the function that's previously inserted by insertSSPDeclarations, + /// if any, otherwise return nullptr. The returned function needs to be + /// called to check the stack protection. + /// If nullptr, the stack validation is inlined before each return + /// instructions. + /// Should be used only when getIRStackGuard returns nullptr. + virtual Value *getSDagStackCheck(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 instrumentation code with one of this + /// instrumentation: + /// 1) with 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 InlineInstrumentation) { // 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 (InlineInstrumentation) { + 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,37 @@ 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.getSDagStackCheck(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; + } + // 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 +2085,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,11 @@ LowerArguments(Fn); } } - if (getAnalysis().shouldEmitSDCheck(*LLVMBB)) - SDB->SPDescriptor.initialize(LLVMBB, FuncInfo->MBBMap[LLVMBB]); + if (getAnalysis().shouldEmitSDCheck(*LLVMBB)) { + bool InlineInstrumentation = !TLI->getSDagStackCheck(*Fn.getParent()); + SDB->SPDescriptor.initialize(LLVMBB, FuncInfo->MBBMap[LLVMBB], + InlineInstrumentation); + } if (Begin != BI) ++NumDAGBlocks; @@ -1552,7 +1555,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 +1596,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 +1621,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 @@ -341,7 +341,16 @@ SupportsSelectionDAGSP &= CreatePrologue(F, M, RI, TLI, AI); } - if (!SupportsSelectionDAGSP) { + // SelectionDAG based code generation. Nothing else needs to be done here. + if (SupportsSelectionDAGSP) + break; + + if (Value* GuardCheck = TLI->getSDagStackCheck(*M)) { + // If the architecture provide a stack check function, call it. + IRBuilder<> B(RI); + LoadInst *Cookie = B.CreateLoad(AI, true); + B.CreateCall(GuardCheck, {Cookie}); + } else { // If we do not support SelectionDAG based tail calls, generate IR level // tail calls. // @@ -403,7 +412,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::getSDagStackCheck(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,13 @@ /// returns the address of that location. Otherwise, returns nullptr. Value *getIRStackGuard(IRBuilder<> &IRB) const override; + bool useLoadStackGuardNode() const override; + + // These methods need to be overriden to implement Stack-Protection + // with the SelectionDAGIsel mechanism. void insertSSPDeclarations(Module &M) const override; + Value *getSDagStackGuard(const Module &M) const override; + Value *getSDagStackCheck(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 +980,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::getSDagStackCheck(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::getSDagStackCheck(M); } Value *X86TargetLowering::getSafeStackPointerLocation(IRBuilder<> &IRB) const {