xref: /llvm-project/llvm/lib/CodeGen/StackProtector.cpp (revision db56e5a89af8a8ed9e3d69da176dfc8e44131fa0)
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 #include "llvm/ADT/SmallPtrSet.h"
18 #include "llvm/ADT/Statistic.h"
19 #include "llvm/Analysis/BranchProbabilityInfo.h"
20 #include "llvm/Analysis/EHPersonalities.h"
21 #include "llvm/CodeGen/Passes.h"
22 #include "llvm/CodeGen/StackProtector.h"
23 #include "llvm/IR/Attributes.h"
24 #include "llvm/IR/BasicBlock.h"
25 #include "llvm/IR/Constants.h"
26 #include "llvm/IR/DataLayout.h"
27 #include "llvm/IR/DebugInfo.h"
28 #include "llvm/IR/DebugLoc.h"
29 #include "llvm/IR/DerivedTypes.h"
30 #include "llvm/IR/Function.h"
31 #include "llvm/IR/IRBuilder.h"
32 #include "llvm/IR/Instruction.h"
33 #include "llvm/IR/Instructions.h"
34 #include "llvm/IR/Intrinsics.h"
35 #include "llvm/IR/MDBuilder.h"
36 #include "llvm/IR/Module.h"
37 #include "llvm/IR/Type.h"
38 #include "llvm/IR/User.h"
39 #include "llvm/Pass.h"
40 #include "llvm/Support/Casting.h"
41 #include "llvm/Support/CommandLine.h"
42 #include "llvm/Target/TargetLowering.h"
43 #include "llvm/Target/TargetMachine.h"
44 #include "llvm/Target/TargetOptions.h"
45 #include "llvm/Target/TargetSubtargetInfo.h"
46 #include <utility>
47 
48 using namespace llvm;
49 
50 #define DEBUG_TYPE "stack-protector"
51 
52 STATISTIC(NumFunProtected, "Number of functions protected");
53 STATISTIC(NumAddrTaken, "Number of local variables that have their address"
54                         " taken.");
55 
56 static cl::opt<bool> EnableSelectionDAGSP("enable-selectiondag-sp",
57                                           cl::init(true), cl::Hidden);
58 
59 char StackProtector::ID = 0;
60 INITIALIZE_TM_PASS(StackProtector, "stack-protector", "Insert stack protectors",
61                 false, true)
62 
63 FunctionPass *llvm::createStackProtectorPass(const TargetMachine *TM) {
64   return new StackProtector(TM);
65 }
66 
67 StackProtector::SSPLayoutKind
68 StackProtector::getSSPLayout(const AllocaInst *AI) const {
69   return AI ? Layout.lookup(AI) : SSPLK_None;
70 }
71 
72 void StackProtector::adjustForColoring(const AllocaInst *From,
73                                        const AllocaInst *To) {
74   // When coloring replaces one alloca with another, transfer the SSPLayoutKind
75   // tag from the remapped to the target alloca. The remapped alloca should
76   // have a size smaller than or equal to the replacement alloca.
77   SSPLayoutMap::iterator I = Layout.find(From);
78   if (I != Layout.end()) {
79     SSPLayoutKind Kind = I->second;
80     Layout.erase(I);
81 
82     // Transfer the tag, but make sure that SSPLK_AddrOf does not overwrite
83     // SSPLK_SmallArray or SSPLK_LargeArray, and make sure that
84     // SSPLK_SmallArray does not overwrite SSPLK_LargeArray.
85     I = Layout.find(To);
86     if (I == Layout.end())
87       Layout.insert(std::make_pair(To, Kind));
88     else if (I->second != SSPLK_LargeArray && Kind != SSPLK_AddrOf)
89       I->second = Kind;
90   }
91 }
92 
93 bool StackProtector::runOnFunction(Function &Fn) {
94   F = &Fn;
95   M = F->getParent();
96   DominatorTreeWrapperPass *DTWP =
97       getAnalysisIfAvailable<DominatorTreeWrapperPass>();
98   DT = DTWP ? &DTWP->getDomTree() : nullptr;
99   TLI = TM->getSubtargetImpl(Fn)->getTargetLowering();
100   HasPrologue = false;
101   HasIRCheck = false;
102 
103   Attribute Attr = Fn.getFnAttribute("stack-protector-buffer-size");
104   if (Attr.isStringAttribute() &&
105       Attr.getValueAsString().getAsInteger(10, SSPBufferSize))
106     return false; // Invalid integer string
107 
108   if (!RequiresStackProtector())
109     return false;
110 
111   // TODO(etienneb): Functions with funclets are not correctly supported now.
112   // Do nothing if this is funclet-based personality.
113   if (Fn.hasPersonalityFn()) {
114     EHPersonality Personality = classifyEHPersonality(Fn.getPersonalityFn());
115     if (isFuncletEHPersonality(Personality))
116       return false;
117   }
118 
119   ++NumFunProtected;
120   return InsertStackProtectors();
121 }
122 
123 /// \param [out] IsLarge is set to true if a protectable array is found and
124 /// it is "large" ( >= ssp-buffer-size).  In the case of a structure with
125 /// multiple arrays, this gets set if any of them is large.
126 bool StackProtector::ContainsProtectableArray(Type *Ty, bool &IsLarge,
127                                               bool Strong,
128                                               bool InStruct) const {
129   if (!Ty)
130     return false;
131   if (ArrayType *AT = dyn_cast<ArrayType>(Ty)) {
132     if (!AT->getElementType()->isIntegerTy(8)) {
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       // However, in strong mode any array, regardless of type and size,
136       // triggers a protector.
137       if (!Strong && (InStruct || !Trip.isOSDarwin()))
138         return false;
139     }
140 
141     // If an array has more than SSPBufferSize bytes of allocated space, then we
142     // emit stack protectors.
143     if (SSPBufferSize <= M->getDataLayout().getTypeAllocSize(AT)) {
144       IsLarge = true;
145       return true;
146     }
147 
148     if (Strong)
149       // Require a protector for all arrays in strong mode
150       return true;
151   }
152 
153   const StructType *ST = dyn_cast<StructType>(Ty);
154   if (!ST)
155     return false;
156 
157   bool NeedsProtector = false;
158   for (StructType::element_iterator I = ST->element_begin(),
159                                     E = ST->element_end();
160        I != E; ++I)
161     if (ContainsProtectableArray(*I, IsLarge, Strong, true)) {
162       // If the element is a protectable array and is large (>= SSPBufferSize)
163       // then we are done.  If the protectable array is not large, then
164       // keep looking in case a subsequent element is a large array.
165       if (IsLarge)
166         return true;
167       NeedsProtector = true;
168     }
169 
170   return NeedsProtector;
171 }
172 
173 bool StackProtector::HasAddressTaken(const Instruction *AI) {
174   for (const User *U : AI->users()) {
175     if (const StoreInst *SI = dyn_cast<StoreInst>(U)) {
176       if (AI == SI->getValueOperand())
177         return true;
178     } else if (const PtrToIntInst *SI = dyn_cast<PtrToIntInst>(U)) {
179       if (AI == SI->getOperand(0))
180         return true;
181     } else if (isa<CallInst>(U)) {
182       return true;
183     } else if (isa<InvokeInst>(U)) {
184       return true;
185     } else if (const SelectInst *SI = dyn_cast<SelectInst>(U)) {
186       if (HasAddressTaken(SI))
187         return true;
188     } else if (const PHINode *PN = dyn_cast<PHINode>(U)) {
189       // Keep track of what PHI nodes we have already visited to ensure
190       // they are only visited once.
191       if (VisitedPHIs.insert(PN).second)
192         if (HasAddressTaken(PN))
193           return true;
194     } else if (const GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(U)) {
195       if (HasAddressTaken(GEP))
196         return true;
197     } else if (const BitCastInst *BI = dyn_cast<BitCastInst>(U)) {
198       if (HasAddressTaken(BI))
199         return true;
200     }
201   }
202   return false;
203 }
204 
205 /// \brief Check whether or not this function needs a stack protector based
206 /// upon the stack protector level.
207 ///
208 /// We use two heuristics: a standard (ssp) and strong (sspstrong).
209 /// The standard heuristic which will add a guard variable to functions that
210 /// call alloca with a either a variable size or a size >= SSPBufferSize,
211 /// functions with character buffers larger than SSPBufferSize, and functions
212 /// with aggregates containing character buffers larger than SSPBufferSize. The
213 /// strong heuristic will add a guard variables to functions that call alloca
214 /// regardless of size, functions with any buffer regardless of type and size,
215 /// functions with aggregates that contain any buffer regardless of type and
216 /// size, and functions that contain stack-based variables that have had their
217 /// address taken.
218 bool StackProtector::RequiresStackProtector() {
219   bool Strong = false;
220   bool NeedsProtector = false;
221   for (const BasicBlock &BB : *F)
222     for (const Instruction &I : BB)
223       if (const CallInst *CI = dyn_cast<CallInst>(&I))
224         if (CI->getCalledFunction() ==
225             Intrinsic::getDeclaration(F->getParent(),
226                                       Intrinsic::stackprotector))
227           HasPrologue = true;
228 
229   if (F->hasFnAttribute(Attribute::SafeStack))
230     return false;
231 
232   if (F->hasFnAttribute(Attribute::StackProtectReq)) {
233     NeedsProtector = true;
234     Strong = true; // Use the same heuristic as strong to determine SSPLayout
235   } else if (F->hasFnAttribute(Attribute::StackProtectStrong))
236     Strong = true;
237   else if (HasPrologue)
238     NeedsProtector = true;
239   else if (!F->hasFnAttribute(Attribute::StackProtect))
240     return false;
241 
242   for (const BasicBlock &BB : *F) {
243     for (const Instruction &I : BB) {
244       if (const AllocaInst *AI = dyn_cast<AllocaInst>(&I)) {
245         if (AI->isArrayAllocation()) {
246           if (const auto *CI = dyn_cast<ConstantInt>(AI->getArraySize())) {
247             if (CI->getLimitedValue(SSPBufferSize) >= SSPBufferSize) {
248               // A call to alloca with size >= SSPBufferSize requires
249               // stack protectors.
250               Layout.insert(std::make_pair(AI, SSPLK_LargeArray));
251               NeedsProtector = true;
252             } else if (Strong) {
253               // Require protectors for all alloca calls in strong mode.
254               Layout.insert(std::make_pair(AI, SSPLK_SmallArray));
255               NeedsProtector = true;
256             }
257           } else {
258             // A call to alloca with a variable size requires protectors.
259             Layout.insert(std::make_pair(AI, SSPLK_LargeArray));
260             NeedsProtector = true;
261           }
262           continue;
263         }
264 
265         bool IsLarge = false;
266         if (ContainsProtectableArray(AI->getAllocatedType(), IsLarge, Strong)) {
267           Layout.insert(std::make_pair(AI, IsLarge ? SSPLK_LargeArray
268                                                    : SSPLK_SmallArray));
269           NeedsProtector = true;
270           continue;
271         }
272 
273         if (Strong && HasAddressTaken(AI)) {
274           ++NumAddrTaken;
275           Layout.insert(std::make_pair(AI, SSPLK_AddrOf));
276           NeedsProtector = true;
277         }
278       }
279     }
280   }
281 
282   return NeedsProtector;
283 }
284 
285 /// Create a stack guard loading and populate whether SelectionDAG SSP is
286 /// supported.
287 static Value *getStackGuard(const TargetLoweringBase *TLI, Module *M,
288                             IRBuilder<> &B,
289                             bool *SupportsSelectionDAGSP = nullptr) {
290   if (Value *Guard = TLI->getIRStackGuard(B))
291     return B.CreateLoad(Guard, true, "StackGuard");
292 
293   // Use SelectionDAG SSP handling, since there isn't an IR guard.
294   //
295   // This is more or less weird, since we optionally output whether we
296   // should perform a SelectionDAG SP here. The reason is that it's strictly
297   // defined as !TLI->getIRStackGuard(B), where getIRStackGuard is also
298   // mutating. There is no way to get this bit without mutating the IR, so
299   // getting this bit has to happen in this right time.
300   //
301   // We could have define a new function TLI::supportsSelectionDAGSP(), but that
302   // will put more burden on the backends' overriding work, especially when it
303   // actually conveys the same information getIRStackGuard() already gives.
304   if (SupportsSelectionDAGSP)
305     *SupportsSelectionDAGSP = true;
306   TLI->insertSSPDeclarations(*M);
307   return B.CreateCall(Intrinsic::getDeclaration(M, Intrinsic::stackguard));
308 }
309 
310 /// Insert code into the entry block that stores the stack guard
311 /// variable onto the stack:
312 ///
313 ///   entry:
314 ///     StackGuardSlot = alloca i8*
315 ///     StackGuard = <stack guard>
316 ///     call void @llvm.stackprotector(StackGuard, StackGuardSlot)
317 ///
318 /// Returns true if the platform/triple supports the stackprotectorcreate pseudo
319 /// node.
320 static bool CreatePrologue(Function *F, Module *M, ReturnInst *RI,
321                            const TargetLoweringBase *TLI, AllocaInst *&AI) {
322   bool SupportsSelectionDAGSP = false;
323   IRBuilder<> B(&F->getEntryBlock().front());
324   PointerType *PtrTy = Type::getInt8PtrTy(RI->getContext());
325   AI = B.CreateAlloca(PtrTy, nullptr, "StackGuardSlot");
326 
327   Value *GuardSlot = getStackGuard(TLI, M, B, &SupportsSelectionDAGSP);
328   B.CreateCall(Intrinsic::getDeclaration(M, Intrinsic::stackprotector),
329                {GuardSlot, AI});
330   return SupportsSelectionDAGSP;
331 }
332 
333 /// InsertStackProtectors - Insert code into the prologue and epilogue of the
334 /// function.
335 ///
336 ///  - The prologue code loads and stores the stack guard onto the stack.
337 ///  - The epilogue checks the value stored in the prologue against the original
338 ///    value. It calls __stack_chk_fail if they differ.
339 bool StackProtector::InsertStackProtectors() {
340   bool SupportsSelectionDAGSP =
341       EnableSelectionDAGSP && !TM->Options.EnableFastISel;
342   AllocaInst *AI = nullptr;       // Place on stack that stores the stack guard.
343 
344   for (Function::iterator I = F->begin(), E = F->end(); I != E;) {
345     BasicBlock *BB = &*I++;
346     ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator());
347     if (!RI)
348       continue;
349 
350     // Generate prologue instrumentation if not already generated.
351     if (!HasPrologue) {
352       HasPrologue = true;
353       SupportsSelectionDAGSP &= CreatePrologue(F, M, RI, TLI, AI);
354     }
355 
356     // SelectionDAG based code generation. Nothing else needs to be done here.
357     // The epilogue instrumentation is postponed to SelectionDAG.
358     if (SupportsSelectionDAGSP)
359       break;
360 
361     // Set HasIRCheck to true, so that SelectionDAG will not generate its own
362     // version. SelectionDAG called 'shouldEmitSDCheck' to check whether
363     // instrumentation has already been generated.
364     HasIRCheck = true;
365 
366     // Generate epilogue instrumentation. The epilogue intrumentation can be
367     // function-based or inlined depending on which mechanism the target is
368     // providing.
369     if (Value* GuardCheck = TLI->getSSPStackGuardCheck(*M)) {
370       // Generate the function-based epilogue instrumentation.
371       // The target provides a guard check function, generate a call to it.
372       IRBuilder<> B(RI);
373       LoadInst *Guard = B.CreateLoad(AI, true, "Guard");
374       CallInst *Call = B.CreateCall(GuardCheck, {Guard});
375       llvm::Function *Function = cast<llvm::Function>(GuardCheck);
376       Call->setAttributes(Function->getAttributes());
377       Call->setCallingConv(Function->getCallingConv());
378     } else {
379       // Generate the epilogue with inline instrumentation.
380       // If we do not support SelectionDAG based tail calls, generate IR level
381       // tail calls.
382       //
383       // For each block with a return instruction, convert this:
384       //
385       //   return:
386       //     ...
387       //     ret ...
388       //
389       // into this:
390       //
391       //   return:
392       //     ...
393       //     %1 = <stack guard>
394       //     %2 = load StackGuardSlot
395       //     %3 = cmp i1 %1, %2
396       //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
397       //
398       //   SP_return:
399       //     ret ...
400       //
401       //   CallStackCheckFailBlk:
402       //     call void @__stack_chk_fail()
403       //     unreachable
404 
405       // Create the FailBB. We duplicate the BB every time since the MI tail
406       // merge pass will merge together all of the various BB into one including
407       // fail BB generated by the stack protector pseudo instruction.
408       BasicBlock *FailBB = CreateFailBB();
409 
410       // Split the basic block before the return instruction.
411       BasicBlock *NewBB = BB->splitBasicBlock(RI->getIterator(), "SP_return");
412 
413       // Update the dominator tree if we need to.
414       if (DT && DT->isReachableFromEntry(BB)) {
415         DT->addNewBlock(NewBB, BB);
416         DT->addNewBlock(FailBB, BB);
417       }
418 
419       // Remove default branch instruction to the new BB.
420       BB->getTerminator()->eraseFromParent();
421 
422       // Move the newly created basic block to the point right after the old
423       // basic block so that it's in the "fall through" position.
424       NewBB->moveAfter(BB);
425 
426       // Generate the stack protector instructions in the old basic block.
427       IRBuilder<> B(BB);
428       Value *Guard = getStackGuard(TLI, M, B);
429       LoadInst *LI2 = B.CreateLoad(AI, true);
430       Value *Cmp = B.CreateICmpEQ(Guard, LI2);
431       auto SuccessProb =
432           BranchProbabilityInfo::getBranchProbStackProtector(true);
433       auto FailureProb =
434           BranchProbabilityInfo::getBranchProbStackProtector(false);
435       MDNode *Weights = MDBuilder(F->getContext())
436                             .createBranchWeights(SuccessProb.getNumerator(),
437                                                  FailureProb.getNumerator());
438       B.CreateCondBr(Cmp, NewBB, FailBB, Weights);
439     }
440   }
441 
442   // Return if we didn't modify any basic blocks. i.e., there are no return
443   // statements in the function.
444   return HasPrologue;
445 }
446 
447 /// CreateFailBB - Create a basic block to jump to when the stack protector
448 /// check fails.
449 BasicBlock *StackProtector::CreateFailBB() {
450   LLVMContext &Context = F->getContext();
451   BasicBlock *FailBB = BasicBlock::Create(Context, "CallStackCheckFailBlk", F);
452   IRBuilder<> B(FailBB);
453   B.SetCurrentDebugLocation(DebugLoc::get(0, 0, F->getSubprogram()));
454   if (Trip.isOSOpenBSD()) {
455     Constant *StackChkFail =
456         M->getOrInsertFunction("__stack_smash_handler",
457                                Type::getVoidTy(Context),
458                                Type::getInt8PtrTy(Context), nullptr);
459 
460     B.CreateCall(StackChkFail, B.CreateGlobalStringPtr(F->getName(), "SSH"));
461   } else {
462     Constant *StackChkFail =
463         M->getOrInsertFunction("__stack_chk_fail", Type::getVoidTy(Context),
464                                nullptr);
465     B.CreateCall(StackChkFail, {});
466   }
467   B.CreateUnreachable();
468   return FailBB;
469 }
470 
471 bool StackProtector::shouldEmitSDCheck(const BasicBlock &BB) const {
472   return HasPrologue && !HasIRCheck && dyn_cast<ReturnInst>(BB.getTerminator());
473 }
474