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/SmallPtrSet.h"
20#include "llvm/ADT/Statistic.h"
21#include "llvm/ADT/Triple.h"
22#include "llvm/Analysis/Dominators.h"
23#include "llvm/IR/Attributes.h"
24#include "llvm/IR/Constants.h"
25#include "llvm/IR/DataLayout.h"
26#include "llvm/IR/DerivedTypes.h"
27#include "llvm/IR/Function.h"
28#include "llvm/IR/Instructions.h"
29#include "llvm/IR/Intrinsics.h"
30#include "llvm/IR/Module.h"
31#include "llvm/Pass.h"
32#include "llvm/Support/CommandLine.h"
33#include "llvm/Target/TargetLowering.h"
34using namespace llvm;
35
36STATISTIC(NumFunProtected, "Number of functions protected");
37STATISTIC(NumAddrTaken, "Number of local variables that have their address"
38                        " taken.");
39
40namespace {
41  class StackProtector : public FunctionPass {
42    /// TLI - Keep a pointer of a TargetLowering to consult for determining
43    /// target type sizes.
44    const TargetLoweringBase *TLI;
45
46    Function *F;
47    Module *M;
48
49    DominatorTree *DT;
50
51    /// VisitedPHIs - The set of PHI nodes visited when determining
52    /// if a variable's reference has been taken.  This set
53    /// is maintained to ensure we don't visit the same PHI node multiple
54    /// times.
55    SmallPtrSet<const PHINode*, 16> VisitedPHIs;
56
57    /// InsertStackProtectors - Insert code into the prologue and epilogue of
58    /// the function.
59    ///
60    ///  - The prologue code loads and stores the stack guard onto the stack.
61    ///  - The epilogue checks the value stored in the prologue against the
62    ///    original value. It calls __stack_chk_fail if they differ.
63    bool InsertStackProtectors();
64
65    /// CreateFailBB - Create a basic block to jump to when the stack protector
66    /// check fails.
67    BasicBlock *CreateFailBB();
68
69    /// ContainsProtectableArray - Check whether the type either is an array or
70    /// contains an array of sufficient size so that we need stack protectors
71    /// for it.
72    bool ContainsProtectableArray(Type *Ty, bool Strong = false,
73                                  bool InStruct = false) const;
74
75    /// \brief Check whether a stack allocation has its address taken.
76    bool HasAddressTaken(const Instruction *AI);
77
78    /// RequiresStackProtector - Check whether or not this function needs a
79    /// stack protector based upon the stack protector level.
80    bool RequiresStackProtector();
81  public:
82    static char ID;             // Pass identification, replacement for typeid.
83    StackProtector() : FunctionPass(ID), TLI(0) {
84      initializeStackProtectorPass(*PassRegistry::getPassRegistry());
85    }
86    StackProtector(const TargetLoweringBase *tli)
87      : FunctionPass(ID), TLI(tli) {
88      initializeStackProtectorPass(*PassRegistry::getPassRegistry());
89    }
90
91    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
92      AU.addPreserved<DominatorTree>();
93    }
94
95    virtual bool runOnFunction(Function &Fn);
96  };
97} // end anonymous namespace
98
99char StackProtector::ID = 0;
100INITIALIZE_PASS(StackProtector, "stack-protector",
101                "Insert stack protectors", false, false)
102
103FunctionPass *llvm::createStackProtectorPass(const TargetLoweringBase *tli) {
104  return new StackProtector(tli);
105}
106
107bool StackProtector::runOnFunction(Function &Fn) {
108  F = &Fn;
109  M = F->getParent();
110  DT = getAnalysisIfAvailable<DominatorTree>();
111
112  if (!RequiresStackProtector()) return false;
113
114  ++NumFunProtected;
115  return InsertStackProtectors();
116}
117
118/// ContainsProtectableArray - Check whether the type either is an array or
119/// contains a char array of sufficient size so that we need stack protectors
120/// for it.
121bool StackProtector::ContainsProtectableArray(Type *Ty, bool Strong,
122                                              bool InStruct) const {
123  if (!Ty) return false;
124  if (ArrayType *AT = dyn_cast<ArrayType>(Ty)) {
125    // In strong mode any array, regardless of type and size, triggers a
126    // protector
127    if (Strong)
128      return true;
129    const TargetMachine &TM = TLI->getTargetMachine();
130    if (!AT->getElementType()->isIntegerTy(8)) {
131      Triple Trip(TM.getTargetTriple());
132
133      // If we're on a non-Darwin platform or we're inside of a structure, don't
134      // add stack protectors unless the array is a character array.
135      if (InStruct || !Trip.isOSDarwin())
136          return false;
137    }
138
139    // If an array has more than SSPBufferSize bytes of allocated space, then we
140    // emit stack protectors.
141    if (TM.Options.SSPBufferSize <= TLI->getDataLayout()->getTypeAllocSize(AT))
142      return true;
143  }
144
145  const StructType *ST = dyn_cast<StructType>(Ty);
146  if (!ST) return false;
147
148  for (StructType::element_iterator I = ST->element_begin(),
149         E = ST->element_end(); I != E; ++I)
150    if (ContainsProtectableArray(*I, Strong, true))
151      return true;
152
153  return false;
154}
155
156bool StackProtector::HasAddressTaken(const Instruction *AI) {
157  for (Value::const_use_iterator UI = AI->use_begin(), UE = AI->use_end();
158        UI != UE; ++UI) {
159    const User *U = *UI;
160    if (const StoreInst *SI = dyn_cast<StoreInst>(U)) {
161      if (AI == SI->getValueOperand())
162        return true;
163    } else if (const PtrToIntInst *SI = dyn_cast<PtrToIntInst>(U)) {
164      if (AI == SI->getOperand(0))
165        return true;
166    } else if (isa<CallInst>(U)) {
167      return true;
168    } else if (isa<InvokeInst>(U)) {
169      return true;
170    } else if (const SelectInst *SI = dyn_cast<SelectInst>(U)) {
171      if (HasAddressTaken(SI))
172        return true;
173    } else if (const PHINode *PN = dyn_cast<PHINode>(U)) {
174      // Keep track of what PHI nodes we have already visited to ensure
175      // they are only visited once.
176      if (VisitedPHIs.insert(PN))
177        if (HasAddressTaken(PN))
178          return true;
179    } else if (const GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(U)) {
180      if (HasAddressTaken(GEP))
181        return true;
182    } else if (const BitCastInst *BI = dyn_cast<BitCastInst>(U)) {
183      if (HasAddressTaken(BI))
184        return true;
185    }
186  }
187  return false;
188}
189
190/// \brief Check whether or not this function needs a stack protector based
191/// upon the stack protector level.
192///
193/// We use two heuristics: a standard (ssp) and strong (sspstrong).
194/// The standard heuristic which will add a guard variable to functions that
195/// call alloca with a either a variable size or a size >= SSPBufferSize,
196/// functions with character buffers larger than SSPBufferSize, and functions
197/// with aggregates containing character buffers larger than SSPBufferSize. The
198/// strong heuristic will add a guard variables to functions that call alloca
199/// regardless of size, functions with any buffer regardless of type and size,
200/// functions with aggregates that contain any buffer regardless of type and
201/// size, and functions that contain stack-based variables that have had their
202/// address taken.
203bool StackProtector::RequiresStackProtector() {
204  bool Strong = false;
205  if (F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
206                                      Attribute::StackProtectReq))
207    return true;
208  else if (F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
209                                           Attribute::StackProtectStrong))
210    Strong = true;
211  else if (!F->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
212                                            Attribute::StackProtect))
213    return false;
214
215  for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
216    BasicBlock *BB = I;
217
218    for (BasicBlock::iterator
219           II = BB->begin(), IE = BB->end(); II != IE; ++II) {
220      if (AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
221        if (AI->isArrayAllocation()) {
222          // SSP-Strong: Enable protectors for any call to alloca, regardless
223          // of size.
224          if (Strong)
225            return true;
226
227          if (const ConstantInt *CI =
228               dyn_cast<ConstantInt>(AI->getArraySize())) {
229            unsigned BufferSize = TLI->getTargetMachine().Options.SSPBufferSize;
230            if (CI->getLimitedValue(BufferSize) >= BufferSize)
231              // A call to alloca with size >= SSPBufferSize requires
232              // stack protectors.
233              return true;
234          } else // A call to alloca with a variable size requires protectors.
235            return true;
236        }
237
238        if (ContainsProtectableArray(AI->getAllocatedType(), Strong))
239          return true;
240
241        if (Strong && HasAddressTaken(AI)) {
242          ++NumAddrTaken;
243          return true;
244        }
245      }
246    }
247  }
248
249  return false;
250}
251
252/// InsertStackProtectors - Insert code into the prologue and epilogue of the
253/// function.
254///
255///  - The prologue code loads and stores the stack guard onto the stack.
256///  - The epilogue checks the value stored in the prologue against the original
257///    value. It calls __stack_chk_fail if they differ.
258bool StackProtector::InsertStackProtectors() {
259  BasicBlock *FailBB = 0;       // The basic block to jump to if check fails.
260  BasicBlock *FailBBDom = 0;    // FailBB's dominator.
261  AllocaInst *AI = 0;           // Place on stack that stores the stack guard.
262  Value *StackGuardVar = 0;  // The stack guard variable.
263
264  for (Function::iterator I = F->begin(), E = F->end(); I != E; ) {
265    BasicBlock *BB = I++;
266    ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator());
267    if (!RI) continue;
268
269    if (!FailBB) {
270      // Insert code into the entry block that stores the __stack_chk_guard
271      // variable onto the stack:
272      //
273      //   entry:
274      //     StackGuardSlot = alloca i8*
275      //     StackGuard = load __stack_chk_guard
276      //     call void @llvm.stackprotect.create(StackGuard, StackGuardSlot)
277      //
278      PointerType *PtrTy = Type::getInt8PtrTy(RI->getContext());
279      unsigned AddressSpace, Offset;
280      if (TLI->getStackCookieLocation(AddressSpace, Offset)) {
281        Constant *OffsetVal =
282          ConstantInt::get(Type::getInt32Ty(RI->getContext()), Offset);
283
284        StackGuardVar = ConstantExpr::getIntToPtr(OffsetVal,
285                                      PointerType::get(PtrTy, AddressSpace));
286      } else {
287        StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", PtrTy);
288      }
289
290      BasicBlock &Entry = F->getEntryBlock();
291      Instruction *InsPt = &Entry.front();
292
293      AI = new AllocaInst(PtrTy, "StackGuardSlot", InsPt);
294      LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsPt);
295
296      Value *Args[] = { LI, AI };
297      CallInst::
298        Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector),
299               Args, "", InsPt);
300
301      // Create the basic block to jump to when the guard check fails.
302      FailBB = CreateFailBB();
303    }
304
305    // For each block with a return instruction, convert this:
306    //
307    //   return:
308    //     ...
309    //     ret ...
310    //
311    // into this:
312    //
313    //   return:
314    //     ...
315    //     %1 = load __stack_chk_guard
316    //     %2 = load StackGuardSlot
317    //     %3 = cmp i1 %1, %2
318    //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
319    //
320    //   SP_return:
321    //     ret ...
322    //
323    //   CallStackCheckFailBlk:
324    //     call void @__stack_chk_fail()
325    //     unreachable
326
327    // Split the basic block before the return instruction.
328    BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
329
330    if (DT && DT->isReachableFromEntry(BB)) {
331      DT->addNewBlock(NewBB, BB);
332      FailBBDom = FailBBDom ? DT->findNearestCommonDominator(FailBBDom, BB) :BB;
333    }
334
335    // Remove default branch instruction to the new BB.
336    BB->getTerminator()->eraseFromParent();
337
338    // Move the newly created basic block to the point right after the old basic
339    // block so that it's in the "fall through" position.
340    NewBB->moveAfter(BB);
341
342    // Generate the stack protector instructions in the old basic block.
343    LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
344    LoadInst *LI2 = new LoadInst(AI, "", true, BB);
345    ICmpInst *Cmp = new ICmpInst(*BB, CmpInst::ICMP_EQ, LI1, LI2, "");
346    BranchInst::Create(NewBB, FailBB, Cmp, BB);
347  }
348
349  // Return if we didn't modify any basic blocks. I.e., there are no return
350  // statements in the function.
351  if (!FailBB) return false;
352
353  if (DT && FailBBDom)
354    DT->addNewBlock(FailBB, FailBBDom);
355
356  return true;
357}
358
359/// CreateFailBB - Create a basic block to jump to when the stack protector
360/// check fails.
361BasicBlock *StackProtector::CreateFailBB() {
362  BasicBlock *FailBB = BasicBlock::Create(F->getContext(),
363                                          "CallStackCheckFailBlk", F);
364  Constant *StackChkFail =
365    M->getOrInsertFunction("__stack_chk_fail",
366                           Type::getVoidTy(F->getContext()), NULL);
367  CallInst::Create(StackChkFail, "", FailBB);
368  new UnreachableInst(F->getContext(), FailBB);
369  return FailBB;
370}
371