contrib/llvm/lib/Transforms/Instrumentation/SafeStack.cpp

   1 //===-- SafeStack.cpp - Safe Stack 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 splits the stack into the safe stack (kept as-is for LLVM backend)
  11 // and the unsafe stack (explicitly allocated and managed through the runtime
  12 // support library).
  13 //
  14 // http://clang.llvm.org/docs/SafeStack.html
  15 //
  16 //===----------------------------------------------------------------------===//
  17
  18 #include "llvm/Transforms/Instrumentation.h"
  19 #include "llvm/ADT/Statistic.h"
  20 #include "llvm/ADT/Triple.h"
  21 #include "llvm/Analysis/ScalarEvolution.h"
  22 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
  23 #include "llvm/CodeGen/Passes.h"
  24 #include "llvm/IR/Constants.h"
  25 #include "llvm/IR/DataLayout.h"
  26 #include "llvm/IR/DerivedTypes.h"
  27 #include "llvm/IR/DIBuilder.h"
  28 #include "llvm/IR/Function.h"
  29 #include "llvm/IR/InstIterator.h"
  30 #include "llvm/IR/Instructions.h"
  31 #include "llvm/IR/IntrinsicInst.h"
  32 #include "llvm/IR/Intrinsics.h"
  33 #include "llvm/IR/IRBuilder.h"
  34 #include "llvm/IR/Module.h"
  35 #include "llvm/Pass.h"
  36 #include "llvm/Support/CommandLine.h"
  37 #include "llvm/Support/Debug.h"
  38 #include "llvm/Support/Format.h"
  39 #include "llvm/Support/MathExtras.h"
  40 #include "llvm/Support/raw_os_ostream.h"
  41 #include "llvm/Target/TargetLowering.h"
  42 #include "llvm/Target/TargetSubtargetInfo.h"
  43 #include "llvm/Transforms/Utils/Local.h"
  44 #include "llvm/Transforms/Utils/ModuleUtils.h"
  45
  46 using namespace llvm;
  47
  48 #define DEBUG_TYPE "safestack"
  49
  50 enum UnsafeStackPtrStorageVal { ThreadLocalUSP, SingleThreadUSP };
  51
  52 static cl::opt<UnsafeStackPtrStorageVal> USPStorage("safe-stack-usp-storage",
  53     cl::Hidden, cl::init(ThreadLocalUSP),
  54     cl::desc("Type of storage for the unsafe stack pointer"),
  55     cl::values(clEnumValN(ThreadLocalUSP, "thread-local",
  56                           "Thread-local storage"),
  57                clEnumValN(SingleThreadUSP, "single-thread",
  58                           "Non-thread-local storage"),
  59                clEnumValEnd));
  60
  61 namespace llvm {
  62
  63 STATISTIC(NumFunctions, "Total number of functions");
  64 STATISTIC(NumUnsafeStackFunctions, "Number of functions with unsafe stack");
  65 STATISTIC(NumUnsafeStackRestorePointsFunctions,
  66           "Number of functions that use setjmp or exceptions");
  67
  68 STATISTIC(NumAllocas, "Total number of allocas");
  69 STATISTIC(NumUnsafeStaticAllocas, "Number of unsafe static allocas");
  70 STATISTIC(NumUnsafeDynamicAllocas, "Number of unsafe dynamic allocas");
  71 STATISTIC(NumUnsafeByValArguments, "Number of unsafe byval arguments");
  72 STATISTIC(NumUnsafeStackRestorePoints, "Number of setjmps and landingpads");
  73
  74 } // namespace llvm
  75
  76 namespace {
  77
  78 /// Rewrite an SCEV expression for a memory access address to an expression that
  79 /// represents offset from the given alloca.
  80 ///
  81 /// The implementation simply replaces all mentions of the alloca with zero.
  82 class AllocaOffsetRewriter : public SCEVRewriteVisitor<AllocaOffsetRewriter> {
  83   const Value *AllocaPtr;
  84
  85 public:
  86   AllocaOffsetRewriter(ScalarEvolution &SE, const Value *AllocaPtr)
  87       : SCEVRewriteVisitor(SE), AllocaPtr(AllocaPtr) {}
  88
  89   const SCEV *visitUnknown(const SCEVUnknown *Expr) {
  90     if (Expr->getValue() == AllocaPtr)
  91       return SE.getZero(Expr->getType());
  92     return Expr;
  93   }
  94 };
  95
  96 /// The SafeStack pass splits the stack of each function into the safe
  97 /// stack, which is only accessed through memory safe dereferences (as
  98 /// determined statically), and the unsafe stack, which contains all
  99 /// local variables that are accessed in ways that we can't prove to
 100 /// be safe.
 101 class SafeStack : public FunctionPass {
 102   const TargetMachine *TM;
 103   const TargetLoweringBase *TL;
 104   const DataLayout *DL;
 105   ScalarEvolution *SE;
 106
 107   Type *StackPtrTy;
 108   Type *IntPtrTy;
 109   Type *Int32Ty;
 110   Type *Int8Ty;
 111
 112   Value *UnsafeStackPtr = nullptr;
 113
 114   /// Unsafe stack alignment. Each stack frame must ensure that the stack is
 115   /// aligned to this value. We need to re-align the unsafe stack if the
 116   /// alignment of any object on the stack exceeds this value.
 117   ///
 118   /// 16 seems like a reasonable upper bound on the alignment of objects that we
 119   /// might expect to appear on the stack on most common targets.
 120   enum { StackAlignment = 16 };
 121
 122   /// \brief Build a value representing a pointer to the unsafe stack pointer.
 123   Value *getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F);
 124
 125   /// \brief Find all static allocas, dynamic allocas, return instructions and
 126   /// stack restore points (exception unwind blocks and setjmp calls) in the
 127   /// given function and append them to the respective vectors.
 128   void findInsts(Function &F, SmallVectorImpl<AllocaInst *> &StaticAllocas,
 129                  SmallVectorImpl<AllocaInst *> &DynamicAllocas,
 130                  SmallVectorImpl<Argument *> &ByValArguments,
 131                  SmallVectorImpl<ReturnInst *> &Returns,
 132                  SmallVectorImpl<Instruction *> &StackRestorePoints);
 133
 134   /// \brief Calculate the allocation size of a given alloca. Returns 0 if the
 135   /// size can not be statically determined.
 136   uint64_t getStaticAllocaAllocationSize(const AllocaInst* AI);
 137
 138   /// \brief Allocate space for all static allocas in \p StaticAllocas,
 139   /// replace allocas with pointers into the unsafe stack and generate code to
 140   /// restore the stack pointer before all return instructions in \p Returns.
 141   ///
 142   /// \returns A pointer to the top of the unsafe stack after all unsafe static
 143   /// allocas are allocated.
 144   Value *moveStaticAllocasToUnsafeStack(IRBuilder<> &IRB, Function &F,
 145                                         ArrayRef<AllocaInst *> StaticAllocas,
 146                                         ArrayRef<Argument *> ByValArguments,
 147                                         ArrayRef<ReturnInst *> Returns);
 148
 149   /// \brief Generate code to restore the stack after all stack restore points
 150   /// in \p StackRestorePoints.
 151   ///
 152   /// \returns A local variable in which to maintain the dynamic top of the
 153   /// unsafe stack if needed.
 154   AllocaInst *
 155   createStackRestorePoints(IRBuilder<> &IRB, Function &F,
 156                            ArrayRef<Instruction *> StackRestorePoints,
 157                            Value *StaticTop, bool NeedDynamicTop);
 158
 159   /// \brief Replace all allocas in \p DynamicAllocas with code to allocate
 160   /// space dynamically on the unsafe stack and store the dynamic unsafe stack
 161   /// top to \p DynamicTop if non-null.
 162   void moveDynamicAllocasToUnsafeStack(Function &F, Value *UnsafeStackPtr,
 163                                        AllocaInst *DynamicTop,
 164                                        ArrayRef<AllocaInst *> DynamicAllocas);
 165
 166   bool IsSafeStackAlloca(const Value *AllocaPtr, uint64_t AllocaSize);
 167
 168   bool IsMemIntrinsicSafe(const MemIntrinsic *MI, const Use &U,
 169                           const Value *AllocaPtr, uint64_t AllocaSize);
 170   bool IsAccessSafe(Value *Addr, uint64_t Size, const Value *AllocaPtr,
 171                     uint64_t AllocaSize);
 172
 173 public:
 174   static char ID; // Pass identification, replacement for typeid.
 175   SafeStack(const TargetMachine *TM)
 176       : FunctionPass(ID), TM(TM), TL(nullptr), DL(nullptr) {
 177     initializeSafeStackPass(*PassRegistry::getPassRegistry());
 178   }
 179   SafeStack() : SafeStack(nullptr) {}
 180
 181   void getAnalysisUsage(AnalysisUsage &AU) const override {
 182     AU.addRequired<ScalarEvolutionWrapperPass>();
 183   }
 184
 185   bool doInitialization(Module &M) override {
 186     DL = &M.getDataLayout();
 187
 188     StackPtrTy = Type::getInt8PtrTy(M.getContext());
 189     IntPtrTy = DL->getIntPtrType(M.getContext());
 190     Int32Ty = Type::getInt32Ty(M.getContext());
 191     Int8Ty = Type::getInt8Ty(M.getContext());
 192
 193     return false;
 194   }
 195
 196   bool runOnFunction(Function &F) override;
 197 }; // class SafeStack
 198
 199 uint64_t SafeStack::getStaticAllocaAllocationSize(const AllocaInst* AI) {
 200   uint64_t Size = DL->getTypeAllocSize(AI->getAllocatedType());
 201   if (AI->isArrayAllocation()) {
 202     auto C = dyn_cast<ConstantInt>(AI->getArraySize());
 203     if (!C)
 204       return 0;
 205     Size *= C->getZExtValue();
 206   }
 207   return Size;
 208 }
 209
 210 bool SafeStack::IsAccessSafe(Value *Addr, uint64_t AccessSize,
 211                              const Value *AllocaPtr, uint64_t AllocaSize) {
 212   AllocaOffsetRewriter Rewriter(*SE, AllocaPtr);
 213   const SCEV *Expr = Rewriter.visit(SE->getSCEV(Addr));
 214
 215   uint64_t BitWidth = SE->getTypeSizeInBits(Expr->getType());
 216   ConstantRange AccessStartRange = SE->getUnsignedRange(Expr);
 217   ConstantRange SizeRange =
 218       ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, AccessSize));
 219   ConstantRange AccessRange = AccessStartRange.add(SizeRange);
 220   ConstantRange AllocaRange =
 221       ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, AllocaSize));
 222   bool Safe = AllocaRange.contains(AccessRange);
 223
 224   DEBUG(dbgs() << "[SafeStack] "
 225                << (isa<AllocaInst>(AllocaPtr) ? "Alloca " : "ByValArgument ")
 226                << *AllocaPtr << "\n"
 227                << "            Access " << *Addr << "\n"
 228                << "            SCEV " << *Expr
 229                << " U: " << SE->getUnsignedRange(Expr)
 230                << ", S: " << SE->getSignedRange(Expr) << "\n"
 231                << "            Range " << AccessRange << "\n"
 232                << "            AllocaRange " << AllocaRange << "\n"
 233                << "            " << (Safe ? "safe" : "unsafe") << "\n");
 234
 235   return Safe;
 236 }
 237
 238 bool SafeStack::IsMemIntrinsicSafe(const MemIntrinsic *MI, const Use &U,
 239                                    const Value *AllocaPtr,
 240                                    uint64_t AllocaSize) {
 241   // All MemIntrinsics have destination address in Arg0 and size in Arg2.
 242   if (MI->getRawDest() != U) return true;
 243   const auto *Len = dyn_cast<ConstantInt>(MI->getLength());
 244   // Non-constant size => unsafe. FIXME: try SCEV getRange.
 245   if (!Len) return false;
 246   return IsAccessSafe(U, Len->getZExtValue(), AllocaPtr, AllocaSize);
 247 }
 248
 249 /// Check whether a given allocation must be put on the safe
 250 /// stack or not. The function analyzes all uses of AI and checks whether it is
 251 /// only accessed in a memory safe way (as decided statically).
 252 bool SafeStack::IsSafeStackAlloca(const Value *AllocaPtr, uint64_t AllocaSize) {
 253   // Go through all uses of this alloca and check whether all accesses to the
 254   // allocated object are statically known to be memory safe and, hence, the
 255   // object can be placed on the safe stack.
 256   SmallPtrSet<const Value *, 16> Visited;
 257   SmallVector<const Value *, 8> WorkList;
 258   WorkList.push_back(AllocaPtr);
 259
 260   // A DFS search through all uses of the alloca in bitcasts/PHI/GEPs/etc.
 261   while (!WorkList.empty()) {
 262     const Value *V = WorkList.pop_back_val();
 263     for (const Use &UI : V->uses()) {
 264       auto I = cast<const Instruction>(UI.getUser());
 265       assert(V == UI.get());
 266
 267       switch (I->getOpcode()) {
 268       case Instruction::Load: {
 269         if (!IsAccessSafe(UI, DL->getTypeStoreSize(I->getType()), AllocaPtr,
 270                           AllocaSize))
 271           return false;
 272         break;
 273       }
 274       case Instruction::VAArg:
 275         // "va-arg" from a pointer is safe.
 276         break;
 277       case Instruction::Store: {
 278         if (V == I->getOperand(0)) {
 279           // Stored the pointer - conservatively assume it may be unsafe.
 280           DEBUG(dbgs() << "[SafeStack] Unsafe alloca: " << *AllocaPtr
 281                        << "\n            store of address: " << *I << "\n");
 282           return false;
 283         }
 284
 285         if (!IsAccessSafe(UI, DL->getTypeStoreSize(I->getOperand(0)->getType()),
 286                           AllocaPtr, AllocaSize))
 287           return false;
 288         break;
 289       }
 290       case Instruction::Ret: {
 291         // Information leak.
 292         return false;
 293       }
 294
 295       case Instruction::Call:
 296       case Instruction::Invoke: {
 297         ImmutableCallSite CS(I);
 298
 299         if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
 300           if (II->getIntrinsicID() == Intrinsic::lifetime_start ||
 301               II->getIntrinsicID() == Intrinsic::lifetime_end)
 302             continue;
 303         }
 304
 305         if (const MemIntrinsic *MI = dyn_cast<MemIntrinsic>(I)) {
 306           if (!IsMemIntrinsicSafe(MI, UI, AllocaPtr, AllocaSize)) {
 307             DEBUG(dbgs() << "[SafeStack] Unsafe alloca: " << *AllocaPtr
 308                          << "\n            unsafe memintrinsic: " << *I
 309                          << "\n");
 310             return false;
 311           }
 312           continue;
 313         }
 314
 315         // LLVM 'nocapture' attribute is only set for arguments whose address
 316         // is not stored, passed around, or used in any other non-trivial way.
 317         // We assume that passing a pointer to an object as a 'nocapture
 318         // readnone' argument is safe.
 319         // FIXME: a more precise solution would require an interprocedural
 320         // analysis here, which would look at all uses of an argument inside
 321         // the function being called.
 322         ImmutableCallSite::arg_iterator B = CS.arg_begin(), E = CS.arg_end();
 323         for (ImmutableCallSite::arg_iterator A = B; A != E; ++A)
 324           if (A->get() == V)
 325             if (!(CS.doesNotCapture(A - B) && (CS.doesNotAccessMemory(A - B) ||
 326                                                CS.doesNotAccessMemory()))) {
 327               DEBUG(dbgs() << "[SafeStack] Unsafe alloca: " << *AllocaPtr
 328                            << "\n            unsafe call: " << *I << "\n");
 329               return false;
 330             }
 331         continue;
 332       }
 333
 334       default:
 335         if (Visited.insert(I).second)
 336           WorkList.push_back(cast<const Instruction>(I));
 337       }
 338     }
 339   }
 340
 341   // All uses of the alloca are safe, we can place it on the safe stack.
 342   return true;
 343 }
 344
 345 Value *SafeStack::getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F) {
 346   // Check if there is a target-specific location for the unsafe stack pointer.
 347   if (TL)
 348     if (Value *V = TL->getSafeStackPointerLocation(IRB))
 349       return V;
 350
 351   // Otherwise, assume the target links with compiler-rt, which provides a
 352   // thread-local variable with a magic name.
 353   Module &M = *F.getParent();
 354   const char *UnsafeStackPtrVar = "__safestack_unsafe_stack_ptr";
 355   auto UnsafeStackPtr =
 356       dyn_cast_or_null<GlobalVariable>(M.getNamedValue(UnsafeStackPtrVar));
 357
 358   bool UseTLS = USPStorage == ThreadLocalUSP;
 359
 360   if (!UnsafeStackPtr) {
 361     auto TLSModel = UseTLS ?
 362         GlobalValue::InitialExecTLSModel :
 363         GlobalValue::NotThreadLocal;
 364     // The global variable is not defined yet, define it ourselves.
 365     // We use the initial-exec TLS model because we do not support the
 366     // variable living anywhere other than in the main executable.
 367     UnsafeStackPtr = new GlobalVariable(
 368         M, StackPtrTy, false, GlobalValue::ExternalLinkage, nullptr,
 369         UnsafeStackPtrVar, nullptr, TLSModel);
 370   } else {
 371     // The variable exists, check its type and attributes.
 372     if (UnsafeStackPtr->getValueType() != StackPtrTy)
 373       report_fatal_error(Twine(UnsafeStackPtrVar) + " must have void* type");
 374     if (UseTLS != UnsafeStackPtr->isThreadLocal())
 375       report_fatal_error(Twine(UnsafeStackPtrVar) + " must " +
 376                          (UseTLS ? "" : "not ") + "be thread-local");
 377   }
 378   return UnsafeStackPtr;
 379 }
 380
 381 void SafeStack::findInsts(Function &F,
 382                           SmallVectorImpl<AllocaInst *> &StaticAllocas,
 383                           SmallVectorImpl<AllocaInst *> &DynamicAllocas,
 384                           SmallVectorImpl<Argument *> &ByValArguments,
 385                           SmallVectorImpl<ReturnInst *> &Returns,
 386                           SmallVectorImpl<Instruction *> &StackRestorePoints) {
 387   for (Instruction &I : instructions(&F)) {
 388     if (auto AI = dyn_cast<AllocaInst>(&I)) {
 389       ++NumAllocas;
 390
 391       uint64_t Size = getStaticAllocaAllocationSize(AI);
 392       if (IsSafeStackAlloca(AI, Size))
 393         continue;
 394
 395       if (AI->isStaticAlloca()) {
 396         ++NumUnsafeStaticAllocas;
 397         StaticAllocas.push_back(AI);
 398       } else {
 399         ++NumUnsafeDynamicAllocas;
 400         DynamicAllocas.push_back(AI);
 401       }
 402     } else if (auto RI = dyn_cast<ReturnInst>(&I)) {
 403       Returns.push_back(RI);
 404     } else if (auto CI = dyn_cast<CallInst>(&I)) {
 405       // setjmps require stack restore.
 406       if (CI->getCalledFunction() && CI->canReturnTwice())
 407         StackRestorePoints.push_back(CI);
 408     } else if (auto LP = dyn_cast<LandingPadInst>(&I)) {
 409       // Exception landing pads require stack restore.
 410       StackRestorePoints.push_back(LP);
 411     } else if (auto II = dyn_cast<IntrinsicInst>(&I)) {
 412       if (II->getIntrinsicID() == Intrinsic::gcroot)
 413         llvm::report_fatal_error(
 414             "gcroot intrinsic not compatible with safestack attribute");
 415     }
 416   }
 417   for (Argument &Arg : F.args()) {
 418     if (!Arg.hasByValAttr())
 419       continue;
 420     uint64_t Size =
 421         DL->getTypeStoreSize(Arg.getType()->getPointerElementType());
 422     if (IsSafeStackAlloca(&Arg, Size))
 423       continue;
 424
 425     ++NumUnsafeByValArguments;
 426     ByValArguments.push_back(&Arg);
 427   }
 428 }
 429
 430 AllocaInst *
 431 SafeStack::createStackRestorePoints(IRBuilder<> &IRB, Function &F,
 432                                     ArrayRef<Instruction *> StackRestorePoints,
 433                                     Value *StaticTop, bool NeedDynamicTop) {
 434   if (StackRestorePoints.empty())
 435     return nullptr;
 436
 437   // We need the current value of the shadow stack pointer to restore
 438   // after longjmp or exception catching.
 439
 440   // FIXME: On some platforms this could be handled by the longjmp/exception
 441   // runtime itself.
 442
 443   AllocaInst *DynamicTop = nullptr;
 444   if (NeedDynamicTop)
 445     // If we also have dynamic alloca's, the stack pointer value changes
 446     // throughout the function. For now we store it in an alloca.
 447     DynamicTop = IRB.CreateAlloca(StackPtrTy, /*ArraySize=*/nullptr,
 448                                   "unsafe_stack_dynamic_ptr");
 449
 450   if (!StaticTop)
 451     // We need the original unsafe stack pointer value, even if there are
 452     // no unsafe static allocas.
 453     StaticTop = IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
 454
 455   if (NeedDynamicTop)
 456     IRB.CreateStore(StaticTop, DynamicTop);
 457
 458   // Restore current stack pointer after longjmp/exception catch.
 459   for (Instruction *I : StackRestorePoints) {
 460     ++NumUnsafeStackRestorePoints;
 461
 462     IRB.SetInsertPoint(I->getNextNode());
 463     Value *CurrentTop = DynamicTop ? IRB.CreateLoad(DynamicTop) : StaticTop;
 464     IRB.CreateStore(CurrentTop, UnsafeStackPtr);
 465   }
 466
 467   return DynamicTop;
 468 }
 469
 470 Value *SafeStack::moveStaticAllocasToUnsafeStack(
 471     IRBuilder<> &IRB, Function &F, ArrayRef<AllocaInst *> StaticAllocas,
 472     ArrayRef<Argument *> ByValArguments, ArrayRef<ReturnInst *> Returns) {
 473   if (StaticAllocas.empty() && ByValArguments.empty())
 474     return nullptr;
 475
 476   DIBuilder DIB(*F.getParent());
 477
 478   // We explicitly compute and set the unsafe stack layout for all unsafe
 479   // static alloca instructions. We save the unsafe "base pointer" in the
 480   // prologue into a local variable and restore it in the epilogue.
 481
 482   // Load the current stack pointer (we'll also use it as a base pointer).
 483   // FIXME: use a dedicated register for it ?
 484   Instruction *BasePointer =
 485       IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
 486   assert(BasePointer->getType() == StackPtrTy);
 487
 488   for (ReturnInst *RI : Returns) {
 489     IRB.SetInsertPoint(RI);
 490     IRB.CreateStore(BasePointer, UnsafeStackPtr);
 491   }
 492
 493   // Compute maximum alignment among static objects on the unsafe stack.
 494   unsigned MaxAlignment = 0;
 495   for (Argument *Arg : ByValArguments) {
 496     Type *Ty = Arg->getType()->getPointerElementType();
 497     unsigned Align = std::max((unsigned)DL->getPrefTypeAlignment(Ty),
 498                               Arg->getParamAlignment());
 499     if (Align > MaxAlignment)
 500       MaxAlignment = Align;
 501   }
 502   for (AllocaInst *AI : StaticAllocas) {
 503     Type *Ty = AI->getAllocatedType();
 504     unsigned Align =
 505         std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI->getAlignment());
 506     if (Align > MaxAlignment)
 507       MaxAlignment = Align;
 508   }
 509
 510   if (MaxAlignment > StackAlignment) {
 511     // Re-align the base pointer according to the max requested alignment.
 512     assert(isPowerOf2_32(MaxAlignment));
 513     IRB.SetInsertPoint(BasePointer->getNextNode());
 514     BasePointer = cast<Instruction>(IRB.CreateIntToPtr(
 515         IRB.CreateAnd(IRB.CreatePtrToInt(BasePointer, IntPtrTy),
 516                       ConstantInt::get(IntPtrTy, ~uint64_t(MaxAlignment - 1))),
 517         StackPtrTy));
 518   }
 519
 520   int64_t StaticOffset = 0; // Current stack top.
 521   IRB.SetInsertPoint(BasePointer->getNextNode());
 522
 523   for (Argument *Arg : ByValArguments) {
 524     Type *Ty = Arg->getType()->getPointerElementType();
 525
 526     uint64_t Size = DL->getTypeStoreSize(Ty);
 527     if (Size == 0)
 528       Size = 1; // Don't create zero-sized stack objects.
 529
 530     // Ensure the object is properly aligned.
 531     unsigned Align = std::max((unsigned)DL->getPrefTypeAlignment(Ty),
 532                               Arg->getParamAlignment());
 533
 534     // Add alignment.
 535     // NOTE: we ensure that BasePointer itself is aligned to >= Align.
 536     StaticOffset += Size;
 537     StaticOffset = RoundUpToAlignment(StaticOffset, Align);
 538
 539     Value *Off = IRB.CreateGEP(BasePointer, // BasePointer is i8*
 540                                ConstantInt::get(Int32Ty, -StaticOffset));
 541     Value *NewArg = IRB.CreateBitCast(Off, Arg->getType(),
 542                                      Arg->getName() + ".unsafe-byval");
 543
 544     // Replace alloc with the new location.
 545     replaceDbgDeclare(Arg, BasePointer, BasePointer->getNextNode(), DIB,
 546                       /*Deref=*/true, -StaticOffset);
 547     Arg->replaceAllUsesWith(NewArg);
 548     IRB.SetInsertPoint(cast<Instruction>(NewArg)->getNextNode());
 549     IRB.CreateMemCpy(Off, Arg, Size, Arg->getParamAlignment());
 550   }
 551
 552   // Allocate space for every unsafe static AllocaInst on the unsafe stack.
 553   for (AllocaInst *AI : StaticAllocas) {
 554     IRB.SetInsertPoint(AI);
 555
 556     Type *Ty = AI->getAllocatedType();
 557     uint64_t Size = getStaticAllocaAllocationSize(AI);
 558     if (Size == 0)
 559       Size = 1; // Don't create zero-sized stack objects.
 560
 561     // Ensure the object is properly aligned.
 562     unsigned Align =
 563         std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI->getAlignment());
 564
 565     // Add alignment.
 566     // NOTE: we ensure that BasePointer itself is aligned to >= Align.
 567     StaticOffset += Size;
 568     StaticOffset = RoundUpToAlignment(StaticOffset, Align);
 569
 570     Value *Off = IRB.CreateGEP(BasePointer, // BasePointer is i8*
 571                                ConstantInt::get(Int32Ty, -StaticOffset));
 572     Value *NewAI = IRB.CreateBitCast(Off, AI->getType(), AI->getName());
 573     if (AI->hasName() && isa<Instruction>(NewAI))
 574       cast<Instruction>(NewAI)->takeName(AI);
 575
 576     // Replace alloc with the new location.
 577     replaceDbgDeclareForAlloca(AI, BasePointer, DIB, /*Deref=*/true, -StaticOffset);
 578     AI->replaceAllUsesWith(NewAI);
 579     AI->eraseFromParent();
 580   }
 581
 582   // Re-align BasePointer so that our callees would see it aligned as
 583   // expected.
 584   // FIXME: no need to update BasePointer in leaf functions.
 585   StaticOffset = RoundUpToAlignment(StaticOffset, StackAlignment);
 586
 587   // Update shadow stack pointer in the function epilogue.
 588   IRB.SetInsertPoint(BasePointer->getNextNode());
 589
 590   Value *StaticTop =
 591       IRB.CreateGEP(BasePointer, ConstantInt::get(Int32Ty, -StaticOffset),
 592                     "unsafe_stack_static_top");
 593   IRB.CreateStore(StaticTop, UnsafeStackPtr);
 594   return StaticTop;
 595 }
 596
 597 void SafeStack::moveDynamicAllocasToUnsafeStack(
 598     Function &F, Value *UnsafeStackPtr, AllocaInst *DynamicTop,
 599     ArrayRef<AllocaInst *> DynamicAllocas) {
 600   DIBuilder DIB(*F.getParent());
 601
 602   for (AllocaInst *AI : DynamicAllocas) {
 603     IRBuilder<> IRB(AI);
 604
 605     // Compute the new SP value (after AI).
 606     Value *ArraySize = AI->getArraySize();
 607     if (ArraySize->getType() != IntPtrTy)
 608       ArraySize = IRB.CreateIntCast(ArraySize, IntPtrTy, false);
 609
 610     Type *Ty = AI->getAllocatedType();
 611     uint64_t TySize = DL->getTypeAllocSize(Ty);
 612     Value *Size = IRB.CreateMul(ArraySize, ConstantInt::get(IntPtrTy, TySize));
 613
 614     Value *SP = IRB.CreatePtrToInt(IRB.CreateLoad(UnsafeStackPtr), IntPtrTy);
 615     SP = IRB.CreateSub(SP, Size);
 616
 617     // Align the SP value to satisfy the AllocaInst, type and stack alignments.
 618     unsigned Align = std::max(
 619         std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI->getAlignment()),
 620         (unsigned)StackAlignment);
 621
 622     assert(isPowerOf2_32(Align));
 623     Value *NewTop = IRB.CreateIntToPtr(
 624         IRB.CreateAnd(SP, ConstantInt::get(IntPtrTy, ~uint64_t(Align - 1))),
 625         StackPtrTy);
 626
 627     // Save the stack pointer.
 628     IRB.CreateStore(NewTop, UnsafeStackPtr);
 629     if (DynamicTop)
 630       IRB.CreateStore(NewTop, DynamicTop);
 631
 632     Value *NewAI = IRB.CreatePointerCast(NewTop, AI->getType());
 633     if (AI->hasName() && isa<Instruction>(NewAI))
 634       NewAI->takeName(AI);
 635
 636     replaceDbgDeclareForAlloca(AI, NewAI, DIB, /*Deref=*/true);
 637     AI->replaceAllUsesWith(NewAI);
 638     AI->eraseFromParent();
 639   }
 640
 641   if (!DynamicAllocas.empty()) {
 642     // Now go through the instructions again, replacing stacksave/stackrestore.
 643     for (inst_iterator It = inst_begin(&F), Ie = inst_end(&F); It != Ie;) {
 644       Instruction *I = &*(It++);
 645       auto II = dyn_cast<IntrinsicInst>(I);
 646       if (!II)
 647         continue;
 648
 649       if (II->getIntrinsicID() == Intrinsic::stacksave) {
 650         IRBuilder<> IRB(II);
 651         Instruction *LI = IRB.CreateLoad(UnsafeStackPtr);
 652         LI->takeName(II);
 653         II->replaceAllUsesWith(LI);
 654         II->eraseFromParent();
 655       } else if (II->getIntrinsicID() == Intrinsic::stackrestore) {
 656         IRBuilder<> IRB(II);
 657         Instruction *SI = IRB.CreateStore(II->getArgOperand(0), UnsafeStackPtr);
 658         SI->takeName(II);
 659         assert(II->use_empty());
 660         II->eraseFromParent();
 661       }
 662     }
 663   }
 664 }
 665
 666 bool SafeStack::runOnFunction(Function &F) {
 667   DEBUG(dbgs() << "[SafeStack] Function: " << F.getName() << "\n");
 668
 669   if (!F.hasFnAttribute(Attribute::SafeStack)) {
 670     DEBUG(dbgs() << "[SafeStack]     safestack is not requested"
 671                     " for this function\n");
 672     return false;
 673   }
 674
 675   if (F.isDeclaration()) {
 676     DEBUG(dbgs() << "[SafeStack]     function definition"
 677                     " is not available\n");
 678     return false;
 679   }
 680
 681   TL = TM ? TM->getSubtargetImpl(F)->getTargetLowering() : nullptr;
 682   SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
 683
 684   {
 685     // Make sure the regular stack protector won't run on this function
 686     // (safestack attribute takes precedence).
 687     AttrBuilder B;
 688     B.addAttribute(Attribute::StackProtect)
 689         .addAttribute(Attribute::StackProtectReq)
 690         .addAttribute(Attribute::StackProtectStrong);
 691     F.removeAttributes(
 692         AttributeSet::FunctionIndex,
 693         AttributeSet::get(F.getContext(), AttributeSet::FunctionIndex, B));
 694   }
 695
 696   ++NumFunctions;
 697
 698   SmallVector<AllocaInst *, 16> StaticAllocas;
 699   SmallVector<AllocaInst *, 4> DynamicAllocas;
 700   SmallVector<Argument *, 4> ByValArguments;
 701   SmallVector<ReturnInst *, 4> Returns;
 702
 703   // Collect all points where stack gets unwound and needs to be restored
 704   // This is only necessary because the runtime (setjmp and unwind code) is
 705   // not aware of the unsafe stack and won't unwind/restore it prorerly.
 706   // To work around this problem without changing the runtime, we insert
 707   // instrumentation to restore the unsafe stack pointer when necessary.
 708   SmallVector<Instruction *, 4> StackRestorePoints;
 709
 710   // Find all static and dynamic alloca instructions that must be moved to the
 711   // unsafe stack, all return instructions and stack restore points.
 712   findInsts(F, StaticAllocas, DynamicAllocas, ByValArguments, Returns,
 713             StackRestorePoints);
 714
 715   if (StaticAllocas.empty() && DynamicAllocas.empty() &&
 716       ByValArguments.empty() && StackRestorePoints.empty())
 717     return false; // Nothing to do in this function.
 718
 719   if (!StaticAllocas.empty() || !DynamicAllocas.empty() ||
 720       !ByValArguments.empty())
 721     ++NumUnsafeStackFunctions; // This function has the unsafe stack.
 722
 723   if (!StackRestorePoints.empty())
 724     ++NumUnsafeStackRestorePointsFunctions;
 725
 726   IRBuilder<> IRB(&F.front(), F.begin()->getFirstInsertionPt());
 727   UnsafeStackPtr = getOrCreateUnsafeStackPtr(IRB, F);
 728
 729   // The top of the unsafe stack after all unsafe static allocas are allocated.
 730   Value *StaticTop = moveStaticAllocasToUnsafeStack(IRB, F, StaticAllocas,
 731                                                     ByValArguments, Returns);
 732
 733   // Safe stack object that stores the current unsafe stack top. It is updated
 734   // as unsafe dynamic (non-constant-sized) allocas are allocated and freed.
 735   // This is only needed if we need to restore stack pointer after longjmp
 736   // or exceptions, and we have dynamic allocations.
 737   // FIXME: a better alternative might be to store the unsafe stack pointer
 738   // before setjmp / invoke instructions.
 739   AllocaInst *DynamicTop = createStackRestorePoints(
 740       IRB, F, StackRestorePoints, StaticTop, !DynamicAllocas.empty());
 741
 742   // Handle dynamic allocas.
 743   moveDynamicAllocasToUnsafeStack(F, UnsafeStackPtr, DynamicTop,
 744                                   DynamicAllocas);
 745
 746   DEBUG(dbgs() << "[SafeStack]     safestack applied\n");
 747   return true;
 748 }
 749
 750 } // anonymous namespace
 751
 752 char SafeStack::ID = 0;
 753 INITIALIZE_TM_PASS_BEGIN(SafeStack, "safe-stack",
 754                          "Safe Stack instrumentation pass", false, false)
 755 INITIALIZE_TM_PASS_END(SafeStack, "safe-stack",
 756                        "Safe Stack instrumentation pass", false, false)
 757
 758 FunctionPass *llvm::createSafeStackPass(const llvm::TargetMachine *TM) {
 759   return new SafeStack(TM);
 760 }