3f5e9802e724b6f2b73b555ea47d6348ff305562
[oota-llvm.git] / 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/Support/CommandLine.h"
27 #include "llvm/Target/TargetData.h"
28 #include "llvm/Target/TargetLowering.h"
29 using namespace llvm;
30
31 // SSPBufferSize - The lower bound for a buffer to be considered for stack
32 // smashing protection.
33 static cl::opt<unsigned>
34 SSPBufferSize("stack-protector-buffer-size", cl::init(8),
35               cl::desc("The lower bound for a buffer to be considered for "
36                        "stack smashing protection."));
37
38 namespace {
39   class VISIBILITY_HIDDEN StackProtector : public FunctionPass {
40     /// Level - The level of stack protection.
41     SSP::StackProtectorLevel Level;
42
43     /// TLI - Keep a pointer of a TargetLowering to consult for determining
44     /// target type sizes.
45     const TargetLowering *TLI;
46
47     Function *F;
48     Module *M;
49
50     /// InsertStackProtectors - Insert code into the prologue and epilogue of
51     /// the function.
52     ///
53     ///  - The prologue code loads and stores the stack guard onto the stack.
54     ///  - The epilogue checks the value stored in the prologue against the
55     ///    original value. It calls __stack_chk_fail if they differ.
56     bool InsertStackProtectors();
57
58     /// CreateFailBB - Create a basic block to jump to when the stack protector
59     /// check fails.
60     BasicBlock *CreateFailBB();
61
62     /// RequiresStackProtector - Check whether or not this function needs a
63     /// stack protector based upon the stack protector level.
64     bool RequiresStackProtector() const;
65   public:
66     static char ID;             // Pass identification, replacement for typeid.
67     StackProtector() : FunctionPass(&ID), Level(SSP::OFF), TLI(0) {}
68     StackProtector(SSP::StackProtectorLevel lvl, const TargetLowering *tli)
69       : FunctionPass(&ID), Level(lvl), TLI(tli) {}
70
71     virtual bool runOnFunction(Function &Fn);
72   };
73 } // end anonymous namespace
74
75 char StackProtector::ID = 0;
76 static RegisterPass<StackProtector>
77 X("stack-protector", "Insert stack protectors");
78
79 FunctionPass *llvm::createStackProtectorPass(SSP::StackProtectorLevel lvl,
80                                              const TargetLowering *tli) {
81   return new StackProtector(lvl, tli);
82 }
83
84 bool StackProtector::runOnFunction(Function &Fn) {
85   F = &Fn;
86   M = F->getParent();
87
88   if (!RequiresStackProtector()) return false;
89   
90   return InsertStackProtectors();
91 }
92
93 /// InsertStackProtectors - Insert code into the prologue and epilogue of the
94 /// function.
95 ///
96 ///  - The prologue code loads and stores the stack guard onto the stack.
97 ///  - The epilogue checks the value stored in the prologue against the original
98 ///    value. It calls __stack_chk_fail if they differ.
99 bool StackProtector::InsertStackProtectors() {
100   Constant *StackGuardVar = 0;  // The global variable for the stack guard.
101   BasicBlock *FailBB = 0;       // The basic block to jump to if check fails.
102
103   // Loop through the basic blocks that have return instructions. Convert this:
104   //
105   //   return:
106   //     ...
107   //     ret ...
108   //
109   // into this:
110   //
111   //   return:
112   //     ...
113   //     %1 = load __stack_chk_guard
114   //     %2 = load <stored stack guard>
115   //     %3 = cmp i1 %1, %2
116   //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
117   //
118   //   SP_return:
119   //     ret ...
120   //
121   //   CallStackCheckFailBlk:
122   //     call void @__stack_chk_fail()
123   //     unreachable
124   //
125   for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
126     BasicBlock *BB = I;
127
128     if (isa<ReturnInst>(BB->getTerminator())) {
129       // Create the basic block to jump to when the guard check fails.
130       if (!FailBB)
131         FailBB = CreateFailBB();
132
133       if (!StackGuardVar)
134         StackGuardVar =
135           M->getOrInsertGlobal("__stack_chk_guard",
136                                PointerType::getUnqual(Type::Int8Ty));
137
138       ReturnInst *RI = cast<ReturnInst>(BB->getTerminator());
139       Function::iterator InsPt = BB; ++InsPt; // Insertion point for new BB.
140       ++I;
141
142       // Split the basic block before the return instruction.
143       BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
144
145       // Move the newly created basic block to the point right after the old basic
146       // block so that it's in the "fall through" position.
147       NewBB->removeFromParent();
148       F->getBasicBlockList().insert(InsPt, NewBB);
149
150       // Generate the stack protector instructions in the old basic block.
151       LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
152       CallInst *CI = CallInst::
153         Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_check),
154                "", BB);
155       ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, CI, LI1, "", BB);
156       BranchInst::Create(NewBB, FailBB, Cmp, BB);
157     }
158   }
159
160   // Return if we didn't modify any basic blocks. I.e., there are no return
161   // statements in the function.
162   if (!FailBB) return false;
163
164   // Insert code into the entry block that stores the __stack_chk_guard variable
165   // onto the stack.
166   BasicBlock &Entry = F->getEntryBlock();
167   Instruction *InsertPt = &Entry.front();
168
169   LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsertPt);
170   CallInst::
171     Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_create),
172            LI, "", InsertPt);
173
174   return true;
175 }
176
177 /// CreateFailBB - Create a basic block to jump to when the stack protector
178 /// check fails.
179 BasicBlock *StackProtector::CreateFailBB() {
180   BasicBlock *FailBB = BasicBlock::Create("CallStackCheckFailBlk", F);
181   Constant *StackChkFail =
182     M->getOrInsertFunction("__stack_chk_fail", Type::VoidTy, NULL);
183   CallInst::Create(StackChkFail, "", FailBB);
184   new UnreachableInst(FailBB);
185   return FailBB;
186 }
187
188 /// RequiresStackProtector - Check whether or not this function needs a stack
189 /// protector based upon the stack protector level. The heuristic we use is to
190 /// add a guard variable to functions that call alloca, and functions with
191 /// buffers larger than 8 bytes.
192 bool StackProtector::RequiresStackProtector() const {
193   switch (Level) {
194   default: return false;
195   case SSP::ALL: return true;
196   case SSP::SOME: {
197     const TargetData *TD = TLI->getTargetData();
198
199     for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
200       BasicBlock *BB = I;
201
202       for (BasicBlock::iterator
203              II = BB->begin(), IE = BB->end(); II != IE; ++II)
204         if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
205           if (!AI->isArrayAllocation()) continue; // Only care about arrays.
206
207           if (ConstantInt *CI = dyn_cast<ConstantInt>(AI->getArraySize())) {
208             const Type *Ty = AI->getAllocatedType();
209             uint64_t TySize = TD->getABITypeSize(Ty);
210
211             // If an array has more than 8 bytes of allocated space, then we
212             // emit stack protectors.
213             if (SSPBufferSize <= TySize * CI->getZExtValue())
214               return true;
215           } else {
216             // This is a call to alloca with a variable size. Default to adding
217             // stack protectors.
218             return true;
219           }
220         }
221     }
222
223     return false;
224   }
225   }
226 }