[stackprotector] Removed stale TODO.
[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/Analysis.h"
19 #include "llvm/CodeGen/Passes.h"
20 #include "llvm/ADT/SmallPtrSet.h"
21 #include "llvm/ADT/Statistic.h"
22 #include "llvm/ADT/Triple.h"
23 #include "llvm/Analysis/Dominators.h"
24 #include "llvm/Analysis/ValueTracking.h"
25 #include "llvm/IR/Attributes.h"
26 #include "llvm/IR/Constants.h"
27 #include "llvm/IR/DataLayout.h"
28 #include "llvm/IR/DerivedTypes.h"
29 #include "llvm/IR/Function.h"
30 #include "llvm/IR/GlobalValue.h"
31 #include "llvm/IR/GlobalVariable.h"
32 #include "llvm/IR/Instructions.h"
33 #include "llvm/IR/IntrinsicInst.h"
34 #include "llvm/IR/Intrinsics.h"
35 #include "llvm/IR/Module.h"
36 #include "llvm/Pass.h"
37 #include "llvm/Support/CommandLine.h"
38 #include "llvm/Target/TargetLowering.h"
39 #include <cstdlib>
40 using namespace llvm;
41
42 STATISTIC(NumFunProtected, "Number of functions protected");
43 STATISTIC(NumAddrTaken, "Number of local variables that have their address"
44                         " taken.");
45
46 static cl::opt<bool>
47 EnableSelectionDAGSP("enable-selectiondag-sp", cl::init(true),
48                      cl::Hidden);
49
50 namespace {
51   class StackProtector : public FunctionPass {
52     const TargetMachine *TM;
53
54     /// TLI - Keep a pointer of a TargetLowering to consult for determining
55     /// target type sizes.
56     const TargetLoweringBase *TLI;
57     const Triple Trip;
58
59     Function *F;
60     Module *M;
61
62     DominatorTree *DT;
63
64     /// \brief The minimum size of buffers that will receive stack smashing
65     /// protection when -fstack-protection is used.
66     unsigned SSPBufferSize;
67
68     /// VisitedPHIs - The set of PHI nodes visited when determining
69     /// if a variable's reference has been taken.  This set 
70     /// is maintained to ensure we don't visit the same PHI node multiple
71     /// times.
72     SmallPtrSet<const PHINode*, 16> VisitedPHIs;
73
74     /// InsertStackProtectors - Insert code into the prologue and epilogue of
75     /// the function.
76     ///
77     ///  - The prologue code loads and stores the stack guard onto the stack.
78     ///  - The epilogue checks the value stored in the prologue against the
79     ///    original value. It calls __stack_chk_fail if they differ.
80     bool InsertStackProtectors();
81
82     /// CreateFailBB - Create a basic block to jump to when the stack protector
83     /// check fails.
84     BasicBlock *CreateFailBB();
85
86     /// ContainsProtectableArray - Check whether the type either is an array or
87     /// contains an array of sufficient size so that we need stack protectors
88     /// for it.
89     bool ContainsProtectableArray(Type *Ty, bool Strong = false,
90                                   bool InStruct = false) const;
91
92     /// \brief Check whether a stack allocation has its address taken.
93     bool HasAddressTaken(const Instruction *AI);
94
95     /// RequiresStackProtector - Check whether or not this function needs a
96     /// stack protector based upon the stack protector level.
97     bool RequiresStackProtector();
98   public:
99     static char ID;             // Pass identification, replacement for typeid.
100     StackProtector() : FunctionPass(ID), TM(0), TLI(0), SSPBufferSize(0) {
101       initializeStackProtectorPass(*PassRegistry::getPassRegistry());
102     }
103     StackProtector(const TargetMachine *TM)
104       : FunctionPass(ID), TM(TM), TLI(0), Trip(TM->getTargetTriple()),
105         SSPBufferSize(8) {
106       initializeStackProtectorPass(*PassRegistry::getPassRegistry());
107     }
108
109     virtual void getAnalysisUsage(AnalysisUsage &AU) const {
110       AU.addPreserved<DominatorTree>();
111     }
112
113     virtual bool runOnFunction(Function &Fn);
114   };
115 } // end anonymous namespace
116
117 char StackProtector::ID = 0;
118 INITIALIZE_PASS(StackProtector, "stack-protector",
119                 "Insert stack protectors", false, false)
120
121 FunctionPass *llvm::createStackProtectorPass(const TargetMachine *TM) {
122   return new StackProtector(TM);
123 }
124
125 bool StackProtector::runOnFunction(Function &Fn) {
126   F = &Fn;
127   M = F->getParent();
128   DT = getAnalysisIfAvailable<DominatorTree>();
129   TLI = TM->getTargetLowering();
130
131   if (!RequiresStackProtector()) return false;
132
133   Attribute Attr =
134     Fn.getAttributes().getAttribute(AttributeSet::FunctionIndex,
135                                     "stack-protector-buffer-size");
136   if (Attr.isStringAttribute())
137     SSPBufferSize = atoi(Attr.getValueAsString().data());
138
139   ++NumFunProtected;
140   return InsertStackProtectors();
141 }
142
143 /// ContainsProtectableArray - Check whether the type either is an array or
144 /// contains a char array of sufficient size so that we need stack protectors
145 /// for it.
146 bool StackProtector::ContainsProtectableArray(Type *Ty, bool Strong,
147                                               bool InStruct) const {
148   if (!Ty) return false;
149   if (ArrayType *AT = dyn_cast<ArrayType>(Ty)) {
150     // In strong mode any array, regardless of type and size, triggers a
151     // protector
152     if (Strong)
153       return true;
154     if (!AT->getElementType()->isIntegerTy(8)) {
155       // If we're on a non-Darwin platform or we're inside of a structure, don't
156       // add stack protectors unless the array is a character array.
157       if (InStruct || !Trip.isOSDarwin())
158           return false;
159     }
160
161     // If an array has more than SSPBufferSize bytes of allocated space, then we
162     // emit stack protectors.
163     if (SSPBufferSize <= TLI->getDataLayout()->getTypeAllocSize(AT))
164       return true;
165   }
166
167   const StructType *ST = dyn_cast<StructType>(Ty);
168   if (!ST) return false;
169
170   for (StructType::element_iterator I = ST->element_begin(),
171          E = ST->element_end(); I != E; ++I)
172     if (ContainsProtectableArray(*I, Strong, true))
173       return true;
174
175   return false;
176 }
177
178 bool StackProtector::HasAddressTaken(const Instruction *AI) {
179   for (Value::const_use_iterator UI = AI->use_begin(), UE = AI->use_end();
180         UI != UE; ++UI) {
181     const User *U = *UI;
182     if (const StoreInst *SI = dyn_cast<StoreInst>(U)) {
183       if (AI == SI->getValueOperand())
184         return true;
185     } else if (const PtrToIntInst *SI = dyn_cast<PtrToIntInst>(U)) {
186       if (AI == SI->getOperand(0))
187         return true;
188     } else if (isa<CallInst>(U)) {
189       return true;
190     } else if (isa<InvokeInst>(U)) {
191       return true;
192     } else if (const SelectInst *SI = dyn_cast<SelectInst>(U)) {
193       if (HasAddressTaken(SI))
194         return true;
195     } else if (const PHINode *PN = dyn_cast<PHINode>(U)) {
196       // Keep track of what PHI nodes we have already visited to ensure
197       // they are only visited once.
198       if (VisitedPHIs.insert(PN))
199         if (HasAddressTaken(PN))
200           return true;
201     } else if (const GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(U)) {
202       if (HasAddressTaken(GEP))
203         return true;
204     } else if (const BitCastInst *BI = dyn_cast<BitCastInst>(U)) {
205       if (HasAddressTaken(BI))
206         return true;
207     }
208   }
209   return false;
210 }
211
212 /// \brief Check whether or not this function needs a stack protector based
213 /// upon the stack protector level.
214 ///
215 /// We use two heuristics: a standard (ssp) and strong (sspstrong).
216 /// The standard heuristic which will add a guard variable to functions that
217 /// call alloca with a either a variable size or a size >= SSPBufferSize,
218 /// functions with character buffers larger than SSPBufferSize, and functions
219 /// with aggregates containing character buffers larger than SSPBufferSize. The
220 /// strong heuristic will add a guard variables to functions that call alloca
221 /// regardless of size, functions with any buffer regardless of type and size,
222 /// functions with aggregates that contain any buffer regardless of type and
223 /// size, and functions that contain stack-based variables that have had their
224 /// address taken.
225 bool StackProtector::RequiresStackProtector() {
226   bool Strong = false;
227   if (F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
228                                       Attribute::StackProtectReq))
229     return true;
230   else if (F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
231                                            Attribute::StackProtectStrong))
232     Strong = true;
233   else if (!F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
234                                             Attribute::StackProtect))
235     return false;
236
237   for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
238     BasicBlock *BB = I;
239
240     for (BasicBlock::iterator
241            II = BB->begin(), IE = BB->end(); II != IE; ++II) {
242       if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
243         if (AI->isArrayAllocation()) {
244           // SSP-Strong: Enable protectors for any call to alloca, regardless
245           // of size.
246           if (Strong)
247             return true;
248   
249           if (const ConstantInt *CI =
250                dyn_cast<ConstantInt>(AI->getArraySize())) {
251             if (CI->getLimitedValue(SSPBufferSize) >= SSPBufferSize)
252               // A call to alloca with size >= SSPBufferSize requires
253               // stack protectors.
254               return true;
255           } else {
256             // A call to alloca with a variable size requires protectors.
257             return true;
258           }
259         }
260
261         if (ContainsProtectableArray(AI->getAllocatedType(), Strong))
262           return true;
263
264         if (Strong && HasAddressTaken(AI)) {
265           ++NumAddrTaken; 
266           return true;
267         }
268       }
269     }
270   }
271
272   return false;
273 }
274
275 static bool InstructionWillNotHaveChain(const Instruction *I) {
276   return !I->mayHaveSideEffects() && !I->mayReadFromMemory() &&
277     isSafeToSpeculativelyExecute(I);
278 }
279
280 /// Identify if RI has a previous instruction in the "Tail Position" and return
281 /// it. Otherwise return 0.
282 ///
283 /// This is based off of the code in llvm::isInTailCallPosition
284 static CallInst *FindPotentialTailCall(BasicBlock *BB, ReturnInst *RI,
285                                        const TargetLoweringBase *TLI) {
286   // Establish a reasonable upper bound on the maximum amount of instructions we
287   // will look through to find a tail call.
288   unsigned SearchCounter = 0;
289   const unsigned MaxSearch = 4;
290   bool NoInterposingChain = true;
291
292   for (BasicBlock::reverse_iterator I = llvm::next(BB->rbegin()), E = BB->rend();
293        I != E && SearchCounter < MaxSearch; ++I) {
294     Instruction *Inst = &*I;
295
296     // Skip over debug intrinsics and do not allow them to affect our MaxSearch
297     // counter.
298     if (isa<DbgInfoIntrinsic>(Inst))
299       continue;
300
301     // If we find a call and the following conditions are satisifed, then we
302     // have found a tail call that satisfies at least the target independent
303     // requirements of a tail call:
304     //
305     // 1. The call site has the tail marker.
306     //
307     // 2. The call site either will not cause the creation of a chain or if a
308     // chain is necessary there are no instructions in between the callsite and
309     // the call which would create an interposing chain.
310     //
311     // 3. The return type of the function does not impede tail call
312     // optimization.
313     if (CallInst *CI = dyn_cast<CallInst>(Inst)) {
314       if (CI->isTailCall() &&
315           (InstructionWillNotHaveChain(CI) || NoInterposingChain) &&
316           returnTypeIsEligibleForTailCall(BB->getParent(), CI, RI, *TLI))
317         return CI;
318     }
319
320     // If we did not find a call see if we have an instruction that may create
321     // an interposing chain.
322     NoInterposingChain = NoInterposingChain && InstructionWillNotHaveChain(Inst);
323
324     // Increment max search.
325     SearchCounter++;
326   }
327
328   return 0;
329 }
330
331 /// Insert code into the entry block that stores the __stack_chk_guard
332 /// variable onto the stack:
333 ///
334 ///   entry:
335 ///     StackGuardSlot = alloca i8*
336 ///     StackGuard = load __stack_chk_guard
337 ///     call void @llvm.stackprotect.create(StackGuard, StackGuardSlot)
338 ///
339 /// Returns true if the platform/triple supports the stackprotectorcreate pseudo
340 /// node.
341 static bool CreatePrologue(Function *F, Module *M, ReturnInst *RI,
342                            const TargetLoweringBase *TLI, const Triple &Trip,
343                            AllocaInst *&AI, Value *&StackGuardVar) {
344   bool SupportsSelectionDAGSP = false;
345   PointerType *PtrTy = Type::getInt8PtrTy(RI->getContext());
346   unsigned AddressSpace, Offset;
347   if (TLI->getStackCookieLocation(AddressSpace, Offset)) {
348     Constant *OffsetVal =
349       ConstantInt::get(Type::getInt32Ty(RI->getContext()), Offset);
350     
351     StackGuardVar = ConstantExpr::getIntToPtr(OffsetVal,
352                                               PointerType::get(PtrTy,
353                                                                AddressSpace));
354   } else if (Trip.getOS() == llvm::Triple::OpenBSD) {
355     StackGuardVar = M->getOrInsertGlobal("__guard_local", PtrTy);
356     cast<GlobalValue>(StackGuardVar)
357       ->setVisibility(GlobalValue::HiddenVisibility);
358   } else {
359     SupportsSelectionDAGSP = true;
360     StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", PtrTy);    
361   }
362   
363   BasicBlock &Entry = F->getEntryBlock();
364   Instruction *InsPt = &Entry.front();
365   
366   AI = new AllocaInst(PtrTy, "StackGuardSlot", InsPt);
367   LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsPt);
368   
369   Value *Args[] = { LI, AI };
370   CallInst::
371     Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector),
372            Args, "", InsPt);
373
374   return SupportsSelectionDAGSP;
375 }
376
377 /// InsertStackProtectors - Insert code into the prologue and epilogue of the
378 /// function.
379 ///
380 ///  - The prologue code loads and stores the stack guard onto the stack.
381 ///  - The epilogue checks the value stored in the prologue against the original
382 ///    value. It calls __stack_chk_fail if they differ.
383 bool StackProtector::InsertStackProtectors() {
384   bool HasPrologue = false;
385   bool SupportsSelectionDAGSP = false;
386   AllocaInst *AI = 0;           // Place on stack that stores the stack guard.
387   Value *StackGuardVar = 0;  // The stack guard variable.
388
389   for (Function::iterator I = F->begin(), E = F->end(); I != E; ) {
390     BasicBlock *BB = I++;
391     ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator());
392     if (!RI) continue;
393
394     if (!HasPrologue) {
395       HasPrologue = true;
396       SupportsSelectionDAGSP = CreatePrologue(F, M, RI, TLI, Trip, AI,
397                                           StackGuardVar);
398     }    
399
400     if (EnableSelectionDAGSP && !TM->Options.EnableFastISel &&
401         SupportsSelectionDAGSP) {
402       // Since we have a potential tail call, insert the special stack check
403       // intrinsic.
404       Instruction *InsertionPt = 0;
405       if (CallInst *CI = FindPotentialTailCall(BB, RI, TLI)) {
406         InsertionPt = CI;
407       } else {        
408         InsertionPt = RI;
409         // At this point we know that BB has a return statement so it *DOES*
410         // have a terminator.
411         assert(InsertionPt != 0 && "BB must have a terminator instruction at "
412                "this point.");
413       }
414
415       Function *Intrinsic =
416         Intrinsic::getDeclaration(M, Intrinsic::stackprotectorcheck);
417       Value *Args[] = { StackGuardVar };
418       CallInst::Create(Intrinsic, Args, "", InsertionPt);
419
420     } else {
421       // If we do not support SelectionDAG based tail calls, generate IR level
422       // tail calls.
423       //
424       // For each block with a return instruction, convert this:
425       //
426       //   return:
427       //     ...
428       //     ret ...
429       //
430       // into this:
431       //
432       //   return:
433       //     ...
434       //     %1 = load __stack_chk_guard
435       //     %2 = load StackGuardSlot
436       //     %3 = cmp i1 %1, %2
437       //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
438       //
439       //   SP_return:
440       //     ret ...
441       //
442       //   CallStackCheckFailBlk:
443       //     call void @__stack_chk_fail()
444       //     unreachable
445
446       // Create the FailBB. We duplicate the BB every time since the MI tail
447       // merge pass will merge together all of the various BB into one including
448       // fail BB generated by the stack protector pseudo instruction. 
449       BasicBlock *FailBB = CreateFailBB();
450       
451       // Split the basic block before the return instruction.
452       BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
453       
454       // Update the dominator tree if we need to.
455       if (DT && DT->isReachableFromEntry(BB)) {
456         DT->addNewBlock(NewBB, BB);
457         DT->addNewBlock(FailBB, BB);
458       }
459       
460       // Remove default branch instruction to the new BB.
461       BB->getTerminator()->eraseFromParent();
462       
463       // Move the newly created basic block to the point right after the old
464       // basic block so that it's in the "fall through" position.
465       NewBB->moveAfter(BB);
466       
467       // Generate the stack protector instructions in the old basic block.
468       LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
469       LoadInst *LI2 = new LoadInst(AI, "", true, BB);
470       ICmpInst *Cmp = new ICmpInst(*BB, CmpInst::ICMP_EQ, LI1, LI2, "");
471       BranchInst::Create(NewBB, FailBB, Cmp, BB);
472     }
473   }
474
475   // Return if we didn't modify any basic blocks. I.e., there are no return
476   // statements in the function.
477   if (!HasPrologue)
478     return false;
479
480   return true;
481 }
482
483 /// CreateFailBB - Create a basic block to jump to when the stack protector
484 /// check fails.
485 BasicBlock *StackProtector::CreateFailBB() {
486   LLVMContext &Context = F->getContext();
487   BasicBlock *FailBB = BasicBlock::Create(Context, "CallStackCheckFailBlk", F);
488   if (Trip.getOS() == llvm::Triple::OpenBSD) {
489     Constant *StackChkFail = M->getOrInsertFunction(
490         "__stack_smash_handler", Type::getVoidTy(Context),
491         Type::getInt8PtrTy(Context), NULL);
492
493     Constant *NameStr = ConstantDataArray::getString(Context, F->getName());
494     Constant *FuncName =
495         new GlobalVariable(*M, NameStr->getType(), true,
496                            GlobalVariable::PrivateLinkage, NameStr, "SSH");
497
498     SmallVector<Constant *, 2> IdxList;
499     IdxList.push_back(ConstantInt::get(Type::getInt8Ty(Context), 0));
500     IdxList.push_back(ConstantInt::get(Type::getInt8Ty(Context), 0));
501
502     SmallVector<Value *, 1> Args;
503     Args.push_back(ConstantExpr::getGetElementPtr(FuncName, IdxList));
504
505     CallInst::Create(StackChkFail, Args, "", FailBB);
506   } else {
507     Constant *StackChkFail = M->getOrInsertFunction(
508         "__stack_chk_fail", Type::getVoidTy(Context), NULL);
509     CallInst::Create(StackChkFail, "", FailBB);
510   }
511   new UnreachableInst(Context, FailBB);
512   return FailBB;
513 }