1//===- SafeStack.cpp - Safe Stack Insertion -------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This pass splits the stack into the safe stack (kept as-is for LLVM backend)
10// and the unsafe stack (explicitly allocated and managed through the runtime
11// support library).
12//
13// http://clang.llvm.org/docs/SafeStack.html
14//
15//===----------------------------------------------------------------------===//
16
17#include "llvm/CodeGen/SafeStack.h"
18#include "SafeStackLayout.h"
19#include "llvm/ADT/APInt.h"
20#include "llvm/ADT/ArrayRef.h"
21#include "llvm/ADT/SmallPtrSet.h"
22#include "llvm/ADT/SmallVector.h"
23#include "llvm/ADT/Statistic.h"
24#include "llvm/Analysis/AssumptionCache.h"
25#include "llvm/Analysis/BranchProbabilityInfo.h"
26#include "llvm/Analysis/DomTreeUpdater.h"
27#include "llvm/Analysis/InlineCost.h"
28#include "llvm/Analysis/LoopInfo.h"
29#include "llvm/Analysis/ScalarEvolution.h"
30#include "llvm/Analysis/ScalarEvolutionExpressions.h"
31#include "llvm/Analysis/StackLifetime.h"
32#include "llvm/Analysis/TargetLibraryInfo.h"
33#include "llvm/CodeGen/TargetLowering.h"
34#include "llvm/CodeGen/TargetPassConfig.h"
35#include "llvm/CodeGen/TargetSubtargetInfo.h"
36#include "llvm/IR/Argument.h"
37#include "llvm/IR/Attributes.h"
38#include "llvm/IR/ConstantRange.h"
39#include "llvm/IR/Constants.h"
40#include "llvm/IR/DIBuilder.h"
41#include "llvm/IR/DataLayout.h"
42#include "llvm/IR/DerivedTypes.h"
43#include "llvm/IR/Dominators.h"
44#include "llvm/IR/Function.h"
45#include "llvm/IR/IRBuilder.h"
46#include "llvm/IR/InstIterator.h"
47#include "llvm/IR/Instruction.h"
48#include "llvm/IR/Instructions.h"
49#include "llvm/IR/IntrinsicInst.h"
50#include "llvm/IR/Intrinsics.h"
51#include "llvm/IR/MDBuilder.h"
52#include "llvm/IR/Metadata.h"
53#include "llvm/IR/Module.h"
54#include "llvm/IR/Type.h"
55#include "llvm/IR/Use.h"
56#include "llvm/IR/Value.h"
57#include "llvm/InitializePasses.h"
58#include "llvm/Pass.h"
59#include "llvm/Support/Casting.h"
60#include "llvm/Support/Debug.h"
61#include "llvm/Support/ErrorHandling.h"
62#include "llvm/Support/raw_ostream.h"
63#include "llvm/Target/TargetMachine.h"
64#include "llvm/Transforms/Utils/BasicBlockUtils.h"
65#include "llvm/Transforms/Utils/Cloning.h"
66#include "llvm/Transforms/Utils/Local.h"
67#include <algorithm>
68#include <cassert>
69#include <cstdint>
70#include <optional>
71#include <string>
72#include <utility>
73
74using namespace llvm;
75using namespace llvm::safestack;
76
77#define DEBUG_TYPE "safe-stack"
78
79namespace llvm {
80
81STATISTIC(NumFunctions, "Total number of functions");
82STATISTIC(NumUnsafeStackFunctions, "Number of functions with unsafe stack");
83STATISTIC(NumUnsafeStackRestorePointsFunctions,
84 "Number of functions that use setjmp or exceptions");
85
86STATISTIC(NumAllocas, "Total number of allocas");
87STATISTIC(NumUnsafeStaticAllocas, "Number of unsafe static allocas");
88STATISTIC(NumUnsafeDynamicAllocas, "Number of unsafe dynamic allocas");
89STATISTIC(NumUnsafeByValArguments, "Number of unsafe byval arguments");
90STATISTIC(NumUnsafeStackRestorePoints, "Number of setjmps and landingpads");
91
92} // namespace llvm
93
94/// Use __safestack_pointer_address even if the platform has a faster way of
95/// access safe stack pointer.
96static cl::opt<bool>
97 SafeStackUsePointerAddress("safestack-use-pointer-address",
98 cl::init(Val: false), cl::Hidden);
99
100static cl::opt<bool> ClColoring("safe-stack-coloring",
101 cl::desc("enable safe stack coloring"),
102 cl::Hidden, cl::init(Val: true));
103
104namespace {
105
106/// The SafeStack pass splits the stack of each function into the safe
107/// stack, which is only accessed through memory safe dereferences (as
108/// determined statically), and the unsafe stack, which contains all
109/// local variables that are accessed in ways that we can't prove to
110/// be safe.
111class SafeStack {
112 Function &F;
113 const TargetLoweringBase &TL;
114 const DataLayout &DL;
115 DomTreeUpdater *DTU;
116 ScalarEvolution &SE;
117
118 Type *StackPtrTy;
119 Type *IntPtrTy;
120 Type *Int32Ty;
121
122 Value *UnsafeStackPtr = nullptr;
123
124 /// Unsafe stack alignment. Each stack frame must ensure that the stack is
125 /// aligned to this value. We need to re-align the unsafe stack if the
126 /// alignment of any object on the stack exceeds this value.
127 ///
128 /// 16 seems like a reasonable upper bound on the alignment of objects that we
129 /// might expect to appear on the stack on most common targets.
130 static constexpr Align StackAlignment = Align::Constant<16>();
131
132 /// Return the value of the stack canary.
133 Value *getStackGuard(IRBuilder<> &IRB, Function &F);
134
135 /// Load stack guard from the frame and check if it has changed.
136 void checkStackGuard(IRBuilder<> &IRB, Function &F, Instruction &RI,
137 AllocaInst *StackGuardSlot, Value *StackGuard);
138
139 /// Find all static allocas, dynamic allocas, return instructions and
140 /// stack restore points (exception unwind blocks and setjmp calls) in the
141 /// given function and append them to the respective vectors.
142 void findInsts(Function &F, SmallVectorImpl<AllocaInst *> &StaticAllocas,
143 SmallVectorImpl<AllocaInst *> &DynamicAllocas,
144 SmallVectorImpl<Argument *> &ByValArguments,
145 SmallVectorImpl<Instruction *> &Returns,
146 SmallVectorImpl<Instruction *> &StackRestorePoints);
147
148 /// Calculate the allocation size of a given alloca. Returns 0 if the
149 /// size can not be statically determined.
150 uint64_t getStaticAllocaAllocationSize(const AllocaInst* AI);
151
152 /// Allocate space for all static allocas in \p StaticAllocas,
153 /// replace allocas with pointers into the unsafe stack.
154 ///
155 /// \returns A pointer to the top of the unsafe stack after all unsafe static
156 /// allocas are allocated.
157 Value *moveStaticAllocasToUnsafeStack(IRBuilder<> &IRB, Function &F,
158 ArrayRef<AllocaInst *> StaticAllocas,
159 ArrayRef<Argument *> ByValArguments,
160 Instruction *BasePointer,
161 AllocaInst *StackGuardSlot);
162
163 /// Generate code to restore the stack after all stack restore points
164 /// in \p StackRestorePoints.
165 ///
166 /// \returns A local variable in which to maintain the dynamic top of the
167 /// unsafe stack if needed.
168 AllocaInst *
169 createStackRestorePoints(IRBuilder<> &IRB, Function &F,
170 ArrayRef<Instruction *> StackRestorePoints,
171 Value *StaticTop, bool NeedDynamicTop);
172
173 /// Replace all allocas in \p DynamicAllocas with code to allocate
174 /// space dynamically on the unsafe stack and store the dynamic unsafe stack
175 /// top to \p DynamicTop if non-null.
176 void moveDynamicAllocasToUnsafeStack(Function &F, Value *UnsafeStackPtr,
177 AllocaInst *DynamicTop,
178 ArrayRef<AllocaInst *> DynamicAllocas);
179
180 bool IsSafeStackAlloca(const Value *AllocaPtr, uint64_t AllocaSize);
181
182 bool IsMemIntrinsicSafe(const MemIntrinsic *MI, const Use &U,
183 const Value *AllocaPtr, uint64_t AllocaSize);
184 bool IsAccessSafe(Value *Addr, uint64_t Size, const Value *AllocaPtr,
185 uint64_t AllocaSize);
186
187 bool ShouldInlinePointerAddress(CallInst &CI);
188 void TryInlinePointerAddress();
189
190public:
191 SafeStack(Function &F, const TargetLoweringBase &TL, const DataLayout &DL,
192 DomTreeUpdater *DTU, ScalarEvolution &SE)
193 : F(F), TL(TL), DL(DL), DTU(DTU), SE(SE),
194 StackPtrTy(DL.getAllocaPtrType(Ctx&: F.getContext())),
195 IntPtrTy(DL.getIntPtrType(C&: F.getContext())),
196 Int32Ty(Type::getInt32Ty(C&: F.getContext())) {}
197
198 // Run the transformation on the associated function.
199 // Returns whether the function was changed.
200 bool run();
201};
202
203constexpr Align SafeStack::StackAlignment;
204
205uint64_t SafeStack::getStaticAllocaAllocationSize(const AllocaInst* AI) {
206 uint64_t Size = DL.getTypeAllocSize(Ty: AI->getAllocatedType());
207 if (AI->isArrayAllocation()) {
208 auto C = dyn_cast<ConstantInt>(Val: AI->getArraySize());
209 if (!C)
210 return 0;
211 Size *= C->getZExtValue();
212 }
213 return Size;
214}
215
216bool SafeStack::IsAccessSafe(Value *Addr, uint64_t AccessSize,
217 const Value *AllocaPtr, uint64_t AllocaSize) {
218 const SCEV *AddrExpr = SE.getSCEV(V: Addr);
219 const auto *Base = dyn_cast<SCEVUnknown>(Val: SE.getPointerBase(V: AddrExpr));
220 if (!Base || Base->getValue() != AllocaPtr) {
221 LLVM_DEBUG(
222 dbgs() << "[SafeStack] "
223 << (isa<AllocaInst>(AllocaPtr) ? "Alloca " : "ByValArgument ")
224 << *AllocaPtr << "\n"
225 << "SCEV " << *AddrExpr << " not directly based on alloca\n");
226 return false;
227 }
228
229 const SCEV *Expr = SE.removePointerBase(S: AddrExpr);
230 uint64_t BitWidth = SE.getTypeSizeInBits(Ty: Expr->getType());
231 ConstantRange AccessStartRange = SE.getUnsignedRange(S: Expr);
232 ConstantRange SizeRange =
233 ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, AccessSize));
234 ConstantRange AccessRange = AccessStartRange.add(Other: SizeRange);
235 ConstantRange AllocaRange =
236 ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, AllocaSize));
237 bool Safe = AllocaRange.contains(CR: AccessRange);
238
239 LLVM_DEBUG(
240 dbgs() << "[SafeStack] "
241 << (isa<AllocaInst>(AllocaPtr) ? "Alloca " : "ByValArgument ")
242 << *AllocaPtr << "\n"
243 << " Access " << *Addr << "\n"
244 << " SCEV " << *Expr
245 << " U: " << SE.getUnsignedRange(Expr)
246 << ", S: " << SE.getSignedRange(Expr) << "\n"
247 << " Range " << AccessRange << "\n"
248 << " AllocaRange " << AllocaRange << "\n"
249 << " " << (Safe ? "safe" : "unsafe") << "\n");
250
251 return Safe;
252}
253
254bool SafeStack::IsMemIntrinsicSafe(const MemIntrinsic *MI, const Use &U,
255 const Value *AllocaPtr,
256 uint64_t AllocaSize) {
257 if (auto MTI = dyn_cast<MemTransferInst>(Val: MI)) {
258 if (MTI->getRawSource() != U && MTI->getRawDest() != U)
259 return true;
260 } else {
261 if (MI->getRawDest() != U)
262 return true;
263 }
264
265 const auto *Len = dyn_cast<ConstantInt>(Val: MI->getLength());
266 // Non-constant size => unsafe. FIXME: try SCEV getRange.
267 if (!Len) return false;
268 return IsAccessSafe(Addr: U, AccessSize: Len->getZExtValue(), AllocaPtr, AllocaSize);
269}
270
271/// Check whether a given allocation must be put on the safe
272/// stack or not. The function analyzes all uses of AI and checks whether it is
273/// only accessed in a memory safe way (as decided statically).
274bool SafeStack::IsSafeStackAlloca(const Value *AllocaPtr, uint64_t AllocaSize) {
275 // Go through all uses of this alloca and check whether all accesses to the
276 // allocated object are statically known to be memory safe and, hence, the
277 // object can be placed on the safe stack.
278 SmallPtrSet<const Value *, 16> Visited;
279 SmallVector<const Value *, 8> WorkList;
280 WorkList.push_back(Elt: AllocaPtr);
281
282 // A DFS search through all uses of the alloca in bitcasts/PHI/GEPs/etc.
283 while (!WorkList.empty()) {
284 const Value *V = WorkList.pop_back_val();
285 for (const Use &UI : V->uses()) {
286 auto I = cast<const Instruction>(Val: UI.getUser());
287 assert(V == UI.get());
288
289 switch (I->getOpcode()) {
290 case Instruction::Load:
291 if (!IsAccessSafe(Addr: UI, AccessSize: DL.getTypeStoreSize(Ty: I->getType()), AllocaPtr,
292 AllocaSize))
293 return false;
294 break;
295
296 case Instruction::VAArg:
297 // "va-arg" from a pointer is safe.
298 break;
299 case Instruction::Store:
300 if (V == I->getOperand(i: 0)) {
301 // Stored the pointer - conservatively assume it may be unsafe.
302 LLVM_DEBUG(dbgs()
303 << "[SafeStack] Unsafe alloca: " << *AllocaPtr
304 << "\n store of address: " << *I << "\n");
305 return false;
306 }
307
308 if (!IsAccessSafe(Addr: UI, AccessSize: DL.getTypeStoreSize(Ty: I->getOperand(i: 0)->getType()),
309 AllocaPtr, AllocaSize))
310 return false;
311 break;
312
313 case Instruction::Ret:
314 // Information leak.
315 return false;
316
317 case Instruction::Call:
318 case Instruction::Invoke: {
319 const CallBase &CS = *cast<CallBase>(Val: I);
320
321 if (I->isLifetimeStartOrEnd())
322 continue;
323
324 if (const MemIntrinsic *MI = dyn_cast<MemIntrinsic>(Val: I)) {
325 if (!IsMemIntrinsicSafe(MI, U: UI, AllocaPtr, AllocaSize)) {
326 LLVM_DEBUG(dbgs()
327 << "[SafeStack] Unsafe alloca: " << *AllocaPtr
328 << "\n unsafe memintrinsic: " << *I << "\n");
329 return false;
330 }
331 continue;
332 }
333
334 // LLVM 'nocapture' attribute is only set for arguments whose address
335 // is not stored, passed around, or used in any other non-trivial way.
336 // We assume that passing a pointer to an object as a 'nocapture
337 // readnone' argument is safe.
338 // FIXME: a more precise solution would require an interprocedural
339 // analysis here, which would look at all uses of an argument inside
340 // the function being called.
341 auto B = CS.arg_begin(), E = CS.arg_end();
342 for (const auto *A = B; A != E; ++A)
343 if (A->get() == V)
344 if (!(CS.doesNotCapture(OpNo: A - B) && (CS.doesNotAccessMemory(OpNo: A - B) ||
345 CS.doesNotAccessMemory()))) {
346 LLVM_DEBUG(dbgs() << "[SafeStack] Unsafe alloca: " << *AllocaPtr
347 << "\n unsafe call: " << *I << "\n");
348 return false;
349 }
350 continue;
351 }
352
353 default:
354 if (Visited.insert(Ptr: I).second)
355 WorkList.push_back(Elt: cast<const Instruction>(Val: I));
356 }
357 }
358 }
359
360 // All uses of the alloca are safe, we can place it on the safe stack.
361 return true;
362}
363
364Value *SafeStack::getStackGuard(IRBuilder<> &IRB, Function &F) {
365 Value *StackGuardVar = TL.getIRStackGuard(IRB);
366 Module *M = F.getParent();
367
368 if (!StackGuardVar) {
369 TL.insertSSPDeclarations(M&: *M);
370 return IRB.CreateIntrinsic(ID: Intrinsic::stackguard, Args: {});
371 }
372
373 return IRB.CreateLoad(Ty: StackPtrTy, Ptr: StackGuardVar, Name: "StackGuard");
374}
375
376void SafeStack::findInsts(Function &F,
377 SmallVectorImpl<AllocaInst *> &StaticAllocas,
378 SmallVectorImpl<AllocaInst *> &DynamicAllocas,
379 SmallVectorImpl<Argument *> &ByValArguments,
380 SmallVectorImpl<Instruction *> &Returns,
381 SmallVectorImpl<Instruction *> &StackRestorePoints) {
382 for (Instruction &I : instructions(F: &F)) {
383 if (auto AI = dyn_cast<AllocaInst>(Val: &I)) {
384 ++NumAllocas;
385
386 uint64_t Size = getStaticAllocaAllocationSize(AI);
387 if (IsSafeStackAlloca(AllocaPtr: AI, AllocaSize: Size))
388 continue;
389
390 if (AI->isStaticAlloca()) {
391 ++NumUnsafeStaticAllocas;
392 StaticAllocas.push_back(Elt: AI);
393 } else {
394 ++NumUnsafeDynamicAllocas;
395 DynamicAllocas.push_back(Elt: AI);
396 }
397 } else if (auto RI = dyn_cast<ReturnInst>(Val: &I)) {
398 if (CallInst *CI = I.getParent()->getTerminatingMustTailCall())
399 Returns.push_back(Elt: CI);
400 else
401 Returns.push_back(Elt: RI);
402 } else if (auto CI = dyn_cast<CallInst>(Val: &I)) {
403 // setjmps require stack restore.
404 if (CI->getCalledFunction() && CI->canReturnTwice())
405 StackRestorePoints.push_back(Elt: CI);
406 } else if (auto LP = dyn_cast<LandingPadInst>(Val: &I)) {
407 // Exception landing pads require stack restore.
408 StackRestorePoints.push_back(Elt: LP);
409 } else if (auto II = dyn_cast<IntrinsicInst>(Val: &I)) {
410 if (II->getIntrinsicID() == Intrinsic::gcroot)
411 report_fatal_error(
412 reason: "gcroot intrinsic not compatible with safestack attribute");
413 }
414 }
415 for (Argument &Arg : F.args()) {
416 if (!Arg.hasByValAttr())
417 continue;
418 uint64_t Size = DL.getTypeStoreSize(Ty: Arg.getParamByValType());
419 if (IsSafeStackAlloca(AllocaPtr: &Arg, AllocaSize: Size))
420 continue;
421
422 ++NumUnsafeByValArguments;
423 ByValArguments.push_back(Elt: &Arg);
424 }
425}
426
427AllocaInst *
428SafeStack::createStackRestorePoints(IRBuilder<> &IRB, Function &F,
429 ArrayRef<Instruction *> StackRestorePoints,
430 Value *StaticTop, bool NeedDynamicTop) {
431 assert(StaticTop && "The stack top isn't set.");
432
433 if (StackRestorePoints.empty())
434 return nullptr;
435
436 // We need the current value of the shadow stack pointer to restore
437 // after longjmp or exception catching.
438
439 // FIXME: On some platforms this could be handled by the longjmp/exception
440 // runtime itself.
441
442 AllocaInst *DynamicTop = nullptr;
443 if (NeedDynamicTop) {
444 // If we also have dynamic alloca's, the stack pointer value changes
445 // throughout the function. For now we store it in an alloca.
446 DynamicTop = IRB.CreateAlloca(Ty: StackPtrTy, /*ArraySize=*/nullptr,
447 Name: "unsafe_stack_dynamic_ptr");
448 IRB.CreateStore(Val: StaticTop, Ptr: DynamicTop);
449 }
450
451 // Restore current stack pointer after longjmp/exception catch.
452 for (Instruction *I : StackRestorePoints) {
453 ++NumUnsafeStackRestorePoints;
454
455 IRB.SetInsertPoint(I->getNextNode());
456 Value *CurrentTop =
457 DynamicTop ? IRB.CreateLoad(Ty: StackPtrTy, Ptr: DynamicTop) : StaticTop;
458 IRB.CreateStore(Val: CurrentTop, Ptr: UnsafeStackPtr);
459 }
460
461 return DynamicTop;
462}
463
464void SafeStack::checkStackGuard(IRBuilder<> &IRB, Function &F, Instruction &RI,
465 AllocaInst *StackGuardSlot, Value *StackGuard) {
466 Value *V = IRB.CreateLoad(Ty: StackPtrTy, Ptr: StackGuardSlot);
467 Value *Cmp = IRB.CreateICmpNE(LHS: StackGuard, RHS: V);
468
469 auto SuccessProb = BranchProbabilityInfo::getBranchProbStackProtector(IsLikely: true);
470 auto FailureProb = BranchProbabilityInfo::getBranchProbStackProtector(IsLikely: false);
471 MDNode *Weights = MDBuilder(F.getContext())
472 .createBranchWeights(TrueWeight: SuccessProb.getNumerator(),
473 FalseWeight: FailureProb.getNumerator());
474 Instruction *CheckTerm =
475 SplitBlockAndInsertIfThen(Cond: Cmp, SplitBefore: &RI, /* Unreachable */ true, BranchWeights: Weights, DTU);
476 IRBuilder<> IRBFail(CheckTerm);
477 // FIXME: respect -fsanitize-trap / -ftrap-function here?
478 FunctionCallee StackChkFail =
479 F.getParent()->getOrInsertFunction(Name: "__stack_chk_fail", RetTy: IRB.getVoidTy());
480 IRBFail.CreateCall(Callee: StackChkFail, Args: {});
481}
482
483/// We explicitly compute and set the unsafe stack layout for all unsafe
484/// static alloca instructions. We save the unsafe "base pointer" in the
485/// prologue into a local variable and restore it in the epilogue.
486Value *SafeStack::moveStaticAllocasToUnsafeStack(
487 IRBuilder<> &IRB, Function &F, ArrayRef<AllocaInst *> StaticAllocas,
488 ArrayRef<Argument *> ByValArguments, Instruction *BasePointer,
489 AllocaInst *StackGuardSlot) {
490 if (StaticAllocas.empty() && ByValArguments.empty())
491 return BasePointer;
492
493 DIBuilder DIB(*F.getParent());
494
495 StackLifetime SSC(F, StaticAllocas, StackLifetime::LivenessType::May);
496 static const StackLifetime::LiveRange NoColoringRange(1, true);
497 if (ClColoring)
498 SSC.run();
499
500 for (const auto *I : SSC.getMarkers()) {
501 auto *Op = dyn_cast<Instruction>(Val: I->getOperand(i_nocapture: 1));
502 const_cast<IntrinsicInst *>(I)->eraseFromParent();
503 // Remove the operand bitcast, too, if it has no more uses left.
504 if (Op && Op->use_empty())
505 Op->eraseFromParent();
506 }
507
508 // Unsafe stack always grows down.
509 StackLayout SSL(StackAlignment);
510 if (StackGuardSlot) {
511 Type *Ty = StackGuardSlot->getAllocatedType();
512 Align Align = std::max(a: DL.getPrefTypeAlign(Ty), b: StackGuardSlot->getAlign());
513 SSL.addObject(V: StackGuardSlot, Size: getStaticAllocaAllocationSize(AI: StackGuardSlot),
514 Alignment: Align, Range: SSC.getFullLiveRange());
515 }
516
517 for (Argument *Arg : ByValArguments) {
518 Type *Ty = Arg->getParamByValType();
519 uint64_t Size = DL.getTypeStoreSize(Ty);
520 if (Size == 0)
521 Size = 1; // Don't create zero-sized stack objects.
522
523 // Ensure the object is properly aligned.
524 Align Align = DL.getPrefTypeAlign(Ty);
525 if (auto A = Arg->getParamAlign())
526 Align = std::max(a: Align, b: *A);
527 SSL.addObject(V: Arg, Size, Alignment: Align, Range: SSC.getFullLiveRange());
528 }
529
530 for (AllocaInst *AI : StaticAllocas) {
531 Type *Ty = AI->getAllocatedType();
532 uint64_t Size = getStaticAllocaAllocationSize(AI);
533 if (Size == 0)
534 Size = 1; // Don't create zero-sized stack objects.
535
536 // Ensure the object is properly aligned.
537 Align Align = std::max(a: DL.getPrefTypeAlign(Ty), b: AI->getAlign());
538
539 SSL.addObject(V: AI, Size, Alignment: Align,
540 Range: ClColoring ? SSC.getLiveRange(AI) : NoColoringRange);
541 }
542
543 SSL.computeLayout();
544 Align FrameAlignment = SSL.getFrameAlignment();
545
546 // FIXME: tell SSL that we start at a less-then-MaxAlignment aligned location
547 // (AlignmentSkew).
548 if (FrameAlignment > StackAlignment) {
549 // Re-align the base pointer according to the max requested alignment.
550 IRB.SetInsertPoint(BasePointer->getNextNode());
551 BasePointer = cast<Instruction>(Val: IRB.CreateIntToPtr(
552 V: IRB.CreateAnd(
553 LHS: IRB.CreatePtrToInt(V: BasePointer, DestTy: IntPtrTy),
554 RHS: ConstantInt::get(Ty: IntPtrTy, V: ~(FrameAlignment.value() - 1))),
555 DestTy: StackPtrTy));
556 }
557
558 IRB.SetInsertPoint(BasePointer->getNextNode());
559
560 if (StackGuardSlot) {
561 unsigned Offset = SSL.getObjectOffset(V: StackGuardSlot);
562 Value *Off =
563 IRB.CreatePtrAdd(Ptr: BasePointer, Offset: ConstantInt::get(Ty: Int32Ty, V: -Offset));
564 Value *NewAI =
565 IRB.CreateBitCast(V: Off, DestTy: StackGuardSlot->getType(), Name: "StackGuardSlot");
566
567 // Replace alloc with the new location.
568 StackGuardSlot->replaceAllUsesWith(V: NewAI);
569 StackGuardSlot->eraseFromParent();
570 }
571
572 for (Argument *Arg : ByValArguments) {
573 unsigned Offset = SSL.getObjectOffset(V: Arg);
574 MaybeAlign Align(SSL.getObjectAlignment(V: Arg));
575 Type *Ty = Arg->getParamByValType();
576
577 uint64_t Size = DL.getTypeStoreSize(Ty);
578 if (Size == 0)
579 Size = 1; // Don't create zero-sized stack objects.
580
581 Value *Off =
582 IRB.CreatePtrAdd(Ptr: BasePointer, Offset: ConstantInt::get(Ty: Int32Ty, V: -Offset));
583 Value *NewArg = IRB.CreateBitCast(V: Off, DestTy: Arg->getType(),
584 Name: Arg->getName() + ".unsafe-byval");
585
586 // Replace alloc with the new location.
587 replaceDbgDeclare(Address: Arg, NewAddress: BasePointer, Builder&: DIB, DIExprFlags: DIExpression::ApplyOffset,
588 Offset: -Offset);
589 Arg->replaceAllUsesWith(V: NewArg);
590 IRB.SetInsertPoint(cast<Instruction>(Val: NewArg)->getNextNode());
591 IRB.CreateMemCpy(Dst: Off, DstAlign: Align, Src: Arg, SrcAlign: Arg->getParamAlign(), Size);
592 }
593
594 // Allocate space for every unsafe static AllocaInst on the unsafe stack.
595 for (AllocaInst *AI : StaticAllocas) {
596 IRB.SetInsertPoint(AI);
597 unsigned Offset = SSL.getObjectOffset(V: AI);
598
599 replaceDbgDeclare(Address: AI, NewAddress: BasePointer, Builder&: DIB, DIExprFlags: DIExpression::ApplyOffset, Offset: -Offset);
600 replaceDbgValueForAlloca(AI, NewAllocaAddress: BasePointer, Builder&: DIB, Offset: -Offset);
601
602 // Replace uses of the alloca with the new location.
603 // Insert address calculation close to each use to work around PR27844.
604 std::string Name = std::string(AI->getName()) + ".unsafe";
605 while (!AI->use_empty()) {
606 Use &U = *AI->use_begin();
607 Instruction *User = cast<Instruction>(Val: U.getUser());
608
609 Instruction *InsertBefore;
610 if (auto *PHI = dyn_cast<PHINode>(Val: User))
611 InsertBefore = PHI->getIncomingBlock(U)->getTerminator();
612 else
613 InsertBefore = User;
614
615 IRBuilder<> IRBUser(InsertBefore);
616 Value *Off =
617 IRBUser.CreatePtrAdd(Ptr: BasePointer, Offset: ConstantInt::get(Ty: Int32Ty, V: -Offset));
618 Value *Replacement =
619 IRBUser.CreateAddrSpaceCast(V: Off, DestTy: AI->getType(), Name);
620
621 if (auto *PHI = dyn_cast<PHINode>(Val: User))
622 // PHI nodes may have multiple incoming edges from the same BB (why??),
623 // all must be updated at once with the same incoming value.
624 PHI->setIncomingValueForBlock(BB: PHI->getIncomingBlock(U), V: Replacement);
625 else
626 U.set(Replacement);
627 }
628
629 AI->eraseFromParent();
630 }
631
632 // Re-align BasePointer so that our callees would see it aligned as
633 // expected.
634 // FIXME: no need to update BasePointer in leaf functions.
635 unsigned FrameSize = alignTo(Size: SSL.getFrameSize(), A: StackAlignment);
636
637 MDBuilder MDB(F.getContext());
638 SmallVector<Metadata *, 2> Data;
639 Data.push_back(Elt: MDB.createString(Str: "unsafe-stack-size"));
640 Data.push_back(Elt: MDB.createConstant(C: ConstantInt::get(Ty: Int32Ty, V: FrameSize)));
641 MDNode *MD = MDTuple::get(Context&: F.getContext(), MDs: Data);
642 F.setMetadata(KindID: LLVMContext::MD_annotation, Node: MD);
643
644 // Update shadow stack pointer in the function epilogue.
645 IRB.SetInsertPoint(BasePointer->getNextNode());
646
647 Value *StaticTop =
648 IRB.CreatePtrAdd(Ptr: BasePointer, Offset: ConstantInt::get(Ty: Int32Ty, V: -FrameSize),
649 Name: "unsafe_stack_static_top");
650 IRB.CreateStore(Val: StaticTop, Ptr: UnsafeStackPtr);
651 return StaticTop;
652}
653
654void SafeStack::moveDynamicAllocasToUnsafeStack(
655 Function &F, Value *UnsafeStackPtr, AllocaInst *DynamicTop,
656 ArrayRef<AllocaInst *> DynamicAllocas) {
657 DIBuilder DIB(*F.getParent());
658
659 for (AllocaInst *AI : DynamicAllocas) {
660 IRBuilder<> IRB(AI);
661
662 // Compute the new SP value (after AI).
663 Value *ArraySize = AI->getArraySize();
664 if (ArraySize->getType() != IntPtrTy)
665 ArraySize = IRB.CreateIntCast(V: ArraySize, DestTy: IntPtrTy, isSigned: false);
666
667 Type *Ty = AI->getAllocatedType();
668 uint64_t TySize = DL.getTypeAllocSize(Ty);
669 Value *Size = IRB.CreateMul(LHS: ArraySize, RHS: ConstantInt::get(Ty: IntPtrTy, V: TySize));
670
671 Value *SP = IRB.CreatePtrToInt(V: IRB.CreateLoad(Ty: StackPtrTy, Ptr: UnsafeStackPtr),
672 DestTy: IntPtrTy);
673 SP = IRB.CreateSub(LHS: SP, RHS: Size);
674
675 // Align the SP value to satisfy the AllocaInst, type and stack alignments.
676 auto Align = std::max(a: std::max(a: DL.getPrefTypeAlign(Ty), b: AI->getAlign()),
677 b: StackAlignment);
678
679 Value *NewTop = IRB.CreateIntToPtr(
680 V: IRB.CreateAnd(LHS: SP,
681 RHS: ConstantInt::get(Ty: IntPtrTy, V: ~uint64_t(Align.value() - 1))),
682 DestTy: StackPtrTy);
683
684 // Save the stack pointer.
685 IRB.CreateStore(Val: NewTop, Ptr: UnsafeStackPtr);
686 if (DynamicTop)
687 IRB.CreateStore(Val: NewTop, Ptr: DynamicTop);
688
689 Value *NewAI = IRB.CreatePointerCast(V: NewTop, DestTy: AI->getType());
690 if (AI->hasName() && isa<Instruction>(Val: NewAI))
691 NewAI->takeName(V: AI);
692
693 replaceDbgDeclare(Address: AI, NewAddress: NewAI, Builder&: DIB, DIExprFlags: DIExpression::ApplyOffset, Offset: 0);
694 AI->replaceAllUsesWith(V: NewAI);
695 AI->eraseFromParent();
696 }
697
698 if (!DynamicAllocas.empty()) {
699 // Now go through the instructions again, replacing stacksave/stackrestore.
700 for (Instruction &I : llvm::make_early_inc_range(Range: instructions(F: &F))) {
701 auto *II = dyn_cast<IntrinsicInst>(Val: &I);
702 if (!II)
703 continue;
704
705 if (II->getIntrinsicID() == Intrinsic::stacksave) {
706 IRBuilder<> IRB(II);
707 Instruction *LI = IRB.CreateLoad(Ty: StackPtrTy, Ptr: UnsafeStackPtr);
708 LI->takeName(V: II);
709 II->replaceAllUsesWith(V: LI);
710 II->eraseFromParent();
711 } else if (II->getIntrinsicID() == Intrinsic::stackrestore) {
712 IRBuilder<> IRB(II);
713 Instruction *SI = IRB.CreateStore(Val: II->getArgOperand(i: 0), Ptr: UnsafeStackPtr);
714 SI->takeName(V: II);
715 assert(II->use_empty());
716 II->eraseFromParent();
717 }
718 }
719 }
720}
721
722bool SafeStack::ShouldInlinePointerAddress(CallInst &CI) {
723 Function *Callee = CI.getCalledFunction();
724 if (CI.hasFnAttr(Kind: Attribute::AlwaysInline) &&
725 isInlineViable(Callee&: *Callee).isSuccess())
726 return true;
727 if (Callee->isInterposable() || Callee->hasFnAttribute(Kind: Attribute::NoInline) ||
728 CI.isNoInline())
729 return false;
730 return true;
731}
732
733void SafeStack::TryInlinePointerAddress() {
734 auto *CI = dyn_cast<CallInst>(Val: UnsafeStackPtr);
735 if (!CI)
736 return;
737
738 if(F.hasOptNone())
739 return;
740
741 Function *Callee = CI->getCalledFunction();
742 if (!Callee || Callee->isDeclaration())
743 return;
744
745 if (!ShouldInlinePointerAddress(CI&: *CI))
746 return;
747
748 InlineFunctionInfo IFI;
749 InlineFunction(CB&: *CI, IFI);
750}
751
752bool SafeStack::run() {
753 assert(F.hasFnAttribute(Attribute::SafeStack) &&
754 "Can't run SafeStack on a function without the attribute");
755 assert(!F.isDeclaration() && "Can't run SafeStack on a function declaration");
756
757 ++NumFunctions;
758
759 SmallVector<AllocaInst *, 16> StaticAllocas;
760 SmallVector<AllocaInst *, 4> DynamicAllocas;
761 SmallVector<Argument *, 4> ByValArguments;
762 SmallVector<Instruction *, 4> Returns;
763
764 // Collect all points where stack gets unwound and needs to be restored
765 // This is only necessary because the runtime (setjmp and unwind code) is
766 // not aware of the unsafe stack and won't unwind/restore it properly.
767 // To work around this problem without changing the runtime, we insert
768 // instrumentation to restore the unsafe stack pointer when necessary.
769 SmallVector<Instruction *, 4> StackRestorePoints;
770
771 // Find all static and dynamic alloca instructions that must be moved to the
772 // unsafe stack, all return instructions and stack restore points.
773 findInsts(F, StaticAllocas, DynamicAllocas, ByValArguments, Returns,
774 StackRestorePoints);
775
776 if (StaticAllocas.empty() && DynamicAllocas.empty() &&
777 ByValArguments.empty() && StackRestorePoints.empty())
778 return false; // Nothing to do in this function.
779
780 if (!StaticAllocas.empty() || !DynamicAllocas.empty() ||
781 !ByValArguments.empty())
782 ++NumUnsafeStackFunctions; // This function has the unsafe stack.
783
784 if (!StackRestorePoints.empty())
785 ++NumUnsafeStackRestorePointsFunctions;
786
787 IRBuilder<> IRB(&F.front(), F.begin()->getFirstInsertionPt());
788 // Calls must always have a debug location, or else inlining breaks. So
789 // we explicitly set a artificial debug location here.
790 if (DISubprogram *SP = F.getSubprogram())
791 IRB.SetCurrentDebugLocation(
792 DILocation::get(Context&: SP->getContext(), Line: SP->getScopeLine(), Column: 0, Scope: SP));
793 if (SafeStackUsePointerAddress) {
794 FunctionCallee Fn = F.getParent()->getOrInsertFunction(
795 Name: "__safestack_pointer_address", RetTy: IRB.getPtrTy(AddrSpace: 0));
796 UnsafeStackPtr = IRB.CreateCall(Callee: Fn);
797 } else {
798 UnsafeStackPtr = TL.getSafeStackPointerLocation(IRB);
799 }
800
801 // Load the current stack pointer (we'll also use it as a base pointer).
802 // FIXME: use a dedicated register for it ?
803 Instruction *BasePointer =
804 IRB.CreateLoad(Ty: StackPtrTy, Ptr: UnsafeStackPtr, isVolatile: false, Name: "unsafe_stack_ptr");
805 assert(BasePointer->getType() == StackPtrTy);
806
807 AllocaInst *StackGuardSlot = nullptr;
808 // FIXME: implement weaker forms of stack protector.
809 if (F.hasFnAttribute(Kind: Attribute::StackProtect) ||
810 F.hasFnAttribute(Kind: Attribute::StackProtectStrong) ||
811 F.hasFnAttribute(Kind: Attribute::StackProtectReq)) {
812 Value *StackGuard = getStackGuard(IRB, F);
813 StackGuardSlot = IRB.CreateAlloca(Ty: StackPtrTy, ArraySize: nullptr);
814 IRB.CreateStore(Val: StackGuard, Ptr: StackGuardSlot);
815
816 for (Instruction *RI : Returns) {
817 IRBuilder<> IRBRet(RI);
818 checkStackGuard(IRB&: IRBRet, F, RI&: *RI, StackGuardSlot, StackGuard);
819 }
820 }
821
822 // The top of the unsafe stack after all unsafe static allocas are
823 // allocated.
824 Value *StaticTop = moveStaticAllocasToUnsafeStack(
825 IRB, F, StaticAllocas, ByValArguments, BasePointer, StackGuardSlot);
826
827 // Safe stack object that stores the current unsafe stack top. It is updated
828 // as unsafe dynamic (non-constant-sized) allocas are allocated and freed.
829 // This is only needed if we need to restore stack pointer after longjmp
830 // or exceptions, and we have dynamic allocations.
831 // FIXME: a better alternative might be to store the unsafe stack pointer
832 // before setjmp / invoke instructions.
833 AllocaInst *DynamicTop = createStackRestorePoints(
834 IRB, F, StackRestorePoints, StaticTop, NeedDynamicTop: !DynamicAllocas.empty());
835
836 // Handle dynamic allocas.
837 moveDynamicAllocasToUnsafeStack(F, UnsafeStackPtr, DynamicTop,
838 DynamicAllocas);
839
840 // Restore the unsafe stack pointer before each return.
841 for (Instruction *RI : Returns) {
842 IRB.SetInsertPoint(RI);
843 IRB.CreateStore(Val: BasePointer, Ptr: UnsafeStackPtr);
844 }
845
846 TryInlinePointerAddress();
847
848 LLVM_DEBUG(dbgs() << "[SafeStack] safestack applied\n");
849 return true;
850}
851
852class SafeStackLegacyPass : public FunctionPass {
853 const TargetMachine *TM = nullptr;
854
855public:
856 static char ID; // Pass identification, replacement for typeid..
857
858 SafeStackLegacyPass() : FunctionPass(ID) {
859 initializeSafeStackLegacyPassPass(*PassRegistry::getPassRegistry());
860 }
861
862 void getAnalysisUsage(AnalysisUsage &AU) const override {
863 AU.addRequired<TargetPassConfig>();
864 AU.addRequired<TargetLibraryInfoWrapperPass>();
865 AU.addRequired<AssumptionCacheTracker>();
866 AU.addPreserved<DominatorTreeWrapperPass>();
867 }
868
869 bool runOnFunction(Function &F) override {
870 LLVM_DEBUG(dbgs() << "[SafeStack] Function: " << F.getName() << "\n");
871
872 if (!F.hasFnAttribute(Kind: Attribute::SafeStack)) {
873 LLVM_DEBUG(dbgs() << "[SafeStack] safestack is not requested"
874 " for this function\n");
875 return false;
876 }
877
878 if (F.isDeclaration()) {
879 LLVM_DEBUG(dbgs() << "[SafeStack] function definition"
880 " is not available\n");
881 return false;
882 }
883
884 TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
885 auto *TL = TM->getSubtargetImpl(F)->getTargetLowering();
886 if (!TL)
887 report_fatal_error(reason: "TargetLowering instance is required");
888
889 auto *DL = &F.getDataLayout();
890 auto &TLI = getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
891 auto &ACT = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
892
893 // Compute DT and LI only for functions that have the attribute.
894 // This is only useful because the legacy pass manager doesn't let us
895 // compute analyzes lazily.
896
897 DominatorTree *DT;
898 bool ShouldPreserveDominatorTree;
899 std::optional<DominatorTree> LazilyComputedDomTree;
900
901 // Do we already have a DominatorTree available from the previous pass?
902 // Note that we should *NOT* require it, to avoid the case where we end up
903 // not needing it, but the legacy PM would have computed it for us anyways.
904 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>()) {
905 DT = &DTWP->getDomTree();
906 ShouldPreserveDominatorTree = true;
907 } else {
908 // Otherwise, we need to compute it.
909 LazilyComputedDomTree.emplace(args&: F);
910 DT = &*LazilyComputedDomTree;
911 ShouldPreserveDominatorTree = false;
912 }
913
914 // Likewise, lazily compute loop info.
915 LoopInfo LI(*DT);
916
917 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
918
919 ScalarEvolution SE(F, TLI, ACT, *DT, LI);
920
921 return SafeStack(F, *TL, *DL, ShouldPreserveDominatorTree ? &DTU : nullptr,
922 SE)
923 .run();
924 }
925};
926
927} // end anonymous namespace
928
929PreservedAnalyses SafeStackPass::run(Function &F,
930 FunctionAnalysisManager &FAM) {
931 LLVM_DEBUG(dbgs() << "[SafeStack] Function: " << F.getName() << "\n");
932
933 if (!F.hasFnAttribute(Kind: Attribute::SafeStack)) {
934 LLVM_DEBUG(dbgs() << "[SafeStack] safestack is not requested"
935 " for this function\n");
936 return PreservedAnalyses::all();
937 }
938
939 if (F.isDeclaration()) {
940 LLVM_DEBUG(dbgs() << "[SafeStack] function definition"
941 " is not available\n");
942 return PreservedAnalyses::all();
943 }
944
945 auto *TL = TM->getSubtargetImpl(F)->getTargetLowering();
946 if (!TL)
947 report_fatal_error(reason: "TargetLowering instance is required");
948
949 auto &DL = F.getDataLayout();
950
951 // preserve DominatorTree
952 auto &DT = FAM.getResult<DominatorTreeAnalysis>(IR&: F);
953 auto &SE = FAM.getResult<ScalarEvolutionAnalysis>(IR&: F);
954 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
955
956 bool Changed = SafeStack(F, *TL, DL, &DTU, SE).run();
957
958 if (!Changed)
959 return PreservedAnalyses::all();
960 PreservedAnalyses PA;
961 PA.preserve<DominatorTreeAnalysis>();
962 return PA;
963}
964
965char SafeStackLegacyPass::ID = 0;
966
967INITIALIZE_PASS_BEGIN(SafeStackLegacyPass, DEBUG_TYPE,
968 "Safe Stack instrumentation pass", false, false)
969INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
970INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
971INITIALIZE_PASS_END(SafeStackLegacyPass, DEBUG_TYPE,
972 "Safe Stack instrumentation pass", false, false)
973
974FunctionPass *llvm::createSafeStackPass() { return new SafeStackLegacyPass(); }
975