Remove unused variable.
[oota-llvm.git] / lib / CodeGen / StackProtector.cpp
index 659f8a01ab7fbc5ef386264ad9bd90fd6432a177..43a6ad8c97a476ca35cfceb604dbe1852d294b22 100644 (file)
@@ -16,6 +16,8 @@
 
 #define DEBUG_TYPE "stack-protector"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/Analysis/Dominators.h"
+#include "llvm/Attributes.h"
 #include "llvm/Constants.h"
 #include "llvm/DerivedTypes.h"
 #include "llvm/Function.h"
@@ -23,7 +25,6 @@
 #include "llvm/Intrinsics.h"
 #include "llvm/Module.h"
 #include "llvm/Pass.h"
-#include "llvm/ADT/APInt.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Target/TargetData.h"
 #include "llvm/Target/TargetLowering.h"
@@ -33,14 +34,11 @@ using namespace llvm;
 // smashing protection.
 static cl::opt<unsigned>
 SSPBufferSize("stack-protector-buffer-size", cl::init(8),
-              cl::desc("The lower bound for a buffer to be considered for "
-                       "stack smashing protection."));
+              cl::desc("Lower bound for a buffer to be considered for "
+                       "stack protection"));
 
 namespace {
-  class VISIBILITY_HIDDEN StackProtector : public FunctionPass {
-    /// Level - The level of stack protection.
-    SSP::StackProtectorLevel Level;
-
+  class StackProtector : public FunctionPass {
     /// TLI - Keep a pointer of a TargetLowering to consult for determining
     /// target type sizes.
     const TargetLowering *TLI;
@@ -48,6 +46,8 @@ namespace {
     Function *F;
     Module *M;
 
+    DominatorTree* DT;
+
     /// InsertStackProtectors - Insert code into the prologue and epilogue of
     /// the function.
     ///
@@ -65,32 +65,75 @@ namespace {
     bool RequiresStackProtector() const;
   public:
     static char ID;             // Pass identification, replacement for typeid.
-    StackProtector() : FunctionPass(&ID), Level(SSP::OFF), TLI(0) {}
-    StackProtector(SSP::StackProtectorLevel lvl, const TargetLowering *tli)
-      : FunctionPass(&ID), Level(lvl), TLI(tli) {}
+    StackProtector() : FunctionPass(ID), TLI(0) {
+      initializeStackProtectorPass(*PassRegistry::getPassRegistry());
+    }
+    StackProtector(const TargetLowering *tli)
+      : FunctionPass(ID), TLI(tli) {
+        initializeStackProtectorPass(*PassRegistry::getPassRegistry());
+      }
+
+    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
+      AU.addPreserved<DominatorTree>();
+    }
 
     virtual bool runOnFunction(Function &Fn);
   };
 } // end anonymous namespace
 
 char StackProtector::ID = 0;
-static RegisterPass<StackProtector>
-X("stack-protector", "Insert stack protectors");
+INITIALIZE_PASS(StackProtector, "stack-protector",
+                "Insert stack protectors", false, false)
 
-FunctionPass *llvm::createStackProtectorPass(SSP::StackProtectorLevel lvl,
-                                             const TargetLowering *tli) {
-  return new StackProtector(lvl, tli);
+FunctionPass *llvm::createStackProtectorPass(const TargetLowering *tli) {
+  return new StackProtector(tli);
 }
 
 bool StackProtector::runOnFunction(Function &Fn) {
   F = &Fn;
   M = F->getParent();
+  DT = getAnalysisIfAvailable<DominatorTree>();
 
   if (!RequiresStackProtector()) return false;
   
   return InsertStackProtectors();
 }
 
+/// RequiresStackProtector - Check whether or not this function needs a stack
+/// protector based upon the stack protector level. The heuristic we use is to
+/// add a guard variable to functions that call alloca, and functions with
+/// buffers larger than SSPBufferSize bytes.
+bool StackProtector::RequiresStackProtector() const {
+  if (F->hasFnAttr(Attribute::StackProtectReq))
+    return true;
+
+  if (!F->hasFnAttr(Attribute::StackProtect))
+    return false;
+
+  const TargetData *TD = TLI->getTargetData();
+
+  for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
+    BasicBlock *BB = I;
+
+    for (BasicBlock::iterator
+           II = BB->begin(), IE = BB->end(); II != IE; ++II)
+      if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
+        if (AI->isArrayAllocation())
+          // This is a call to alloca with a variable size. Emit stack
+          // protectors.
+          return true;
+
+        if (ArrayType *AT = dyn_cast<ArrayType>(AI->getAllocatedType()))
+          // If an array has more than SSPBufferSize bytes of allocated space,
+          // then we emit stack protectors.
+          if (SSPBufferSize <= TD->getTypeAllocSize(AT))
+            return true;
+      }
+  }
+
+  return false;
+}
+
 /// InsertStackProtectors - Insert code into the prologue and epilogue of the
 /// function.
 ///
@@ -98,121 +141,115 @@ bool StackProtector::runOnFunction(Function &Fn) {
 ///  - The epilogue checks the value stored in the prologue against the original
 ///    value. It calls __stack_chk_fail if they differ.
 bool StackProtector::InsertStackProtectors() {
-  std::vector<BasicBlock*> ReturnBBs;
-
-  for (Function::iterator I = F->begin(); I != F->end(); ++I)
-    if (isa<ReturnInst>(I->getTerminator()))
-      ReturnBBs.push_back(I);
-
-  // If this function doesn't return, don't bother with stack protectors.
-  if (ReturnBBs.empty()) return false;
-
-  // Insert code into the entry block that stores the __stack_chk_guard variable
-  // onto the stack.
-  BasicBlock &Entry = F->getEntryBlock();
-  Instruction *InsertPt = &Entry.front();
-
-  const PointerType *GuardTy = PointerType::getUnqual(Type::Int8Ty);
-
-  // The global variable for the stack guard.
-  Constant *StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", GuardTy);
-  LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsertPt);
-  CallInst::
-    Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_prologue),
-           LI, "", InsertPt);
-
-  // Create the basic block to jump to when the guard check fails.
-  BasicBlock *FailBB = CreateFailBB();
-
-  // Loop through the basic blocks that have return instructions. Convert this:
-  //
-  //   return:
-  //     ...
-  //     ret ...
-  //
-  // into this:
-  //
-  //   return:
-  //     ...
-  //     %1 = load __stack_chk_guard
-  //     %2 = load <stored stack guard>
-  //     %3 = cmp i1 %1, %2
-  //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
-  //
-  //   SP_return:
-  //     ret ...
-  //
-  //   CallStackCheckFailBlk:
-  //     call void @__stack_chk_fail()
-  //     unreachable
-  //
-  for (std::vector<BasicBlock*>::iterator
-         I = ReturnBBs.begin(), E = ReturnBBs.end(); I != E; ++I) {
-    BasicBlock *BB = *I;
-    ReturnInst *RI = cast<ReturnInst>(BB->getTerminator());
-    Function::iterator InsPt = BB; ++InsPt; // Insertion point for new BB.
+  BasicBlock *FailBB = 0;       // The basic block to jump to if check fails.
+  BasicBlock *FailBBDom = 0;    // FailBB's dominator.
+  AllocaInst *AI = 0;           // Place on stack that stores the stack guard.
+  Value *StackGuardVar = 0;  // The stack guard variable.
+
+  for (Function::iterator I = F->begin(), E = F->end(); I != E; ) {
+    BasicBlock *BB = I++;
+    ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator());
+    if (!RI) continue;
+
+    if (!FailBB) {
+      // Insert code into the entry block that stores the __stack_chk_guard
+      // variable onto the stack:
+      //
+      //   entry:
+      //     StackGuardSlot = alloca i8*
+      //     StackGuard = load __stack_chk_guard
+      //     call void @llvm.stackprotect.create(StackGuard, StackGuardSlot)
+      // 
+      PointerType *PtrTy = Type::getInt8PtrTy(RI->getContext());
+      unsigned AddressSpace, Offset;
+      if (TLI->getStackCookieLocation(AddressSpace, Offset)) {
+        Constant *OffsetVal =
+          ConstantInt::get(Type::getInt32Ty(RI->getContext()), Offset);
+        
+        StackGuardVar = ConstantExpr::getIntToPtr(OffsetVal,
+                                      PointerType::get(PtrTy, AddressSpace));
+      } else {
+        StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", PtrTy); 
+      }
+
+      BasicBlock &Entry = F->getEntryBlock();
+      Instruction *InsPt = &Entry.front();
+
+      AI = new AllocaInst(PtrTy, "StackGuardSlot", InsPt);
+      LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsPt);
+
+      Value *Args[] = { LI, AI };
+      CallInst::
+        Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector),
+               Args, "", InsPt);
+
+      // Create the basic block to jump to when the guard check fails.
+      FailBB = CreateFailBB();
+    }
+
+    // For each block with a return instruction, convert this:
+    //
+    //   return:
+    //     ...
+    //     ret ...
+    //
+    // into this:
+    //
+    //   return:
+    //     ...
+    //     %1 = load __stack_chk_guard
+    //     %2 = load StackGuardSlot
+    //     %3 = cmp i1 %1, %2
+    //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
+    //
+    //   SP_return:
+    //     ret ...
+    //
+    //   CallStackCheckFailBlk:
+    //     call void @__stack_chk_fail()
+    //     unreachable
 
     // Split the basic block before the return instruction.
     BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
 
+    if (DT && DT->isReachableFromEntry(BB)) {
+      DT->addNewBlock(NewBB, BB);
+      FailBBDom = FailBBDom ? DT->findNearestCommonDominator(FailBBDom, BB) :BB;
+    }
+
+    // Remove default branch instruction to the new BB.
+    BB->getTerminator()->eraseFromParent();
+
     // 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->removeFromParent();
-    F->getBasicBlockList().insert(InsPt, NewBB);
+    NewBB->moveAfter(BB);
 
     // Generate the stack protector instructions in the old basic block.
     LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
-    CallInst *CI = CallInst::
-      Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_epilogue),
-             "", BB);
-    ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, CI, LI1, "", BB);
+    LoadInst *LI2 = new LoadInst(AI, "", true, BB);
+    ICmpInst *Cmp = new ICmpInst(*BB, CmpInst::ICMP_EQ, LI1, LI2, "");
     BranchInst::Create(NewBB, FailBB, Cmp, BB);
   }
 
+  // Return if we didn't modify any basic blocks. I.e., there are no return
+  // statements in the function.
+  if (!FailBB) return false;
+
+  if (DT && FailBBDom)
+    DT->addNewBlock(FailBB, FailBBDom);
+
   return true;
 }
 
 /// CreateFailBB - Create a basic block to jump to when the stack protector
 /// check fails.
 BasicBlock *StackProtector::CreateFailBB() {
-  BasicBlock *FailBB = BasicBlock::Create("CallStackCheckFailBlk", F);
+  BasicBlock *FailBB = BasicBlock::Create(F->getContext(),
+                                          "CallStackCheckFailBlk", F);
   Constant *StackChkFail =
-    M->getOrInsertFunction("__stack_chk_fail", Type::VoidTy, NULL);
+    M->getOrInsertFunction("__stack_chk_fail",
+                           Type::getVoidTy(F->getContext()), NULL);
   CallInst::Create(StackChkFail, "", FailBB);
-  new UnreachableInst(FailBB);
+  new UnreachableInst(F->getContext(), FailBB);
   return FailBB;
 }
-
-/// RequiresStackProtector - Check whether or not this function needs a stack
-/// protector based upon the stack protector level.
-bool StackProtector::RequiresStackProtector() const {
-  switch (Level) {
-  default: return false;
-  case SSP::ALL: return true;
-  case SSP::SOME: {
-    // If the size of the local variables allocated on the stack is greater than
-    // SSPBufferSize, then we require a stack protector.
-    uint64_t StackSize = 0;
-    const TargetData *TD = TLI->getTargetData();
-
-    for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
-      BasicBlock *BB = I;
-
-      for (BasicBlock::iterator
-             II = BB->begin(), IE = BB->end(); II != IE; ++II)
-        if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
-          if (ConstantInt *CI = dyn_cast<ConstantInt>(AI->getArraySize())) {
-            const Type *Ty = AI->getAllocatedType();
-            uint64_t TySize = TD->getABITypeSize(Ty);
-            StackSize += TySize * CI->getZExtValue(); // Total allocated size.
-
-            if (SSPBufferSize <= StackSize)
-              return true;
-          }
-        }
-    }
-
-    return false;
-  }
-  }
-}