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 TargetLoweringBase *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 TargetLoweringBase *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 TargetLoweringBase *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   // FIXME: Dummy SSP-strong implementation.  Default to required until
145   // strong heuristic is implemented.
146   if (F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
147                                         Attribute::StackProtectStrong))
148     return true;
149 
150   if (!F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
151                                        Attribute::StackProtect))
152     return false;
153 
154   for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
155     BasicBlock *BB = I;
156 
157     for (BasicBlock::iterator
158            II = BB->begin(), IE = BB->end(); II != IE; ++II)
159       if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
160         if (AI->isArrayAllocation())
161           // This is a call to alloca with a variable size. Emit stack
162           // protectors.
163           return true;
164 
165         if (ContainsProtectableArray(AI->getAllocatedType()))
166           return true;
167       }
168   }
169 
170   return false;
171 }
172 
173 /// InsertStackProtectors - Insert code into the prologue and epilogue of the
174 /// function.
175 ///
176 ///  - The prologue code loads and stores the stack guard onto the stack.
177 ///  - The epilogue checks the value stored in the prologue against the original
178 ///    value. It calls __stack_chk_fail if they differ.
179 bool StackProtector::InsertStackProtectors() {
180   BasicBlock *FailBB = 0;       // The basic block to jump to if check fails.
181   BasicBlock *FailBBDom = 0;    // FailBB's dominator.
182   AllocaInst *AI = 0;           // Place on stack that stores the stack guard.
183   Value *StackGuardVar = 0;  // The stack guard variable.
184 
185   for (Function::iterator I = F->begin(), E = F->end(); I != E; ) {
186     BasicBlock *BB = I++;
187     ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator());
188     if (!RI) continue;
189 
190     if (!FailBB) {
191       // Insert code into the entry block that stores the __stack_chk_guard
192       // variable onto the stack:
193       //
194       //   entry:
195       //     StackGuardSlot = alloca i8*
196       //     StackGuard = load __stack_chk_guard
197       //     call void @llvm.stackprotect.create(StackGuard, StackGuardSlot)
198       //
199       PointerType *PtrTy = Type::getInt8PtrTy(RI->getContext());
200       unsigned AddressSpace, Offset;
201       if (TLI->getStackCookieLocation(AddressSpace, Offset)) {
202         Constant *OffsetVal =
203           ConstantInt::get(Type::getInt32Ty(RI->getContext()), Offset);
204 
205         StackGuardVar = ConstantExpr::getIntToPtr(OffsetVal,
206                                       PointerType::get(PtrTy, AddressSpace));
207       } else {
208         StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", PtrTy);
209       }
210 
211       BasicBlock &Entry = F->getEntryBlock();
212       Instruction *InsPt = &Entry.front();
213 
214       AI = new AllocaInst(PtrTy, "StackGuardSlot", InsPt);
215       LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsPt);
216 
217       Value *Args[] = { LI, AI };
218       CallInst::
219         Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector),
220                Args, "", InsPt);
221 
222       // Create the basic block to jump to when the guard check fails.
223       FailBB = CreateFailBB();
224     }
225 
226     // For each block with a return instruction, convert this:
227     //
228     //   return:
229     //     ...
230     //     ret ...
231     //
232     // into this:
233     //
234     //   return:
235     //     ...
236     //     %1 = load __stack_chk_guard
237     //     %2 = load StackGuardSlot
238     //     %3 = cmp i1 %1, %2
239     //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
240     //
241     //   SP_return:
242     //     ret ...
243     //
244     //   CallStackCheckFailBlk:
245     //     call void @__stack_chk_fail()
246     //     unreachable
247 
248     // Split the basic block before the return instruction.
249     BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
250 
251     if (DT && DT->isReachableFromEntry(BB)) {
252       DT->addNewBlock(NewBB, BB);
253       FailBBDom = FailBBDom ? DT->findNearestCommonDominator(FailBBDom, BB) :BB;
254     }
255 
256     // Remove default branch instruction to the new BB.
257     BB->getTerminator()->eraseFromParent();
258 
259     // Move the newly created basic block to the point right after the old basic
260     // block so that it's in the "fall through" position.
261     NewBB->moveAfter(BB);
262 
263     // Generate the stack protector instructions in the old basic block.
264     LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
265     LoadInst *LI2 = new LoadInst(AI, "", true, BB);
266     ICmpInst *Cmp = new ICmpInst(*BB, CmpInst::ICMP_EQ, LI1, LI2, "");
267     BranchInst::Create(NewBB, FailBB, Cmp, BB);
268   }
269 
270   // Return if we didn't modify any basic blocks. I.e., there are no return
271   // statements in the function.
272   if (!FailBB) return false;
273 
274   if (DT && FailBBDom)
275     DT->addNewBlock(FailBB, FailBBDom);
276 
277   return true;
278 }
279 
280 /// CreateFailBB - Create a basic block to jump to when the stack protector
281 /// check fails.
282 BasicBlock *StackProtector::CreateFailBB() {
283   BasicBlock *FailBB = BasicBlock::Create(F->getContext(),
284                                           "CallStackCheckFailBlk", F);
285   Constant *StackChkFail =
286     M->getOrInsertFunction("__stack_chk_fail",
287                            Type::getVoidTy(F->getContext()), NULL);
288   CallInst::Create(StackChkFail, "", FailBB);
289   new UnreachableInst(F->getContext(), FailBB);
290   return FailBB;
291 }
292