Small simplification. Use the iterator already present as the insertion point.
[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   // Loop through the basic blocks that have return instructions. Convert this:
101   //
102   //   return:
103   //     ...
104   //     ret ...
105   //
106   // into this:
107   //
108   //   return:
109   //     ...
110   //     %1 = load __stack_chk_guard
111   //     %2 = load <stored stack guard>
112   //     %3 = cmp i1 %1, %2
113   //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
114   //
115   //   SP_return:
116   //     ret ...
117   //
118   //   CallStackCheckFailBlk:
119   //     call void @__stack_chk_fail()
120   //     unreachable
121   //
122   BasicBlock *FailBB = 0;       // The basic block to jump to if check fails.
123   AllocaInst *AI = 0;           // Place on stack that stores the stack guard.
124   Constant *StackGuardVar = 0;  // The stack guard variable.
125
126   for (Function::iterator I = F->begin(), E = F->end(); I != E; ) {
127     BasicBlock *BB = I;
128
129     if (ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator())) {
130       if (!FailBB) {
131         // Insert code into the entry block that stores the __stack_chk_guard
132         // variable onto the stack.
133         PointerType *PtrTy = PointerType::getUnqual(Type::Int8Ty);
134         StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", PtrTy);
135
136         BasicBlock &Entry = F->getEntryBlock();
137         Instruction *InsPt = &Entry.front();
138
139         AI = new AllocaInst(PtrTy, "StackGuardSlot", InsPt);
140         LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsPt);
141
142         Value *Args[] = { LI, AI };
143         CallInst::
144           Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_create),
145                  &Args[0], array_endof(Args), "", InsPt);
146
147         // Create the basic block to jump to when the guard check fails.
148         FailBB = CreateFailBB();
149       }
150
151       ++I; // Skip to the next block so that we don't resplit the return block.
152
153       // Split the basic block before the return instruction.
154       BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
155
156       // Move the newly created basic block to the point right after the old
157       // basic block so that it's in the "fall through" position.
158       NewBB->removeFromParent();
159       F->getBasicBlockList().insert(I, NewBB);
160
161       // Generate the stack protector instructions in the old basic block.
162       LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
163       CallInst *CI = CallInst::
164         Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_check),
165                AI, "", BB);
166       ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, CI, LI1, "", BB);
167       BranchInst::Create(NewBB, FailBB, Cmp, BB);
168     } else {
169       ++I;
170     }
171   }
172
173   // Return if we didn't modify any basic blocks. I.e., there are no return
174   // statements in the function.
175   if (!FailBB) return false;
176
177   return true;
178 }
179
180 /// CreateFailBB - Create a basic block to jump to when the stack protector
181 /// check fails.
182 BasicBlock *StackProtector::CreateFailBB() {
183   BasicBlock *FailBB = BasicBlock::Create("CallStackCheckFailBlk", F);
184   Constant *StackChkFail =
185     M->getOrInsertFunction("__stack_chk_fail", Type::VoidTy, NULL);
186   CallInst::Create(StackChkFail, "", FailBB);
187   new UnreachableInst(FailBB);
188   return FailBB;
189 }
190
191 /// RequiresStackProtector - Check whether or not this function needs a stack
192 /// protector based upon the stack protector level. The heuristic we use is to
193 /// add a guard variable to functions that call alloca, and functions with
194 /// buffers larger than 8 bytes.
195 bool StackProtector::RequiresStackProtector() const {
196   switch (Level) {
197   default: return false;
198   case SSP::ALL: return true;
199   case SSP::SOME: {
200     const TargetData *TD = TLI->getTargetData();
201
202     for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
203       BasicBlock *BB = I;
204
205       for (BasicBlock::iterator
206              II = BB->begin(), IE = BB->end(); II != IE; ++II)
207         if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
208           if (AI->isArrayAllocation())
209             // This is a call to alloca with a variable size. Emit stack
210             // protectors.
211             return true;
212
213           if (const ArrayType *AT = dyn_cast<ArrayType>(AI->getAllocatedType()))
214             // If an array has more than 8 bytes of allocated space, then we
215             // emit stack protectors.
216             if (SSPBufferSize <= TD->getABITypeSize(AT))
217               return true;
218         }
219     }
220
221     return false;
222   }
223   }
224 }