d724bd71f8042d7282b3811c17ee28e9256a7e86
[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/Attributes.h"
20 #include "llvm/Constants.h"
21 #include "llvm/DerivedTypes.h"
22 #include "llvm/Function.h"
23 #include "llvm/Instructions.h"
24 #include "llvm/Intrinsics.h"
25 #include "llvm/Module.h"
26 #include "llvm/Pass.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     /// TLI - Keep a pointer of a TargetLowering to consult for determining
42     /// target type sizes.
43     const TargetLowering *TLI;
44
45     Function *F;
46     Module *M;
47
48     /// InsertStackProtectors - Insert code into the prologue and epilogue of
49     /// the function.
50     ///
51     ///  - The prologue code loads and stores the stack guard onto the stack.
52     ///  - The epilogue checks the value stored in the prologue against the
53     ///    original value. It calls __stack_chk_fail if they differ.
54     bool InsertStackProtectors();
55
56     /// CreateFailBB - Create a basic block to jump to when the stack protector
57     /// check fails.
58     BasicBlock *CreateFailBB();
59
60     /// RequiresStackProtector - Check whether or not this function needs a
61     /// stack protector based upon the stack protector level.
62     bool RequiresStackProtector() const;
63   public:
64     static char ID;             // Pass identification, replacement for typeid.
65     StackProtector() : FunctionPass(&ID), TLI(0) {}
66     StackProtector(const TargetLowering *tli)
67       : FunctionPass(&ID), TLI(tli) {}
68
69     virtual bool runOnFunction(Function &Fn);
70   };
71 } // end anonymous namespace
72
73 char StackProtector::ID = 0;
74 static RegisterPass<StackProtector>
75 X("stack-protector", "Insert stack protectors");
76
77 FunctionPass *llvm::createStackProtectorPass(const TargetLowering *tli) {
78   return new StackProtector(tli);
79 }
80
81 bool StackProtector::runOnFunction(Function &Fn) {
82   F = &Fn;
83   M = F->getParent();
84
85   if (!RequiresStackProtector()) return false;
86   
87   return InsertStackProtectors();
88 }
89
90 /// InsertStackProtectors - Insert code into the prologue and epilogue of the
91 /// function.
92 ///
93 ///  - The prologue code loads and stores the stack guard onto the stack.
94 ///  - The epilogue checks the value stored in the prologue against the original
95 ///    value. It calls __stack_chk_fail if they differ.
96 bool StackProtector::InsertStackProtectors() {
97   // Loop through the basic blocks that have return instructions. Convert this:
98   //
99   //   return:
100   //     ...
101   //     ret ...
102   //
103   // into this:
104   //
105   //   return:
106   //     ...
107   //     %1 = load __stack_chk_guard
108   //     %2 = load <stored stack guard>
109   //     %3 = cmp i1 %1, %2
110   //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
111   //
112   //   SP_return:
113   //     ret ...
114   //
115   //   CallStackCheckFailBlk:
116   //     call void @__stack_chk_fail()
117   //     unreachable
118   //
119   BasicBlock *FailBB = 0;       // The basic block to jump to if check fails.
120   AllocaInst *AI = 0;           // Place on stack that stores the stack guard.
121   Constant *StackGuardVar = 0;  // The stack guard variable.
122
123   for (Function::iterator I = F->begin(), E = F->end(); I != E; ) {
124     BasicBlock *BB = I;
125
126     if (ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator())) {
127       if (!FailBB) {
128         // Insert code into the entry block that stores the __stack_chk_guard
129         // variable onto the stack.
130         PointerType *PtrTy = PointerType::getUnqual(Type::Int8Ty);
131         StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", PtrTy);
132
133         BasicBlock &Entry = F->getEntryBlock();
134         Instruction *InsPt = &Entry.front();
135
136         AI = new AllocaInst(PtrTy, "StackGuardSlot", InsPt);
137         LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsPt);
138
139         Value *Args[] = { LI, AI };
140         CallInst::
141           Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_create),
142                  &Args[0], array_endof(Args), "", InsPt);
143
144         // Create the basic block to jump to when the guard check fails.
145         FailBB = CreateFailBB();
146       }
147
148       ++I; // Skip to the next block so that we don't resplit the return block.
149
150       // Split the basic block before the return instruction.
151       BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
152
153       // Move the newly created basic block to the point right after the old
154       // basic block so that it's in the "fall through" position.
155       NewBB->removeFromParent();
156       F->getBasicBlockList().insert(I, NewBB);
157
158       // Generate the stack protector instructions in the old basic block.
159       LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
160       CallInst *CI = CallInst::
161         Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_check),
162                AI, "", BB);
163       ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, CI, LI1, "", BB);
164       BranchInst::Create(NewBB, FailBB, Cmp, BB);
165     } else {
166       ++I;
167     }
168   }
169
170   // Return if we didn't modify any basic blocks. I.e., there are no return
171   // statements in the function.
172   if (!FailBB) return false;
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   if (F->hasFnAttr(Attribute::StackProtectReq))
194       return true;
195
196   if (F->hasFnAttr(Attribute::StackProtect)) {
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())
206             // This is a call to alloca with a variable size. Emit stack
207             // protectors.
208             return true;
209
210           if (const ArrayType *AT = dyn_cast<ArrayType>(AI->getAllocatedType()))
211             // If an array has more than 8 bytes of allocated space, then we
212             // emit stack protectors.
213             if (SSPBufferSize <= TD->getABITypeSize(AT))
214               return true;
215         }
216     }
217
218     return false;
219   }
220
221   return false;
222 }