Enhance stack protector

Reviewed By: LuoYuanke

Differential Revision: https://reviews.llvm.org/D139254

1 files changed, 55 insertions, 14 deletions

diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp
index f974bb6..9a1063e 100644
--- a/llvm/lib/CodeGen/StackProtector.cpp
+++ b/llvm/lib/CodeGen/StackProtector.cpp
@@ -415,11 +415,11 @@ 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, ReturnInst *RI,
+static bool CreatePrologue(Function *F, Module *M, Instruction *CheckLoc,
                            const TargetLoweringBase *TLI, AllocaInst *&AI) {
   bool SupportsSelectionDAGSP = false;
   IRBuilder<> B(&F->getEntryBlock().front());
-  PointerType *PtrTy = Type::getInt8PtrTy(RI->getContext());
+  PointerType *PtrTy = Type::getInt8PtrTy(CheckLoc->getContext());
   AI = B.CreateAlloca(PtrTy, nullptr, "StackGuardSlot");
 
   Value *GuardSlot = getStackGuard(TLI, M, B, &SupportsSelectionDAGSP);
@@ -442,16 +442,34 @@ bool StackProtector::InsertStackProtectors() {
       TLI->useStackGuardXorFP() ||
       (EnableSelectionDAGSP && !TM->Options.EnableFastISel);
   AllocaInst *AI = nullptr; // Place on stack that stores the stack guard.
+  bool RecalculateDT = false;
+  BasicBlock *FailBB = nullptr;
 
   for (BasicBlock &BB : llvm::make_early_inc_range(*F)) {
-    ReturnInst *RI = dyn_cast<ReturnInst>(BB.getTerminator());
-    if (!RI)
+    // This is stack protector auto generated check BB, skip it.
+    if (&BB == FailBB)
+      continue;
+    Instruction *CheckLoc = dyn_cast<ReturnInst>(BB.getTerminator());
+    if (!CheckLoc) {
+      for (auto &Inst : BB) {
+        auto *CB = dyn_cast<CallBase>(&Inst);
+        if (!CB)
+          continue;
+        if (!CB->doesNotReturn())
+          continue;
+        // Do stack check before non-return calls (e.g: __cxa_throw)
+        CheckLoc = CB;
+        break;
+      }
+    }
+
+    if (!CheckLoc)
       continue;
 
     // Generate prologue instrumentation if not already generated.
     if (!HasPrologue) {
       HasPrologue = true;
-      SupportsSelectionDAGSP &= CreatePrologue(F, M, RI, TLI, AI);
+      SupportsSelectionDAGSP &= CreatePrologue(F, M, CheckLoc, TLI, AI);
     }
 
     // SelectionDAG based code generation. Nothing else needs to be done here.
@@ -477,8 +495,7 @@ bool StackProtector::InsertStackProtectors() {
     // verifier guarantees that a tail call is either directly before the
     // return or with a single correct bitcast of the return value in between so
     // we don't need to worry about many situations here.
-    Instruction *CheckLoc = RI;
-    Instruction *Prev = RI->getPrevNonDebugInstruction();
+    Instruction *Prev = CheckLoc->getPrevNonDebugInstruction();
     if (Prev && isa<CallInst>(Prev) && cast<CallInst>(Prev)->isTailCall())
       CheckLoc = Prev;
     else if (Prev) {
@@ -528,18 +545,13 @@ bool StackProtector::InsertStackProtectors() {
       // Create the FailBB. We duplicate the BB every time since the MI tail
       // merge pass will merge together all of the various BB into one including
       // fail BB generated by the stack protector pseudo instruction.
-      BasicBlock *FailBB = CreateFailBB();
+      if (!FailBB)
+        FailBB = CreateFailBB();
 
       // Split the basic block before the return instruction.
       BasicBlock *NewBB =
           BB.splitBasicBlock(CheckLoc->getIterator(), "SP_return");
 
-      // Update the dominator tree if we need to.
-      if (DT && DT->isReachableFromEntry(&BB)) {
-        DT->addNewBlock(NewBB, &BB);
-        DT->addNewBlock(FailBB, &BB);
-      }
-
       // Remove default branch instruction to the new BB.
       BB.getTerminator()->eraseFromParent();
 
@@ -560,9 +572,38 @@ bool StackProtector::InsertStackProtectors() {
                             .createBranchWeights(SuccessProb.getNumerator(),
                                                  FailureProb.getNumerator());
       B.CreateCondBr(Cmp, NewBB, FailBB, Weights);
+
+      // Update the dominator tree if we need to.
+      if (DT && DT->isReachableFromEntry(&BB))
+        RecalculateDT = true;
     }
   }
 
+  // TODO: Refine me, use faster way to update DT.
+  // Now we have spilt the BB, some like:
+  // ===================================
+  // BB:
+  //   RetOrNoReturnCall
+  // ==>
+  // BB:
+  //  CondBr
+  // NewBB:
+  //   RetOrNoReturnCall
+  // FailBB: (*)
+  //   HandleStackCheckFail
+  // ===================================
+  // The faster way should cover:
+  // For NewBB, it should success the old BB's dominatees.
+  // 1) return: it didn't have dominatee
+  // 2) no-return call: there may has dominatees.
+  //
+  // For FailBB, it may be created before, So
+  // 1) if it has 1 Predecessors, add it into DT.
+  // 2) if it has 2 Predecessors, it should has no dominator, remove it from DT.
+  // 3) if it has 3 or more Predecessors, DT has removed it, do nothing.
+  if (RecalculateDT)
+    DT->recalculate(*F);
+
   // Return if we didn't modify any basic blocks. i.e., there are no return
   // statements in the function.
   return HasPrologue;