lib/CodeGen/StackProtector.cpp

   1 //===-- StackProtector.cpp - Stack Protector Insertion --------------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This pass inserts stack protectors into functions which need them. A variable
  11 // with a random value in it is stored onto the stack before the local variables
  12 // are allocated. Upon exiting the block, the stored value is checked. If it's
  13 // changed, then there was some sort of violation and the program aborts.
  14 //
  15 //===----------------------------------------------------------------------===//
  16
  17 #define DEBUG_TYPE "stack-protector"
  18 #include "llvm/CodeGen/Passes.h"
  19 #include "llvm/Constants.h"
  20 #include "llvm/DerivedTypes.h"
  21 #include "llvm/Function.h"
  22 #include "llvm/Instructions.h"
  23 #include "llvm/Intrinsics.h"
  24 #include "llvm/Module.h"
  25 #include "llvm/Pass.h"
  26 #include "llvm/ADT/APInt.h"
  27 #include "llvm/Support/CommandLine.h"
  28 #include "llvm/Target/TargetData.h"
  29 #include "llvm/Target/TargetLowering.h"
  30 using namespace llvm;
  31
  32 // SSPBufferSize - The lower bound for a buffer to be considered for stack
  33 // smashing protection.
  34 static cl::opt<unsigned>
  35 SSPBufferSize("stack-protector-buffer-size", cl::init(8),
  36               cl::desc("The lower bound for a buffer to be considered for "
  37                        "stack smashing protection."));
  38
  39 namespace {
  40   class VISIBILITY_HIDDEN StackProtector : public FunctionPass {
  41     /// Level - The level of stack protection.
  42     SSP::StackProtectorLevel Level;
  43
  44     /// TLI - Keep a pointer of a TargetLowering to consult for determining
  45     /// target type sizes.
  46     const TargetLowering *TLI;
  47
  48     Function *F;
  49     Module *M;
  50
  51     /// InsertStackProtectors - Insert code into the prologue and epilogue of
  52     /// the function.
  53     ///
  54     ///  - The prologue code loads and stores the stack guard onto the stack.
  55     ///  - The epilogue checks the value stored in the prologue against the
  56     ///    original value. It calls __stack_chk_fail if they differ.
  57     bool InsertStackProtectors();
  58
  59     /// CreateFailBB - Create a basic block to jump to when the stack protector
  60     /// check fails.
  61     BasicBlock *CreateFailBB();
  62
  63     /// RequiresStackProtector - Check whether or not this function needs a
  64     /// stack protector based upon the stack protector level.
  65     bool RequiresStackProtector() const;
  66   public:
  67     static char ID;             // Pass identification, replacement for typeid.
  68     StackProtector() : FunctionPass(&ID), Level(SSP::OFF), TLI(0) {}
  69     StackProtector(SSP::StackProtectorLevel lvl, const TargetLowering *tli)
  70       : FunctionPass(&ID), Level(lvl), TLI(tli) {}
  71
  72     virtual bool runOnFunction(Function &Fn);
  73   };
  74 } // end anonymous namespace
  75
  76 char StackProtector::ID = 0;
  77 static RegisterPass<StackProtector>
  78 X("stack-protector", "Insert stack protectors");
  79
  80 FunctionPass *llvm::createStackProtectorPass(SSP::StackProtectorLevel lvl,
  81                                              const TargetLowering *tli) {
  82   return new StackProtector(lvl, tli);
  83 }
  84
  85 bool StackProtector::runOnFunction(Function &Fn) {
  86   F = &Fn;
  87   M = F->getParent();
  88
  89   if (!RequiresStackProtector()) return false;
  90
  91   return InsertStackProtectors();
  92 }
  93
  94 /// InsertStackProtectors - Insert code into the prologue and epilogue of the
  95 /// function.
  96 ///
  97 ///  - The prologue code loads and stores the stack guard onto the stack.
  98 ///  - The epilogue checks the value stored in the prologue against the original
  99 ///    value. It calls __stack_chk_fail if they differ.
 100 bool StackProtector::InsertStackProtectors() {
 101   std::vector<BasicBlock*> ReturnBBs;
 102
 103   for (Function::iterator I = F->begin(); I != F->end(); ++I)
 104     if (isa<ReturnInst>(I->getTerminator()))
 105       ReturnBBs.push_back(I);
 106
 107   // If this function doesn't return, don't bother with stack protectors.
 108   if (ReturnBBs.empty()) return false;
 109
 110   // Insert code into the entry block that stores the __stack_chk_guard variable
 111   // onto the stack.
 112   BasicBlock &Entry = F->getEntryBlock();
 113   Instruction *InsertPt = &Entry.front();
 114
 115   const PointerType *GuardTy = PointerType::getUnqual(Type::Int8Ty);
 116
 117   // The global variable for the stack guard.
 118   Constant *StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", GuardTy);
 119   LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsertPt);
 120   CallInst::
 121     Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_prologue),
 122            LI, "", InsertPt);
 123
 124   // Create the basic block to jump to when the guard check fails.
 125   BasicBlock *FailBB = CreateFailBB();
 126
 127   // Loop through the basic blocks that have return instructions. Convert this:
 128   //
 129   //   return:
 130   //     ...
 131   //     ret ...
 132   //
 133   // into this:
 134   //
 135   //   return:
 136   //     ...
 137   //     %1 = load __stack_chk_guard
 138   //     %2 = load <stored stack guard>
 139   //     %3 = cmp i1 %1, %2
 140   //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
 141   //
 142   //   SP_return:
 143   //     ret ...
 144   //
 145   //   CallStackCheckFailBlk:
 146   //     call void @__stack_chk_fail()
 147   //     unreachable
 148   //
 149   for (std::vector<BasicBlock*>::iterator
 150          I = ReturnBBs.begin(), E = ReturnBBs.end(); I != E; ++I) {
 151     BasicBlock *BB = *I;
 152     ReturnInst *RI = cast<ReturnInst>(BB->getTerminator());
 153     Function::iterator InsPt = BB; ++InsPt; // Insertion point for new BB.
 154
 155     // Split the basic block before the return instruction.
 156     BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
 157
 158     // Move the newly created basic block to the point right after the old basic
 159     // block so that it's in the "fall through" position.
 160     NewBB->removeFromParent();
 161     F->getBasicBlockList().insert(InsPt, NewBB);
 162
 163     // Generate the stack protector instructions in the old basic block.
 164     LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
 165     CallInst *CI = CallInst::
 166       Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_epilogue),
 167              "", BB);
 168     ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, CI, LI1, "", BB);
 169     BranchInst::Create(NewBB, FailBB, Cmp, BB);
 170   }
 171
 172   return true;
 173 }
 174
 175 /// CreateFailBB - Create a basic block to jump to when the stack protector
 176 /// check fails.
 177 BasicBlock *StackProtector::CreateFailBB() {
 178   BasicBlock *FailBB = BasicBlock::Create("CallStackCheckFailBlk", F);
 179   Constant *StackChkFail =
 180     M->getOrInsertFunction("__stack_chk_fail", Type::VoidTy, NULL);
 181   CallInst::Create(StackChkFail, "", FailBB);
 182   new UnreachableInst(FailBB);
 183   return FailBB;
 184 }
 185
 186 /// RequiresStackProtector - Check whether or not this function needs a stack
 187 /// protector based upon the stack protector level.
 188 bool StackProtector::RequiresStackProtector() const {
 189   switch (Level) {
 190   default: return false;
 191   case SSP::ALL: return true;
 192   case SSP::SOME: {
 193     // If the size of the local variables allocated on the stack is greater than
 194     // SSPBufferSize, then we require a stack protector.
 195     uint64_t StackSize = 0;
 196     const TargetData *TD = TLI->getTargetData();
 197
 198     for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
 199       BasicBlock *BB = I;
 200
 201       for (BasicBlock::iterator
 202              II = BB->begin(), IE = BB->end(); II != IE; ++II)
 203         if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
 204           if (ConstantInt *CI = dyn_cast<ConstantInt>(AI->getArraySize())) {
 205             const Type *Ty = AI->getAllocatedType();
 206             uint64_t TySize = TD->getABITypeSize(Ty);
 207             StackSize += TySize * CI->getZExtValue(); // Total allocated size.
 208
 209             if (SSPBufferSize <= StackSize)
 210               return true;
 211           }
 212         }
 213     }
 214
 215     return false;
 216   }
 217   }
 218 }