Move all of the header files which are involved in modelling the LLVM IR
[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/ADT/Triple.h"
20 #include "llvm/Analysis/Dominators.h"
21 #include "llvm/IR/Attributes.h"
22 #include "llvm/IR/Constants.h"
23 #include "llvm/IR/DataLayout.h"
24 #include "llvm/IR/DerivedTypes.h"
25 #include "llvm/IR/Function.h"
26 #include "llvm/IR/Instructions.h"
27 #include "llvm/IR/Intrinsics.h"
28 #include "llvm/IR/Module.h"
29 #include "llvm/Pass.h"
30 #include "llvm/Support/CommandLine.h"
31 #include "llvm/Target/TargetLowering.h"
32 #include "llvm/Target/TargetOptions.h"
33 using namespace llvm;
34
35 namespace {
36   class StackProtector : public FunctionPass {
37     /// TLI - Keep a pointer of a TargetLowering to consult for determining
38     /// target type sizes.
39     const TargetLowering *TLI;
40
41     Function *F;
42     Module *M;
43
44     DominatorTree *DT;
45
46     /// InsertStackProtectors - Insert code into the prologue and epilogue of
47     /// the function.
48     ///
49     ///  - The prologue code loads and stores the stack guard onto the stack.
50     ///  - The epilogue checks the value stored in the prologue against the
51     ///    original value. It calls __stack_chk_fail if they differ.
52     bool InsertStackProtectors();
53
54     /// CreateFailBB - Create a basic block to jump to when the stack protector
55     /// check fails.
56     BasicBlock *CreateFailBB();
57
58     /// ContainsProtectableArray - Check whether the type either is an array or
59     /// contains an array of sufficient size so that we need stack protectors
60     /// for it.
61     bool ContainsProtectableArray(Type *Ty, bool InStruct = false) const;
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), TLI(0) {
69       initializeStackProtectorPass(*PassRegistry::getPassRegistry());
70     }
71     StackProtector(const TargetLowering *tli)
72       : FunctionPass(ID), TLI(tli) {
73       initializeStackProtectorPass(*PassRegistry::getPassRegistry());
74     }
75
76     virtual void getAnalysisUsage(AnalysisUsage &AU) const {
77       AU.addPreserved<DominatorTree>();
78     }
79
80     virtual bool runOnFunction(Function &Fn);
81   };
82 } // end anonymous namespace
83
84 char StackProtector::ID = 0;
85 INITIALIZE_PASS(StackProtector, "stack-protector",
86                 "Insert stack protectors", false, false)
87
88 FunctionPass *llvm::createStackProtectorPass(const TargetLowering *tli) {
89   return new StackProtector(tli);
90 }
91
92 bool StackProtector::runOnFunction(Function &Fn) {
93   F = &Fn;
94   M = F->getParent();
95   DT = getAnalysisIfAvailable<DominatorTree>();
96
97   if (!RequiresStackProtector()) return false;
98
99   return InsertStackProtectors();
100 }
101
102 /// ContainsProtectableArray - Check whether the type either is an array or
103 /// contains a char array of sufficient size so that we need stack protectors
104 /// for it.
105 bool StackProtector::ContainsProtectableArray(Type *Ty, bool InStruct) const {
106   if (!Ty) return false;
107   if (ArrayType *AT = dyn_cast<ArrayType>(Ty)) {
108     const TargetMachine &TM = TLI->getTargetMachine();
109     if (!AT->getElementType()->isIntegerTy(8)) {
110       Triple Trip(TM.getTargetTriple());
111
112       // If we're on a non-Darwin platform or we're inside of a structure, don't
113       // add stack protectors unless the array is a character array.
114       if (InStruct || !Trip.isOSDarwin())
115           return false;
116     }
117
118     // If an array has more than SSPBufferSize bytes of allocated space, then we
119     // emit stack protectors.
120     if (TM.Options.SSPBufferSize <= TLI->getDataLayout()->getTypeAllocSize(AT))
121       return true;
122   }
123
124   const StructType *ST = dyn_cast<StructType>(Ty);
125   if (!ST) return false;
126
127   for (StructType::element_iterator I = ST->element_begin(),
128          E = ST->element_end(); I != E; ++I)
129     if (ContainsProtectableArray(*I, true))
130       return true;
131
132   return false;
133 }
134
135 /// RequiresStackProtector - Check whether or not this function needs a stack
136 /// protector based upon the stack protector level. The heuristic we use is to
137 /// add a guard variable to functions that call alloca, and functions with
138 /// buffers larger than SSPBufferSize bytes.
139 bool StackProtector::RequiresStackProtector() const {
140   if (F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
141                                       Attribute::StackProtectReq))
142     return true;
143
144   if (!F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
145                                        Attribute::StackProtect))
146     return false;
147
148   for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
149     BasicBlock *BB = I;
150
151     for (BasicBlock::iterator
152            II = BB->begin(), IE = BB->end(); II != IE; ++II)
153       if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
154         if (AI->isArrayAllocation())
155           // This is a call to alloca with a variable size. Emit stack
156           // protectors.
157           return true;
158
159         if (ContainsProtectableArray(AI->getAllocatedType()))
160           return true;
161       }
162   }
163
164   return false;
165 }
166
167 /// InsertStackProtectors - Insert code into the prologue and epilogue of the
168 /// function.
169 ///
170 ///  - The prologue code loads and stores the stack guard onto the stack.
171 ///  - The epilogue checks the value stored in the prologue against the original
172 ///    value. It calls __stack_chk_fail if they differ.
173 bool StackProtector::InsertStackProtectors() {
174   BasicBlock *FailBB = 0;       // The basic block to jump to if check fails.
175   BasicBlock *FailBBDom = 0;    // FailBB's dominator.
176   AllocaInst *AI = 0;           // Place on stack that stores the stack guard.
177   Value *StackGuardVar = 0;  // The stack guard variable.
178
179   for (Function::iterator I = F->begin(), E = F->end(); I != E; ) {
180     BasicBlock *BB = I++;
181     ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator());
182     if (!RI) continue;
183
184     if (!FailBB) {
185       // Insert code into the entry block that stores the __stack_chk_guard
186       // variable onto the stack:
187       //
188       //   entry:
189       //     StackGuardSlot = alloca i8*
190       //     StackGuard = load __stack_chk_guard
191       //     call void @llvm.stackprotect.create(StackGuard, StackGuardSlot)
192       //
193       PointerType *PtrTy = Type::getInt8PtrTy(RI->getContext());
194       unsigned AddressSpace, Offset;
195       if (TLI->getStackCookieLocation(AddressSpace, Offset)) {
196         Constant *OffsetVal =
197           ConstantInt::get(Type::getInt32Ty(RI->getContext()), Offset);
198
199         StackGuardVar = ConstantExpr::getIntToPtr(OffsetVal,
200                                       PointerType::get(PtrTy, AddressSpace));
201       } else {
202         StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", PtrTy);
203       }
204
205       BasicBlock &Entry = F->getEntryBlock();
206       Instruction *InsPt = &Entry.front();
207
208       AI = new AllocaInst(PtrTy, "StackGuardSlot", InsPt);
209       LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsPt);
210
211       Value *Args[] = { LI, AI };
212       CallInst::
213         Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector),
214                Args, "", InsPt);
215
216       // Create the basic block to jump to when the guard check fails.
217       FailBB = CreateFailBB();
218     }
219
220     // For each block with a return instruction, convert this:
221     //
222     //   return:
223     //     ...
224     //     ret ...
225     //
226     // into this:
227     //
228     //   return:
229     //     ...
230     //     %1 = load __stack_chk_guard
231     //     %2 = load StackGuardSlot
232     //     %3 = cmp i1 %1, %2
233     //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
234     //
235     //   SP_return:
236     //     ret ...
237     //
238     //   CallStackCheckFailBlk:
239     //     call void @__stack_chk_fail()
240     //     unreachable
241
242     // Split the basic block before the return instruction.
243     BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
244
245     if (DT && DT->isReachableFromEntry(BB)) {
246       DT->addNewBlock(NewBB, BB);
247       FailBBDom = FailBBDom ? DT->findNearestCommonDominator(FailBBDom, BB) :BB;
248     }
249
250     // Remove default branch instruction to the new BB.
251     BB->getTerminator()->eraseFromParent();
252
253     // Move the newly created basic block to the point right after the old basic
254     // block so that it's in the "fall through" position.
255     NewBB->moveAfter(BB);
256
257     // Generate the stack protector instructions in the old basic block.
258     LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
259     LoadInst *LI2 = new LoadInst(AI, "", true, BB);
260     ICmpInst *Cmp = new ICmpInst(*BB, CmpInst::ICMP_EQ, LI1, LI2, "");
261     BranchInst::Create(NewBB, FailBB, Cmp, BB);
262   }
263
264   // Return if we didn't modify any basic blocks. I.e., there are no return
265   // statements in the function.
266   if (!FailBB) return false;
267
268   if (DT && FailBBDom)
269     DT->addNewBlock(FailBB, FailBBDom);
270
271   return true;
272 }
273
274 /// CreateFailBB - Create a basic block to jump to when the stack protector
275 /// check fails.
276 BasicBlock *StackProtector::CreateFailBB() {
277   BasicBlock *FailBB = BasicBlock::Create(F->getContext(),
278                                           "CallStackCheckFailBlk", F);
279   Constant *StackChkFail =
280     M->getOrInsertFunction("__stack_chk_fail",
281                            Type::getVoidTy(F->getContext()), NULL);
282   CallInst::Create(StackChkFail, "", FailBB);
283   new UnreachableInst(F->getContext(), FailBB);
284   return FailBB;
285 }