[StackProtector] Don't crash when passing post-SP IR back in

[loladiro]

Summary

In julia, we use LLVM as a JIT compiler. A fairly common occurrence
is that we get an LLVM crash assertion in the backend and then extract
the problematic IR module from the core dump. However, this IR already
had the StackProtector pass run on it. Unfortunately, this means
that running llc on it will crash in ISel (because StackProtector
gets confused and assumes that the existence of the prologue means
that we will use SDag-based stack protection, but the necessary
intrinsic prototypes aren't actually defined). This is unfortunate,
because it also prevents using bugpoint on this IR without manually
stripping out the StackProtector artifacts first. This patch aims
to improve that situation by teaching StackProtector to recognize
whether a given prologue means that the SDag based stack protection
mechanism is available. If it is not, this patch does cause the
SP transformation to be applied twice, but that's a lot better
than crashing, because it often still allows the original bug
we're investigating to be reproduced, as well as allowing tools
like bugpoint to run properly.

I had to adjust one existing test. The reason for this is that
previously this test would trigger the SDAG-based stack-protector
despite not using llvm.stackguard. With the introduction of
the latter, this pattern is not actually generated by the
StackProtector pass anymore. For now I simply deleted the checks
for the precise generated instructions. Alternatively, it might
make sense to delete the test entirely, since the pattern it
tests for is no longer used.

[timshen]

To help myself recovering the memory better, I attempt to rephrase the problem:

• The IR pass StackProtector always generate the prologue, but not always generate the epilogue.
• Whether SDAG should generate the epilogue is indicated by StackProtector::shouldEmitSDCheck.
• During a re-run on llc, it doesn't matter whether we run StackProtector IR pass as part of the llc pipeline, the result is not expected. If we do run StackProtector, it inserts the prologue again; if we don't run StackProtector, StackProtector::shouldEmitSDCheck doesn't return the expected value, as the data members are not set.

If my memory recovey is successful and I understand the problem correctly, the best solution to me is to separate the mutation and analysis code in StackProtector. Suppose we split it into two passes "StackProtectorAnalysis" and "InsertStackProtectorPrologue", will it solve the problem? During the re-run, we only run StackProtectorAnalysis, but never again insert prologue. StackProtectorAnalysis should be kept as-is, without more special logic to detect whether the prologue is already generated (as this patch does). What do you think?

[loladiro]

If my memory recovey is successful and I understand the problem correctly, the best solution to me is to separate the mutation and analysis code in StackProtector. Suppose we split it into two passes "StackProtectorAnalysis" and "InsertStackProtectorPrologue", will it solve the problem? During the re-run, we only run StackProtectorAnalysis, but never again insert prologue. StackProtectorAnalysis should be kept as-is, without more special logic to detect whether the prologue is already generated (as this patch does). What do you think?

I think your description is accurate, with the small addition that if shouldEmitSDCheck is wrong (returns true when it shouldn't), the backend crashes because it expects the middle end to have inserted appropriate global declarations which are not there. I think it makes a lot of sense to separate the analysis and the mutation, but I'm not that's entirely sufficient without also giving the analysis the code to detect whether the code was already generated. Also, at the moment the mutation pass is implicitly added to the pass pipeline. Were you proposing that after the split it no longer is? If so, I think that would be quite breaking to users that expect it to be added implicitly. If not, I'm not sure it fixes the issue. Maybe I'm misunderstanding what you are proposing.

[timshen]

My feeling is that getSDagStackGuard/getSSPStackGuardCheck gets mismatched with shouldEmitSDCheck/SupportsSelectionDAGSP. Would you explain the invariance between them if you already investigated it? It seems shouldEmitSDCheck needs to be fixed too, but the current patch doesn't fix it.

Also, the function getSSPStackGuardCheck is weird, since we already have getSDagStackGuard.

include/llvm/CodeGen/StackProtector.h

private:
   /// if a variable's reference has been taken.  This set
   /// is maintained to ensure we don't visit the same PHI node multiple
   /// times.
   SmallPtrSet<const PHINode *, 16> VisitedPHIs;
 
   // A prologue is generated.
   bool HasPrologue = false;
 
+  // A stackguard call exists in the function.
+  bool HasStackGuard = false;
+
   // IR checking code is generated.
   bool HasIRCheck = false;
 
+  // Place on stack that stores the stack guard.
+  AllocaInst *StackGuardAlloca = nullptr;
+
   /// InsertStackProtectors - Insert code into the prologue and epilogue of
   /// the function.
   ///
   ///  - 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.
   bool InsertStackProtectors();
 

lib/CodeGen/StackProtector.cpp

bool StackProtector::runOnFunction(Function &Fn) {
   M = F->getParent();
   DominatorTreeWrapperPass *DTWP =
       getAnalysisIfAvailable<DominatorTreeWrapperPass>();
   DT = DTWP ? &DTWP->getDomTree() : nullptr;
   TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
   Trip = TM->getTargetTriple();
   TLI = TM->getSubtargetImpl(Fn)->getTargetLowering();
   HasPrologue = false;
+  HasStackGuard = false;
   HasIRCheck = false;
+  StackGuardAlloca = nullptr;
 
   Attribute Attr = Fn.getFnAttribute("stack-protector-buffer-size");
   if (Attr.isStringAttribute() &&
       Attr.getValueAsString().getAsInteger(10, SSPBufferSize))
     return false; // Invalid integer string
 
   if (!RequiresStackProtector())
     return false;

 /// strong heuristic will add a guard variables to functions that call alloca
 /// regardless of size, functions with any buffer regardless of type and size,
 /// functions with aggregates that contain any buffer regardless of type and
 /// size, and functions that contain stack-based variables that have had their
 /// address taken.
 bool StackProtector::RequiresStackProtector() {
   bool Strong = false;
   bool NeedsProtector = false;
-  for (const BasicBlock &BB : *F)
-    for (const Instruction &I : BB)
-      if (const CallInst *CI = dyn_cast<CallInst>(&I))
+  for (const BasicBlock &BB : *F) {
+    for (const Instruction &I : BB) {
+      if (const CallInst *CI = dyn_cast<CallInst>(&I)) {
         if (CI->getCalledFunction() ==
             Intrinsic::getDeclaration(F->getParent(),
-                                      Intrinsic::stackprotector))
+                                      Intrinsic::stackprotector)) {
+          StackGuardAlloca = cast<AllocaInst>(CI->getOperand(1));
           HasPrologue = true;
+        }
+        else if (CI->getCalledFunction() ==
+            Intrinsic::getDeclaration(F->getParent(),
+                                      Intrinsic::stackguard)) {
+          HasStackGuard = true;
+        }
+      }
+    }
+  }
 
   if (F->hasFnAttribute(Attribute::SafeStack))
     return false;
 
   // We are constructing the OptimizationRemarkEmitter on the fly rather than
   // using the analysis pass to avoid building DominatorTree and LoopInfo which
   // are not available this late in the IR pipeline.
   OptimizationRemarkEmitter ORE(F);

 ///  - 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.
 bool StackProtector::InsertStackProtectors() {
   // 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);
-  AllocaInst *AI = nullptr;       // Place on stack that stores the stack guard.
+      (TLI->useStackGuardXorFP() ||
+       (EnableSelectionDAGSP && !TM->Options.EnableFastISel)) &&
+      (!HasPrologue || HasStackGuard);
 
   for (Function::iterator I = F->begin(), E = F->end(); I != E;) {
     BasicBlock *BB = &*I++;
     ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator());
     if (!RI)
       continue;
 
     // Generate prologue instrumentation if not already generated.
     if (!HasPrologue) {
       HasPrologue = true;
-      SupportsSelectionDAGSP &= CreatePrologue(F, M, RI, TLI, AI);
+      SupportsSelectionDAGSP &= CreatePrologue(F, M, RI, TLI, StackGuardAlloca);
     }
 
     // SelectionDAG based code generation. Nothing else needs to be done here.
     // The epilogue instrumentation is postponed to SelectionDAG.
     if (SupportsSelectionDAGSP)
       break;
 
     // Set HasIRCheck to true, so that SelectionDAG will not generate its own
     // version. SelectionDAG called 'shouldEmitSDCheck' to check whether
     // instrumentation has already been generated.
     HasIRCheck = true;
 
     // Generate epilogue instrumentation. The epilogue intrumentation can be
     // function-based or inlined depending on which mechanism 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 *Guard = B.CreateLoad(AI, true, "Guard");
+      LoadInst *Guard = B.CreateLoad(StackGuardAlloca, true, "Guard");
       CallInst *Call = B.CreateCall(GuardCheck, {Guard});
       llvm::Function *Function = cast<llvm::Function>(GuardCheck);
       Call->setAttributes(Function->getAttributes());
       Call->setCallingConv(Function->getCallingConv());
     } else {
       // Generate the epilogue with inline instrumentation.
       // If we do not support SelectionDAG based tail calls, generate IR level
       // tail calls.
if (Value* GuardCheck = TLI->getSSPStackGuardCheck(*M)) {
 
       // Move the newly created basic block to the point right after the old
       // basic block so that it's in the "fall through" position.
       NewBB->moveAfter(BB);
 
       // Generate the stack protector instructions in the old basic block.
       IRBuilder<> B(BB);
       Value *Guard = getStackGuard(TLI, M, B);
-      LoadInst *LI2 = B.CreateLoad(AI, true);
+      LoadInst *LI2 = B.CreateLoad(StackGuardAlloca, true);
       Value *Cmp = B.CreateICmpEQ(Guard, LI2);
       auto SuccessProb =
           BranchProbabilityInfo::getBranchProbStackProtector(true);
       auto FailureProb =
           BranchProbabilityInfo::getBranchProbStackProtector(false);
       MDNode *Weights = MDBuilder(F->getContext())
                             .createBranchWeights(SuccessProb.getNumerator(),
                                                  FailureProb.getNumerator());

test/CodeGen/AArch64/stack-guard-remat-bitcast.ll

 ; RUN: llc < %s -mtriple=arm64-apple-ios -relocation-model=pic -disable-fp-elim | FileCheck %s
 
 @__stack_chk_guard = external global i64*
 
 ; PR20558
 
 ; CHECK: adrp [[R0:x[0-9]+]], ___stack_chk_guard@GOTPAGE
 ; CHECK: ldr  [[R1:x[0-9]+]], {{\[}}[[R0]], ___stack_chk_guard@GOTPAGEOFF{{\]}}
 ; Load the stack guard for the second time, just in case the previous value gets spilled.
 ; CHECK: adrp [[GUARD_PAGE:x[0-9]+]], ___stack_chk_guard@GOTPAGE
 ; CHECK: ldr  [[R2:x[0-9]+]], {{\[}}[[R1]]{{\]}}
-; CHECK: stur [[R2]], {{\[}}x29, [[SLOT0:[0-9#\-]+]]{{\]}}
-; CHECK: ldur [[R3:x[0-9]+]], {{\[}}x29, [[SLOT0]]{{\]}}
-; CHECK: ldr  [[GUARD_ADDR:x[0-9]+]], {{\[}}[[GUARD_PAGE]], ___stack_chk_guard@GOTPAGEOFF{{\]}}
-; CHECK: ldr  [[GUARD:x[0-9]+]], {{\[}}[[GUARD_ADDR]]{{\]}}
-; CHECK: cmp  [[GUARD]], [[R3]]
-; CHECK: b.ne LBB
 
 define i32 @test_stack_guard_remat2() {
 entry:
   %StackGuardSlot = alloca i8*
   %StackGuard = load i8*, i8** bitcast (i64** @__stack_chk_guard to i8**)
   call void @llvm.stackprotector(i8* %StackGuard, i8** %StackGuardSlot)
   %container = alloca [32 x i8], align 1
-  call void @llvm.stackprotectorcheck(i8** bitcast (i64** @__stack_chk_guard to i8**))
   ret i32 -1
 }
 
 declare void @llvm.stackprotector(i8*, i8**)
-declare void @llvm.stackprotectorcheck(i8**)

test/CodeGen/X86/stack-protector-twice.ll

+; RUN: llc -mtriple=x86_64-pc-linux-gnu < %s -o -
+
+declare void @__stack_chk_fail()
+declare void @llvm.stackprotector(i8*, i8**)
+
+; Check that we don't crash when running a function already instrumented by
+; StackProtector through llc
+define void @instrumented_func() sspstrong
+{
+top:
+  %StackGuardSlot = alloca i8*
+  %StackGuard = load volatile i8*, i8* addrspace(257)* inttoptr (i32 40 to i8* addrspace(257)*)
+  call void @llvm.stackprotector(i8* %StackGuard, i8** %StackGuardSlot)
+  %StackGuard2 = load volatile i8*, i8* addrspace(257)* inttoptr (i32 40 to i8* addrspace(257)*)
+  %loaded = load volatile i8*, i8** %StackGuardSlot
+  %cmp = icmp eq i8* %StackGuard2, %loaded
+  br i1 %cmp, label %SP_return, label %CallStackCheckFailBlk
+
+SP_return:
+  ret void
+
+CallStackCheckFailBlk:
+  call void @__stack_chk_fail()
+  unreachable
+}