1 files changed, 85 insertions, 49 deletions

diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp
index 2a8234a..cf6ed09 100644
--- a/llvm/lib/CodeGen/StackProtector.cpp
+++ b/llvm/lib/CodeGen/StackProtector.cpp
@@ -49,7 +49,6 @@
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <optional>
-#include <utility>
 
 using namespace llvm;
 
@@ -70,13 +69,15 @@ static cl::opt<bool> DisableCheckNoReturn("disable-check-noreturn-call",
 ///  - The prologue code loads and stores the stack guard onto the stack.
 ///  - The epilogue checks the value stored in the prologue against the original
 ///    value. It calls __stack_chk_fail if they differ.
-static bool InsertStackProtectors(const TargetMachine *TM, Function *F,
-                                  DomTreeUpdater *DTU, bool &HasPrologue,
-                                  bool &HasIRCheck);
+static bool InsertStackProtectors(const TargetLowering &TLI,
+                                  const LibcallLoweringInfo &Libcalls,
+                                  Function *F, DomTreeUpdater *DTU,
+                                  bool &HasPrologue, bool &HasIRCheck);
 
 /// CreateFailBB - Create a basic block to jump to when the stack protector
 /// check fails.
-static BasicBlock *CreateFailBB(Function *F, const TargetLowering &TLI);
+static BasicBlock *CreateFailBB(Function *F,
+                                const LibcallLoweringInfo &Libcalls);
 
 bool SSPLayoutInfo::shouldEmitSDCheck(const BasicBlock &BB) const {
   return HasPrologue && !HasIRCheck && isa<ReturnInst>(BB.getTerminator());
@@ -132,8 +133,23 @@ PreservedAnalyses StackProtectorPass::run(Function &F,
       return PreservedAnalyses::all();
   }
 
+  auto &MAMProxy = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
+  const LibcallLoweringModuleAnalysisResult *LibcallLowering =
+      MAMProxy.getCachedResult<LibcallLoweringModuleAnalysis>(*F.getParent());
+
+  if (!LibcallLowering) {
+    F.getContext().emitError("'" + LibcallLoweringModuleAnalysis::name() +
+                             "' analysis required");
+    return PreservedAnalyses::all();
+  }
+
+  const TargetSubtargetInfo *STI = TM->getSubtargetImpl(F);
+  const TargetLowering *TLI = STI->getTargetLowering();
+  const LibcallLoweringInfo &Libcalls =
+      LibcallLowering->getLibcallLowering(*STI);
+
   ++NumFunProtected;
-  bool Changed = InsertStackProtectors(TM, &F, DT ? &DTU : nullptr,
+  bool Changed = InsertStackProtectors(*TLI, Libcalls, &F, DT ? &DTU : nullptr,
                                        Info.HasPrologue, Info.HasIRCheck);
 #ifdef EXPENSIVE_CHECKS
   assert((!DT ||
@@ -151,12 +167,11 @@ PreservedAnalyses StackProtectorPass::run(Function &F,
 
 char StackProtector::ID = 0;
 
-StackProtector::StackProtector() : FunctionPass(ID) {
-  initializeStackProtectorPass(*PassRegistry::getPassRegistry());
-}
+StackProtector::StackProtector() : FunctionPass(ID) {}
 
 INITIALIZE_PASS_BEGIN(StackProtector, DEBUG_TYPE,
                       "Insert stack protectors", false, true)
+INITIALIZE_PASS_DEPENDENCY(LibcallLoweringInfoWrapper)
 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_END(StackProtector, DEBUG_TYPE,
@@ -165,6 +180,7 @@ INITIALIZE_PASS_END(StackProtector, DEBUG_TYPE,
 FunctionPass *llvm::createStackProtectorPass() { return new StackProtector(); }
 
 void StackProtector::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addRequired<LibcallLoweringInfoWrapper>();
   AU.addRequired<TargetPassConfig>();
   AU.addPreserved<DominatorTreeWrapperPass>();
 }
@@ -191,9 +207,16 @@ bool StackProtector::runOnFunction(Function &Fn) {
       return false;
   }
 
+  const TargetSubtargetInfo *Subtarget = TM->getSubtargetImpl(Fn);
+  const LibcallLoweringInfo &Libcalls =
+      getAnalysis<LibcallLoweringInfoWrapper>().getLibcallLowering(*M,
+                                                                   *Subtarget);
+
+  const TargetLowering *TLI = Subtarget->getTargetLowering();
+
   ++NumFunProtected;
   bool Changed =
-      InsertStackProtectors(TM, F, DTU ? &*DTU : nullptr,
+      InsertStackProtectors(*TLI, Libcalls, F, DTU ? &*DTU : nullptr,
                             LayoutInfo.HasPrologue, LayoutInfo.HasIRCheck);
 #ifdef EXPENSIVE_CHECKS
   assert((!DTU ||
@@ -432,6 +455,11 @@ bool SSPLayoutAnalysis::requiresStackProtector(Function *F,
   for (const BasicBlock &BB : *F) {
     for (const Instruction &I : BB) {
       if (const AllocaInst *AI = dyn_cast<AllocaInst>(&I)) {
+        if (const MDNode *MD = AI->getMetadata("stack-protector")) {
+          const auto *CI = mdconst::dyn_extract<ConstantInt>(MD->getOperand(0));
+          if (CI->isZero())
+            continue;
+        }
         if (AI->isArrayAllocation()) {
           auto RemarkBuilder = [&]() {
             return OptimizationRemark(DEBUG_TYPE, "StackProtectorAllocaOrArray",
@@ -491,22 +519,23 @@ bool SSPLayoutAnalysis::requiresStackProtector(Function *F,
           continue;
         }
 
-        if (Strong &&
-            HasAddressTaken(
-                AI, M->getDataLayout().getTypeAllocSize(AI->getAllocatedType()),
-                M, VisitedPHIs)) {
-          ++NumAddrTaken;
-          if (!Layout)
-            return true;
-          Layout->insert(std::make_pair(AI, MachineFrameInfo::SSPLK_AddrOf));
-          ORE.emit([&]() {
-            return OptimizationRemark(DEBUG_TYPE, "StackProtectorAddressTaken",
-                                      &I)
-                   << "Stack protection applied to function "
-                   << ore::NV("Function", F)
-                   << " due to the address of a local variable being taken";
-          });
-          NeedsProtector = true;
+        if (Strong) {
+          std::optional<TypeSize> AllocSize =
+              AI->getAllocationSize(M->getDataLayout());
+          if (!AllocSize || HasAddressTaken(AI, *AllocSize, M, VisitedPHIs)) {
+            ++NumAddrTaken;
+            if (!Layout)
+              return true;
+            Layout->insert(std::make_pair(AI, MachineFrameInfo::SSPLK_AddrOf));
+            ORE.emit([&]() {
+              return OptimizationRemark(DEBUG_TYPE,
+                                        "StackProtectorAddressTaken", &I)
+                     << "Stack protection applied to function "
+                     << ore::NV("Function", F)
+                     << " due to the address of a local variable being taken";
+            });
+            NeedsProtector = true;
+          }
         }
         // Clear any PHIs that we visited, to make sure we examine all uses of
         // any subsequent allocas that we look at.
@@ -520,10 +549,11 @@ bool SSPLayoutAnalysis::requiresStackProtector(Function *F,
 
 /// Create a stack guard loading and populate whether SelectionDAG SSP is
 /// supported.
-static Value *getStackGuard(const TargetLoweringBase *TLI, Module *M,
+static Value *getStackGuard(const TargetLoweringBase &TLI,
+                            const LibcallLoweringInfo &Libcalls, Module *M,
                             IRBuilder<> &B,
                             bool *SupportsSelectionDAGSP = nullptr) {
-  Value *Guard = TLI->getIRStackGuard(B);
+  Value *Guard = TLI.getIRStackGuard(B, Libcalls);
   StringRef GuardMode = M->getStackProtectorGuard();
   if ((GuardMode == "tls" || GuardMode.empty()) && Guard)
     return B.CreateLoad(B.getPtrTy(), Guard, true, "StackGuard");
@@ -541,7 +571,7 @@ static Value *getStackGuard(const TargetLoweringBase *TLI, Module *M,
   // actually conveys the same information getIRStackGuard() already gives.
   if (SupportsSelectionDAGSP)
     *SupportsSelectionDAGSP = true;
-  TLI->insertSSPDeclarations(*M);
+  TLI.insertSSPDeclarations(*M, Libcalls);
   return B.CreateIntrinsic(Intrinsic::stackguard, {});
 }
 
@@ -556,29 +586,32 @@ static Value *getStackGuard(const TargetLoweringBase *TLI, Module *M,
 /// Returns true if the platform/triple supports the stackprotectorcreate pseudo
 /// node.
 static bool CreatePrologue(Function *F, Module *M, Instruction *CheckLoc,
-                           const TargetLoweringBase *TLI, AllocaInst *&AI) {
+                           const TargetLoweringBase *TLI,
+                           const LibcallLoweringInfo &Libcalls,
+                           AllocaInst *&AI) {
   bool SupportsSelectionDAGSP = false;
   IRBuilder<> B(&F->getEntryBlock().front());
   PointerType *PtrTy = PointerType::getUnqual(CheckLoc->getContext());
   AI = B.CreateAlloca(PtrTy, nullptr, "StackGuardSlot");
 
-  Value *GuardSlot = getStackGuard(TLI, M, B, &SupportsSelectionDAGSP);
+  Value *GuardSlot =
+      getStackGuard(*TLI, Libcalls, M, B, &SupportsSelectionDAGSP);
   B.CreateIntrinsic(Intrinsic::stackprotector, {GuardSlot, AI});
   return SupportsSelectionDAGSP;
 }
 
-bool InsertStackProtectors(const TargetMachine *TM, Function *F,
+bool InsertStackProtectors(const TargetLowering &TLI,
+                           const LibcallLoweringInfo &Libcalls, Function *F,
                            DomTreeUpdater *DTU, bool &HasPrologue,
                            bool &HasIRCheck) {
   auto *M = F->getParent();
-  auto *TLI = TM->getSubtargetImpl(*F)->getTargetLowering();
 
   // If the target wants to XOR the frame pointer into the guard value, it's
   // impossible to emit the check in IR, so the target *must* support stack
   // protection in SDAG.
   bool SupportsSelectionDAGSP =
-      TLI->useStackGuardXorFP() ||
-      (EnableSelectionDAGSP && !TM->Options.EnableFastISel);
+      TLI.useStackGuardXorFP() ||
+      (EnableSelectionDAGSP && !TLI.getTargetMachine().Options.EnableFastISel);
   AllocaInst *AI = nullptr; // Place on stack that stores the stack guard.
   BasicBlock *FailBB = nullptr;
 
@@ -611,7 +644,8 @@ bool InsertStackProtectors(const TargetMachine *TM, Function *F,
     // Generate prologue instrumentation if not already generated.
     if (!HasPrologue) {
       HasPrologue = true;
-      SupportsSelectionDAGSP &= CreatePrologue(F, M, CheckLoc, TLI, AI);
+      SupportsSelectionDAGSP &=
+          CreatePrologue(F, M, CheckLoc, &TLI, Libcalls, AI);
     }
 
     // SelectionDAG based code generation. Nothing else needs to be done here.
@@ -636,13 +670,13 @@ bool InsertStackProtectors(const TargetMachine *TM, Function *F,
     // inserted before the call rather than between it and the return.
     Instruction *Prev = CheckLoc->getPrevNode();
     if (auto *CI = dyn_cast_if_present<CallInst>(Prev))
-      if (CI->isTailCall() && isInTailCallPosition(*CI, *TM))
+      if (CI->isTailCall() && isInTailCallPosition(*CI, TLI.getTargetMachine()))
         CheckLoc = Prev;
 
     // Generate epilogue instrumentation. The epilogue intrumentation can be
     // function-based or inlined depending on which mechanism the target is
     // providing.
-    if (Function *GuardCheck = TLI->getSSPStackGuardCheck(*M)) {
+    if (Function *GuardCheck = TLI.getSSPStackGuardCheck(*M, Libcalls)) {
       // Generate the function-based epilogue instrumentation.
       // The target provides a guard check function, generate a call to it.
       IRBuilder<> B(CheckLoc);
@@ -681,10 +715,10 @@ bool InsertStackProtectors(const TargetMachine *TM, Function *F,
       // merge pass will merge together all of the various BB into one including
       // fail BB generated by the stack protector pseudo instruction.
       if (!FailBB)
-        FailBB = CreateFailBB(F, *TLI);
+        FailBB = CreateFailBB(F, Libcalls);
 
       IRBuilder<> B(CheckLoc);
-      Value *Guard = getStackGuard(TLI, M, B);
+      Value *Guard = getStackGuard(TLI, Libcalls, M, B);
       LoadInst *LI2 = B.CreateLoad(B.getPtrTy(), AI, true);
       auto *Cmp = cast<ICmpInst>(B.CreateICmpNE(Guard, LI2));
       auto SuccessProb =
@@ -714,7 +748,7 @@ bool InsertStackProtectors(const TargetMachine *TM, Function *F,
   return HasPrologue;
 }
 
-BasicBlock *CreateFailBB(Function *F, const TargetLowering &TLI) {
+BasicBlock *CreateFailBB(Function *F, const LibcallLoweringInfo &Libcalls) {
   auto *M = F->getParent();
   LLVMContext &Context = F->getContext();
   BasicBlock *FailBB = BasicBlock::Create(Context, "CallStackCheckFailBlk", F);
@@ -725,14 +759,16 @@ BasicBlock *CreateFailBB(Function *F, const TargetLowering &TLI) {
   FunctionCallee StackChkFail;
   SmallVector<Value *, 1> Args;
 
-  if (const char *ChkFailName =
-          TLI.getLibcallName(RTLIB::STACKPROTECTOR_CHECK_FAIL)) {
-    StackChkFail =
-        M->getOrInsertFunction(ChkFailName, Type::getVoidTy(Context));
-  } else if (const char *SSHName =
-                 TLI.getLibcallName(RTLIB::STACK_SMASH_HANDLER)) {
-    StackChkFail = M->getOrInsertFunction(SSHName, Type::getVoidTy(Context),
-                                          PointerType::getUnqual(Context));
+  if (RTLIB::LibcallImpl ChkFailImpl =
+          Libcalls.getLibcallImpl(RTLIB::STACKPROTECTOR_CHECK_FAIL)) {
+    StackChkFail = M->getOrInsertFunction(
+        RTLIB::RuntimeLibcallsInfo::getLibcallImplName(ChkFailImpl),
+        Type::getVoidTy(Context));
+  } else if (RTLIB::LibcallImpl SSHImpl =
+                 Libcalls.getLibcallImpl(RTLIB::STACK_SMASH_HANDLER)) {
+    StackChkFail = M->getOrInsertFunction(
+        RTLIB::RuntimeLibcallsInfo::getLibcallImplName(SSHImpl),
+        Type::getVoidTy(Context), PointerType::getUnqual(Context));
     Args.push_back(B.CreateGlobalString(F->getName(), "SSH"));
   } else {
     Context.emitError("no libcall available for stack protector");