[safestack] Support safestack_call_for_usp attribute

[mlemay-intel]

Summary

This patch adds support for the safestack_call_for_usp function attribute. This
attribute forces invocation of a __safestack_pointer_address function to locate
the unsafe stack pointer. Incidentally, this is already used by default for
non-X86 Android targets. The original intent of this patch was to enable
runtime libraries to use a simple, single-threaded USP during initialization and
then to switch over to using a thread-local USP at some point. It may be
necessary for some functions to be invoked both before and after that point, so
they need to support both modes of operation.

[mlemay-intel]

Changed to rely on function attribute rather than -mllvm -safe-stack-usp-init flag.

[mlemay-intel]

Rebase.

[mlemay-intel]

Rebased.

[mehdi_amini]

LangRef?

[mlemay-intel]

Updated LangRef.

[eugenis]

Do you really need all this complexity, and even a magic function attribute (which is btw entirely safestack-specific, so should have "safestack" in the name)?

We've had the same problem on Android, and were able to solve it entirely on the libc side:
https://android-review.googlesource.com/#/c/170988/23/linker/linker_main.cpp

Alternatively, consider moving this code into a runtime library function. Then SafeStack instrumentation would just call a function in the prologue instead of reading from a global. This would be good for code size.

[mlemay-intel]

Changed name of attribute to safestack_call_for_usp.
Invoke a __safestack_pointer_address function instead of generating inline code to select and access multiple USPs.

[mlemay-intel]

In D19761#575717, @eugenis wrote:

Do you really need all this complexity, and even a magic function attribute (which is btw entirely safestack-specific, so should have "safestack" in the name)?

You're correct that there are simpler ways to handle runtime initialization. I have been able to modify musl libc to successfully initialize based on this revised patch. I also no longer need D19853. I use -force-attribute instead.

We've had the same problem on Android, and were able to solve it entirely on the libc side:
https://android-review.googlesource.com/#/c/170988/23/linker/linker_main.cpp

I'll try to move initialization of a basic thread control block containing the USP early enough in musl libc to avoid the need for this patch. I'll report back after I've had a chance to try that.

Alternatively, consider moving this code into a runtime library function. Then SafeStack instrumentation would just call a function in the prologue instead of reading from a global. This would be good for code size.

This revised patch uses the existing code that calls out to __safestack_pointer_address.

[mlemay-intel]

I was able to setup a temporary thread control block early enough in musl libc initialization to obviate the need for this patch.

docs/LangRef.rst

``argmemonly``
     arguments.
     Note that ``argmemonly`` can be used together with ``readonly`` attribute
     in order to specify that function reads only from its arguments.
 ``returns_twice``
     This attribute indicates that this function can return twice. The C
     ``setjmp`` is an example of such a function. The compiler disables
     some optimizations (like tail calls) in the caller of these
     functions.
+``runtime_init``
+    This attribute indicates that this function may be used during
+    initialization of a dynamic linker or language runtime. The `SafeStack
+    <http://clang.llvm.org/docs/SafeStack.html>`_ pass compiles such functions
+    to be usable both before and after the default (e.g. thread-local) unsafe
+    stack pointer is configured, assuming that appropriate runtime support is
+    available.
 ``safestack``
     This attribute indicates that
     `SafeStack <http://clang.llvm.org/docs/SafeStack.html>`_
     protection is enabled for this function.
 
     If a function that has a ``safestack`` attribute is inlined into a
     function that doesn't have a ``safestack`` attribute or which has an
     ``ssp``, ``sspstrong`` or ``sspreq`` attribute, then the resulting

include/llvm/Bitcode/LLVMBitCodes.h

enum AttributeKindCodes {
   ATTR_KIND_SAFESTACK = 44,
   ATTR_KIND_ARGMEMONLY = 45,
   ATTR_KIND_SWIFT_SELF = 46,
   ATTR_KIND_SWIFT_ERROR = 47,
   ATTR_KIND_NO_RECURSE = 48,
   ATTR_KIND_INACCESSIBLEMEM_ONLY = 49,
   ATTR_KIND_INACCESSIBLEMEM_OR_ARGMEMONLY = 50,
   ATTR_KIND_ALLOC_SIZE = 51,
-  ATTR_KIND_WRITEONLY = 52
+  ATTR_KIND_WRITEONLY = 52,
+  ATTR_KIND_RUNTIME_INIT = 53
 };
 
 enum ComdatSelectionKindCodes {
   COMDAT_SELECTION_KIND_ANY = 1,
   COMDAT_SELECTION_KIND_EXACT_MATCH = 2,
   COMDAT_SELECTION_KIND_LARGEST = 3,
   COMDAT_SELECTION_KIND_NO_DUPLICATES = 4,
   COMDAT_SELECTION_KIND_SAME_SIZE = 5,
 };
 
 } // End bitc namespace
 } // End llvm namespace
 
 #endif

include/llvm/IR/Attributes.td

 def ReadOnly : EnumAttr<"readonly">;
 
 /// Return value is always equal to this argument.
 def Returned : EnumAttr<"returned">;
 
 /// Function can return twice.
 def ReturnsTwice : EnumAttr<"returns_twice">;
 
+/// Function may be used during initialization of dynamic linker or language
+/// runtime.  SafeStack will compile function to be usable both before and after
+/// the default (e.g. thread-local) unsafe stack pointer is configured, assuming
+/// that appropriate runtime support is available.
+def RuntimeInit : EnumAttr<"runtime_init">;
+
 /// Safe Stack protection.
 def SafeStack : EnumAttr<"safestack">;
 
 /// Sign extended before/after call.
 def SExt : EnumAttr<"signext">;
 
 /// Alignment of stack for function (3 bits)  stored as log2 of alignment with
 /// +1 bias 0 means unaligned (different from alignstack=(1)).

include/llvm/Target/TargetLowering.h

public:
 
   /// If the target has a standard stack protection check function that
   /// performs validation and error handling, returns the function. Otherwise,
   /// returns nullptr. Must be previously inserted by insertSSPDeclarations.
   /// Should be used only when getIRStackGuard returns nullptr.
   virtual Value *getSSPStackGuardCheck(const Module &M) const;
 
 protected:
+  GlobalVariable *getOrCreateStackPtrVar(Module &M, const char *Name,
+                                         GlobalValue::ThreadLocalMode TlsModel) const;
+
   Value *getDefaultSafeStackPointerLocation(IRBuilder<> &IRB,
                                             bool UseTLS) const;
 
 public:
   /// Returns the target-specific address of the unsafe stack pointer.
   virtual Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const;
 
+  /// Returns the address of the unsafe stack pointer used during runtime
+  /// initialization.
+  Value *getSafeStackInitPointerLocation(IRBuilder<> &IRB) const;
+
   /// Returns true if a cast between SrcAS and DestAS is a noop.
   virtual bool isNoopAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const {
     return false;
   }
 
   /// Return true if the pointer arguments to CI should be aligned by aligning
   /// the object whose address is being passed. If so then MinSize is set to the
   /// minimum size the object must be to be aligned and PrefAlign is set to the

lib/AsmParser/LLLexer.cpp

#define KEYWORD(STR)                                                           \
   KEYWORD(noreturn);
   KEYWORD(nounwind);
   KEYWORD(optnone);
   KEYWORD(optsize);
   KEYWORD(readnone);
   KEYWORD(readonly);
   KEYWORD(returned);
   KEYWORD(returns_twice);
+  KEYWORD(runtime_init);
   KEYWORD(signext);
   KEYWORD(sret);
   KEYWORD(ssp);
   KEYWORD(sspreq);
   KEYWORD(sspstrong);
   KEYWORD(safestack);
   KEYWORD(sanitize_address);
   KEYWORD(sanitize_thread);

lib/AsmParser/LLParser.cpp

while (true) {
     case lltok::kw_norecurse: B.addAttribute(Attribute::NoRecurse); break;
     case lltok::kw_nounwind: B.addAttribute(Attribute::NoUnwind); break;
     case lltok::kw_optnone: B.addAttribute(Attribute::OptimizeNone); break;
     case lltok::kw_optsize: B.addAttribute(Attribute::OptimizeForSize); break;
     case lltok::kw_readnone: B.addAttribute(Attribute::ReadNone); break;
     case lltok::kw_readonly: B.addAttribute(Attribute::ReadOnly); break;
     case lltok::kw_returns_twice:
       B.addAttribute(Attribute::ReturnsTwice); break;
+    case lltok::kw_runtime_init: B.addAttribute(Attribute::RuntimeInit); break;
     case lltok::kw_ssp: B.addAttribute(Attribute::StackProtect); break;
     case lltok::kw_sspreq: B.addAttribute(Attribute::StackProtectReq); break;
     case lltok::kw_sspstrong:
       B.addAttribute(Attribute::StackProtectStrong); break;
     case lltok::kw_safestack: B.addAttribute(Attribute::SafeStack); break;
     case lltok::kw_sanitize_address:
       B.addAttribute(Attribute::SanitizeAddress); break;
     case lltok::kw_sanitize_thread:
while (true) {
     case lltok::kw_noinline:
     case lltok::kw_nonlazybind:
     case lltok::kw_noredzone:
     case lltok::kw_noreturn:
     case lltok::kw_nounwind:
     case lltok::kw_optnone:
     case lltok::kw_optsize:
     case lltok::kw_returns_twice:
+    case lltok::kw_runtime_init:
     case lltok::kw_sanitize_address:
     case lltok::kw_sanitize_memory:
     case lltok::kw_sanitize_thread:
     case lltok::kw_ssp:
     case lltok::kw_sspreq:
     case lltok::kw_sspstrong:
     case lltok::kw_safestack:
     case lltok::kw_uwtable:
while (true) {
     case lltok::kw_noinline:
     case lltok::kw_nonlazybind:
     case lltok::kw_noredzone:
     case lltok::kw_noreturn:
     case lltok::kw_nounwind:
     case lltok::kw_optnone:
     case lltok::kw_optsize:
     case lltok::kw_returns_twice:
+    case lltok::kw_runtime_init:
     case lltok::kw_sanitize_address:
     case lltok::kw_sanitize_memory:
     case lltok::kw_sanitize_thread:
     case lltok::kw_ssp:
     case lltok::kw_sspreq:
     case lltok::kw_sspstrong:
     case lltok::kw_safestack:
     case lltok::kw_uwtable:

lib/AsmParser/LLToken.h

enum Kind {
   kw_noreturn,
   kw_nounwind,
   kw_optnone,
   kw_optsize,
   kw_readnone,
   kw_readonly,
   kw_returned,
   kw_returns_twice,
+  kw_runtime_init,
   kw_signext,
   kw_ssp,
   kw_sspreq,
   kw_sspstrong,
   kw_safestack,
   kw_sret,
   kw_sanitize_thread,
   kw_sanitize_memory,

lib/Bitcode/Reader/BitcodeReader.cpp

static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
   case bitc::ATTR_KIND_READ_NONE:
     return Attribute::ReadNone;
   case bitc::ATTR_KIND_READ_ONLY:
     return Attribute::ReadOnly;
   case bitc::ATTR_KIND_RETURNED:
     return Attribute::Returned;
   case bitc::ATTR_KIND_RETURNS_TWICE:
     return Attribute::ReturnsTwice;
+  case bitc::ATTR_KIND_RUNTIME_INIT:
+    return Attribute::RuntimeInit;
   case bitc::ATTR_KIND_S_EXT:
     return Attribute::SExt;
   case bitc::ATTR_KIND_STACK_ALIGNMENT:
     return Attribute::StackAlignment;
   case bitc::ATTR_KIND_STACK_PROTECT:
     return Attribute::StackProtect;
   case bitc::ATTR_KIND_STACK_PROTECT_REQ:
     return Attribute::StackProtectReq;

lib/Bitcode/Writer/BitcodeWriter.cpp

static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
   case Attribute::ReadNone:
     return bitc::ATTR_KIND_READ_NONE;
   case Attribute::ReadOnly:
     return bitc::ATTR_KIND_READ_ONLY;
   case Attribute::Returned:
     return bitc::ATTR_KIND_RETURNED;
   case Attribute::ReturnsTwice:
     return bitc::ATTR_KIND_RETURNS_TWICE;
+  case Attribute::RuntimeInit:
+     return bitc::ATTR_KIND_RUNTIME_INIT;
   case Attribute::SExt:
     return bitc::ATTR_KIND_S_EXT;
   case Attribute::StackAlignment:
     return bitc::ATTR_KIND_STACK_ALIGNMENT;
   case Attribute::StackProtect:
     return bitc::ATTR_KIND_STACK_PROTECT;
   case Attribute::StackProtectReq:
     return bitc::ATTR_KIND_STACK_PROTECT_REQ;

lib/CodeGen/SafeStack.cpp

class SafeStack : public FunctionPass {
   const DataLayout *DL;
   ScalarEvolution *SE;
 
   Type *StackPtrTy;
   Type *IntPtrTy;
   Type *Int32Ty;
   Type *Int8Ty;
 
+  /// Pointer to default unsafe stack pointer
   Value *UnsafeStackPtr = nullptr;
+  /// Pointer to single-threaded unsafe stack pointer to be used during program
+  /// initialization
+  Value *UnsafeStackPtrInit = nullptr;
+  /// Load of init unsafe stack pointer at beginning of function
+  Value *UnsafeStackPtrInitLoad = nullptr;
 
   /// Unsafe stack alignment. Each stack frame must ensure that the stack is
   /// aligned to this value. We need to re-align the unsafe stack if the
   /// alignment of any object on the stack exceeds this value.
   ///
   /// 16 seems like a reasonable upper bound on the alignment of objects that we
   /// might expect to appear on the stack on most common targets.
   enum { StackAlignment = 16 };
 
+  /// \brief Load the unsafe stack pointer.
+  ///
+  /// \param Top Set to true if this load is at the beginning/top of the
+  ///            function.
+  Instruction *loadUnsafeStackPtr(IRBuilder<> &IRB, bool Top = false);
+
+  /// \brief Store the updated value of the unsafe stack pointer.
+  void storeUnsafeStackPtr(IRBuilder<> &IRB, Value *UpdatedUSP,
+                           Value *NameSource = nullptr);
+
   /// \brief Return the value of the stack canary.
   Value *getStackGuard(IRBuilder<> &IRB, Function &F);
 
   /// \brief Load stack guard from the frame and check if it has changed.
   void checkStackGuard(IRBuilder<> &IRB, Function &F, ReturnInst &RI,
                        AllocaInst *StackGuardSlot, Value *StackGuard);
 
   /// \brief Find all static allocas, dynamic allocas, return instructions and
class SafeStack : public FunctionPass {
   AllocaInst *
   createStackRestorePoints(IRBuilder<> &IRB, Function &F,
                            ArrayRef<Instruction *> StackRestorePoints,
                            Value *StaticTop, bool NeedDynamicTop);
 
   /// \brief Replace all allocas in \p DynamicAllocas with code to allocate
   /// space dynamically on the unsafe stack and store the dynamic unsafe stack
   /// top to \p DynamicTop if non-null.
-  void moveDynamicAllocasToUnsafeStack(Function &F, Value *UnsafeStackPtr,
-                                       AllocaInst *DynamicTop,
+  void moveDynamicAllocasToUnsafeStack(Function &F, AllocaInst *DynamicTop,
                                        ArrayRef<AllocaInst *> DynamicAllocas);
 
   bool IsSafeStackAlloca(const Value *AllocaPtr, uint64_t AllocaSize);
 
   bool IsMemIntrinsicSafe(const MemIntrinsic *MI, const Use &U,
                           const Value *AllocaPtr, uint64_t AllocaSize);
   bool IsAccessSafe(Value *Addr, uint64_t Size, const Value *AllocaPtr,
                     uint64_t AllocaSize);
for (const Use &UI : V->uses()) {
       }
     }
   }
 
   // All uses of the alloca are safe, we can place it on the safe stack.
   return true;
 }
 
+Instruction *SafeStack::loadUnsafeStackPtr(IRBuilder<> &IRB, bool Top) {
+  if (UnsafeStackPtrInit == nullptr)
+    return IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
+
+  auto IsInit = IRB.CreateIsNotNull(UnsafeStackPtrInitLoad);
+
+  auto PHI = IRB.CreatePHI(StackPtrTy, 2);
+  TerminatorInst *ThenTerm, *ElseTerm;
+  SplitBlockAndInsertIfThenElse(IsInit, PHI, &ThenTerm, &ElseTerm);
+  IRBuilder<> InitIRB(ThenTerm);
+  IRBuilder<> DefaultIRB(ElseTerm);
+  auto Tail = ThenTerm->getSuccessor(0);
+  IRB.SetInsertPoint(Tail, Tail->getFirstInsertionPt());
+
+  auto InitLoad = Top?
+      // There is no need to load this twice at the top of the function.
+      UnsafeStackPtrInitLoad :
+      InitIRB.CreateLoad(UnsafeStackPtrInit, false, "unsafe_stack_ptr_init");
+  auto DefaultLoad = DefaultIRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
+
+  PHI->addIncoming(InitLoad, ThenTerm->getParent());
+  PHI->addIncoming(DefaultLoad, ElseTerm->getParent());
+
+  return PHI;
+}
+
+void SafeStack::storeUnsafeStackPtr(IRBuilder<> &IRB, Value *UpdatedUSP,
+                                    Value *NameSource) {
+  if (UnsafeStackPtrInit == nullptr) {
+    auto DefaultStore = IRB.CreateStore(UpdatedUSP, UnsafeStackPtr);
+    if (NameSource != nullptr)
+      DefaultStore->takeName(NameSource);
+    return;
+  }
+
+  auto IsInit = IRB.CreateIsNotNull(UnsafeStackPtrInitLoad);
+  TerminatorInst *ThenTerm, *ElseTerm;
+  SplitBlockAndInsertIfThenElse(IsInit, &*IRB.GetInsertPoint(),
+                                &ThenTerm, &ElseTerm);
+  IRBuilder<> InitIRB(ThenTerm);
+  IRBuilder<> DefaultIRB(ElseTerm);
+  auto Tail = ThenTerm->getSuccessor(0);
+  IRB.SetInsertPoint(Tail, Tail->getFirstInsertionPt());
+
+  auto InitStore = InitIRB.CreateStore(UpdatedUSP, UnsafeStackPtrInit);
+  auto DefaultStore = DefaultIRB.CreateStore(UpdatedUSP, UnsafeStackPtr);
+
+  if (NameSource == nullptr)
+    return;
+
+  InitStore->takeName(NameSource);
+  DefaultStore->takeName(NameSource);
+}
+
 Value *SafeStack::getStackGuard(IRBuilder<> &IRB, Function &F) {
   Value *StackGuardVar = TL->getIRStackGuard(IRB);
   if (!StackGuardVar)
     StackGuardVar =
         F.getParent()->getOrInsertGlobal("__stack_chk_guard", StackPtrTy);
   return IRB.CreateLoad(StackGuardVar, "StackGuard");
 }
 
SafeStack::createStackRestorePoints(IRBuilder<> &IRB, Function &F,
   }
 
   // Restore current stack pointer after longjmp/exception catch.
   for (Instruction *I : StackRestorePoints) {
     ++NumUnsafeStackRestorePoints;
 
     IRB.SetInsertPoint(I->getNextNode());
     Value *CurrentTop = DynamicTop ? IRB.CreateLoad(DynamicTop) : StaticTop;
-    IRB.CreateStore(CurrentTop, UnsafeStackPtr);
+    storeUnsafeStackPtr(IRB, CurrentTop);
   }
 
   return DynamicTop;
 }
 
 void SafeStack::checkStackGuard(IRBuilder<> &IRB, Function &F, ReturnInst &RI,
                                 AllocaInst *StackGuardSlot, Value *StackGuard) {
   Value *V = IRB.CreateLoad(StackGuardSlot);
Value *SafeStack::moveStaticAllocasToUnsafeStack(
   unsigned FrameSize = alignTo(SSL.getFrameSize(), StackAlignment);
 
   // Update shadow stack pointer in the function epilogue.
   IRB.SetInsertPoint(BasePointer->getNextNode());
 
   Value *StaticTop =
       IRB.CreateGEP(BasePointer, ConstantInt::get(Int32Ty, -FrameSize),
                     "unsafe_stack_static_top");
-  IRB.CreateStore(StaticTop, UnsafeStackPtr);
+  storeUnsafeStackPtr(IRB, StaticTop);
   return StaticTop;
 }
 
 void SafeStack::moveDynamicAllocasToUnsafeStack(
-    Function &F, Value *UnsafeStackPtr, AllocaInst *DynamicTop,
+    Function &F, AllocaInst *DynamicTop,
     ArrayRef<AllocaInst *> DynamicAllocas) {
   DIBuilder DIB(*F.getParent());
 
   for (AllocaInst *AI : DynamicAllocas) {
     IRBuilder<> IRB(AI);
 
     // Compute the new SP value (after AI).
     Value *ArraySize = AI->getArraySize();
     if (ArraySize->getType() != IntPtrTy)
       ArraySize = IRB.CreateIntCast(ArraySize, IntPtrTy, false);
 
     Type *Ty = AI->getAllocatedType();
     uint64_t TySize = DL->getTypeAllocSize(Ty);
     Value *Size = IRB.CreateMul(ArraySize, ConstantInt::get(IntPtrTy, TySize));
 
-    Value *SP = IRB.CreatePtrToInt(IRB.CreateLoad(UnsafeStackPtr), IntPtrTy);
+    Value *SP = IRB.CreatePtrToInt(loadUnsafeStackPtr(IRB), IntPtrTy);
     SP = IRB.CreateSub(SP, Size);
 
     // Align the SP value to satisfy the AllocaInst, type and stack alignments.
     unsigned Align = std::max(
         std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI->getAlignment()),
         (unsigned)StackAlignment);
 
     assert(isPowerOf2_32(Align));
     Value *NewTop = IRB.CreateIntToPtr(
         IRB.CreateAnd(SP, ConstantInt::get(IntPtrTy, ~uint64_t(Align - 1))),
         StackPtrTy);
 
     // Save the stack pointer.
-    IRB.CreateStore(NewTop, UnsafeStackPtr);
+    storeUnsafeStackPtr(IRB, NewTop);
     if (DynamicTop)
       IRB.CreateStore(NewTop, DynamicTop);
 
     Value *NewAI = IRB.CreatePointerCast(NewTop, AI->getType());
     if (AI->hasName() && isa<Instruction>(NewAI))
       NewAI->takeName(AI);
 
     replaceDbgDeclareForAlloca(AI, NewAI, DIB, /*Deref=*/true);
     AI->replaceAllUsesWith(NewAI);
     AI->eraseFromParent();
   }
 
   if (!DynamicAllocas.empty()) {
     // Now go through the instructions again, replacing stacksave/stackrestore.
+    SmallPtrSet<IntrinsicInst *, 4> StackSaves, StackRestores;
     for (inst_iterator It = inst_begin(&F), Ie = inst_end(&F); It != Ie;) {
       Instruction *I = &*(It++);
       auto II = dyn_cast<IntrinsicInst>(I);
       if (!II)
         continue;
 
-      if (II->getIntrinsicID() == Intrinsic::stacksave) {
-        IRBuilder<> IRB(II);
-        Instruction *LI = IRB.CreateLoad(UnsafeStackPtr);
-        LI->takeName(II);
-        II->replaceAllUsesWith(LI);
-        II->eraseFromParent();
-      } else if (II->getIntrinsicID() == Intrinsic::stackrestore) {
-        IRBuilder<> IRB(II);
-        Instruction *SI = IRB.CreateStore(II->getArgOperand(0), UnsafeStackPtr);
-        SI->takeName(II);
-        assert(II->use_empty());
-        II->eraseFromParent();
-      }
-    }
-  }
-}
+      if (II->getIntrinsicID() == Intrinsic::stacksave)
+        StackSaves.insert(II);
+      else if (II->getIntrinsicID() == Intrinsic::stackrestore)
+        StackRestores.insert(II);
+    }
+
+    for (auto II : StackSaves) {
+      Instruction *LI = loadUnsafeStackPtr(IRB);
+    }
+
+    for (auto II : StackRestores) {
+      storeUnsafeStackPtr(IRB, II->getArgOperand(0), II);
 
 bool SafeStack::runOnFunction(Function &F) {
   DEBUG(dbgs() << "[SafeStack] Function: " << F.getName() << "\n");
 
   if (!F.hasFnAttribute(Attribute::SafeStack)) {
     DEBUG(dbgs() << "[SafeStack]     safestack is not requested"
                     " for this function\n");
     return false;
if (!StaticAllocas.empty() || !DynamicAllocas.empty() ||
       !ByValArguments.empty())
     ++NumUnsafeStackFunctions; // This function has the unsafe stack.
 
   if (!StackRestorePoints.empty())
     ++NumUnsafeStackRestorePointsFunctions;
 
   IRBuilder<> IRB(&F.front(), F.begin()->getFirstInsertionPt());
   UnsafeStackPtr = TL->getSafeStackPointerLocation(IRB);
+  UnsafeStackPtrInit = nullptr;
+  if (F.hasFnAttribute(Attribute::RuntimeInit)) {
+    // Initialize UnsafeStackPtrInit with a value representing a pointer to the
+    // single-threaded unsafe stack pointer to be used during program
+    // initialization.
+    UnsafeStackPtrInit = TL->getSafeStackInitPointerLocation(IRB);
+    // Initialize UnsafeStackPtrInitLoad with an instruction that loads
+    // UnsafeStackPtrInit. That loaded value is initially used to select
+    // between that init pointer and the default pointer, depending on whether
+    // the init pointer is null. That loaded value is also used to determine
+    // where to later store the updated unsafe stack pointer.
+    UnsafeStackPtrInitLoad =
+        IRB.CreateLoad(UnsafeStackPtrInit, false, "unsafe_stack_ptr_init");
+  }
 
   // Load the current stack pointer (we'll also use it as a base pointer).
   // FIXME: use a dedicated register for it ?
-  Instruction *BasePointer =
-      IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
+  Instruction *BasePointer = loadUnsafeStackPtr(IRB, true);
   assert(BasePointer->getType() == StackPtrTy);
 
   AllocaInst *StackGuardSlot = nullptr;
   // FIXME: implement weaker forms of stack protector.
   if (F.hasFnAttribute(Attribute::StackProtect) ||
       F.hasFnAttribute(Attribute::StackProtectStrong) ||
       F.hasFnAttribute(Attribute::StackProtectReq)) {
     Value *StackGuard = getStackGuard(IRB, F);
bool SafeStack::runOnFunction(Function &F) {
   // This is only needed if we need to restore stack pointer after longjmp
   // or exceptions, and we have dynamic allocations.
   // FIXME: a better alternative might be to store the unsafe stack pointer
   // before setjmp / invoke instructions.
   AllocaInst *DynamicTop = createStackRestorePoints(
       IRB, F, StackRestorePoints, StaticTop, !DynamicAllocas.empty());
 
   // Handle dynamic allocas.
-  moveDynamicAllocasToUnsafeStack(F, UnsafeStackPtr, DynamicTop,
-                                  DynamicAllocas);
+  moveDynamicAllocasToUnsafeStack(F, DynamicTop, DynamicAllocas);
 
   // Restore the unsafe stack pointer before each return.
   for (ReturnInst *RI : Returns) {
     IRB.SetInsertPoint(RI);
-    IRB.CreateStore(BasePointer, UnsafeStackPtr);
+    storeUnsafeStackPtr(IRB, BasePointer);
   }
 
   DEBUG(dbgs() << "[SafeStack]     safestack applied\n");
   return true;
 }
 
 } // anonymous namespace
 

lib/CodeGen/TargetLoweringBase.cpp

while (true) {
     if (MTy == LK.second)
       return std::make_pair(Cost, MTy.getSimpleVT());
 
     // Keep legalizing the type.
     MTy = LK.second;
   }
 }
 
+GlobalVariable *TargetLoweringBase::getOrCreateStackPtrVar(Module &M,
+                                                           const char *Name,
+                                                           GlobalValue::ThreadLocalMode TlsModel) const {
+  GlobalVariable *V = dyn_cast_or_null<GlobalVariable>(M.getNamedValue(Name));
+  Type *StackPtrTy = Type::getInt8PtrTy(M.getContext());
+  if (!V) {
+    // The global variable is not defined yet, define it ourselves.
+    V = new GlobalVariable(
+        M, StackPtrTy, false, GlobalValue::ExternalLinkage, nullptr, Name,
+        nullptr, TlsModel);
+  } else {
+    // The variable exists, check its type and attributes.
+    if (V->getValueType() != StackPtrTy)
+      report_fatal_error(Twine(Name) + " must have void* type");
+    if (TlsModel == GlobalValue::NotThreadLocal) {
+      if (V->isThreadLocal())
+        report_fatal_error(Twine(Name) + " must not be thread-local");
+    } else
+      if (!V->isThreadLocal())
+        report_fatal_error(Twine(Name) + " must be thread-local");
+  }
+
+  return V;
+}
+
 Value *TargetLoweringBase::getDefaultSafeStackPointerLocation(IRBuilder<> &IRB,
                                                               bool UseTLS) const {
   // compiler-rt provides a variable with a magic name.  Targets that do not
   // link with compiler-rt may also provide such a variable.
   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
-  const char *UnsafeStackPtrVar = "__safestack_unsafe_stack_ptr";
-  auto UnsafeStackPtr =
-      dyn_cast_or_null<GlobalVariable>(M->getNamedValue(UnsafeStackPtrVar));
-
-  Type *StackPtrTy = Type::getInt8PtrTy(M->getContext());
-
-  if (!UnsafeStackPtr) {
-    auto TLSModel = UseTLS ?
-        GlobalValue::InitialExecTLSModel :
-        GlobalValue::NotThreadLocal;
-    // The global variable is not defined yet, define it ourselves.
+  const char *UnsafeStackPtrName = "__safestack_unsafe_stack_ptr";
       // We use the initial-exec TLS model because we do not support the
       // variable living anywhere other than in the main executable.
-    UnsafeStackPtr = new GlobalVariable(
-        *M, StackPtrTy, false, GlobalValue::ExternalLinkage, nullptr,
-        UnsafeStackPtrVar, nullptr, TLSModel);
-  } else {
-    // The variable exists, check its type and attributes.
-    if (UnsafeStackPtr->getValueType() != StackPtrTy)
-      report_fatal_error(Twine(UnsafeStackPtrVar) + " must have void* type");
-    if (UseTLS != UnsafeStackPtr->isThreadLocal())
-      report_fatal_error(Twine(UnsafeStackPtrVar) + " must " +
-                         (UseTLS ? "" : "not ") + "be thread-local");
-  }
-  return UnsafeStackPtr;
+      GlobalValue::InitialExecTLSModel :
+      GlobalValue::NotThreadLocal;
+  return getOrCreateStackPtrVar(*M, UnsafeStackPtrName, TLSModel);
 }
 
 Value *TargetLoweringBase::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
   if (!TM.getTargetTriple().isAndroid())
     return getDefaultSafeStackPointerLocation(IRB, true);
 
   // Android provides a libc function to retrieve the address of the current
   // thread's unsafe stack pointer.
   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
   Type *StackPtrTy = Type::getInt8PtrTy(M->getContext());
   Value *Fn = M->getOrInsertFunction("__safestack_pointer_address",
                                      StackPtrTy->getPointerTo(0), nullptr);
   return IRB.CreateCall(Fn);
 }
 
+Value *TargetLoweringBase::getSafeStackInitPointerLocation(IRBuilder<> &IRB) const {
+  const char *UnsafeStackPtrInitName = "__safestack_unsafe_stack_ptr_init";
+  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
+  return getOrCreateStackPtrVar(
+      *M, UnsafeStackPtrInitName, GlobalValue::NotThreadLocal);
+}
+
 //===----------------------------------------------------------------------===//
 //  Loop Strength Reduction hooks
 //===----------------------------------------------------------------------===//
 
 /// isLegalAddressingMode - Return true if the addressing mode represented
 /// by AM is legal for this target, for a load/store of the specified type.
 bool TargetLoweringBase::isLegalAddressingMode(const DataLayout &DL,
                                                const AddrMode &AM, Type *Ty,

lib/IR/Attributes.cpp

std::string Attribute::getAsString(bool InAttrGrp) const {
   if (hasAttribute(Attribute::ReadOnly))
     return "readonly";
   if (hasAttribute(Attribute::WriteOnly))
     return "writeonly";
   if (hasAttribute(Attribute::Returned))
     return "returned";
   if (hasAttribute(Attribute::ReturnsTwice))
     return "returns_twice";
+  if (hasAttribute(Attribute::RuntimeInit))
+    return "runtime_init";
   if (hasAttribute(Attribute::SExt))
     return "signext";
   if (hasAttribute(Attribute::StackProtect))
     return "ssp";
   if (hasAttribute(Attribute::StackProtectReq))
     return "sspreq";
   if (hasAttribute(Attribute::StackProtectStrong))
     return "sspstrong";
uint64_t AttributeImpl::getAttrMask(Attribute::AttrKind Val) {
   case Attribute::Convergent:      return 1ULL << 46;
   case Attribute::SafeStack:       return 1ULL << 47;
   case Attribute::NoRecurse:       return 1ULL << 48;
   case Attribute::InaccessibleMemOnly:         return 1ULL << 49;
   case Attribute::InaccessibleMemOrArgMemOnly: return 1ULL << 50;
   case Attribute::SwiftSelf:       return 1ULL << 51;
   case Attribute::SwiftError:      return 1ULL << 52;
   case Attribute::WriteOnly:       return 1ULL << 53;
+  case Attribute::RuntimeInit:     return 1ULL << 54;
   case Attribute::Dereferenceable:
     llvm_unreachable("dereferenceable attribute not supported in raw format");
     break;
   case Attribute::DereferenceableOrNull:
     llvm_unreachable("dereferenceable_or_null attribute not supported in raw "
                      "format");
     break;
   case Attribute::ArgMemOnly:

lib/IR/Verifier.cpp

if (I->getKindAsEnum() == Attribute::NoReturn ||
         I->getKindAsEnum() == Attribute::NoRedZone ||
         I->getKindAsEnum() == Attribute::NoImplicitFloat ||
         I->getKindAsEnum() == Attribute::Naked ||
         I->getKindAsEnum() == Attribute::InlineHint ||
         I->getKindAsEnum() == Attribute::StackAlignment ||
         I->getKindAsEnum() == Attribute::UWTable ||
         I->getKindAsEnum() == Attribute::NonLazyBind ||
         I->getKindAsEnum() == Attribute::ReturnsTwice ||
+        I->getKindAsEnum() == Attribute::RuntimeInit ||
         I->getKindAsEnum() == Attribute::SanitizeAddress ||
         I->getKindAsEnum() == Attribute::SanitizeThread ||
         I->getKindAsEnum() == Attribute::SanitizeMemory ||
         I->getKindAsEnum() == Attribute::MinSize ||
         I->getKindAsEnum() == Attribute::NoDuplicate ||
         I->getKindAsEnum() == Attribute::Builtin ||
         I->getKindAsEnum() == Attribute::NoBuiltin ||
         I->getKindAsEnum() == Attribute::Cold ||

lib/Transforms/IPO/ForceFunctionAttrs.cpp

return StringSwitch<Attribute::AttrKind>(Kind)
       .Case("norecurse", Attribute::NoRecurse)
       .Case("nounwind", Attribute::NoUnwind)
       .Case("optnone", Attribute::OptimizeNone)
       .Case("optsize", Attribute::OptimizeForSize)
       .Case("readnone", Attribute::ReadNone)
       .Case("readonly", Attribute::ReadOnly)
       .Case("argmemonly", Attribute::ArgMemOnly)
       .Case("returns_twice", Attribute::ReturnsTwice)
+      .Case("runtime_init", Attribute::RuntimeInit)
       .Case("safestack", Attribute::SafeStack)
       .Case("sanitize_address", Attribute::SanitizeAddress)
       .Case("sanitize_memory", Attribute::SanitizeMemory)
       .Case("sanitize_thread", Attribute::SanitizeThread)
       .Case("ssp", Attribute::StackProtect)
       .Case("sspreq", Attribute::StackProtectReq)
       .Case("sspstrong", Attribute::StackProtectStrong)
       .Case("uwtable", Attribute::UWTable)

test/Transforms/SafeStack/runtime-init.ll

+; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+
+; CHECK: @__safestack_unsafe_stack_ptr = external thread_local(initialexec) global i8*
+; CHECK: @__safestack_unsafe_stack_ptr_init = external global i8*
+
+; Load from an array at a fixed offset with overflow.
+define i8 @StaticArrayFixedUnsafe() nounwind uwtable safestack runtime_init {
+entry:
+  ; CHECK-LABEL: define i8 @StaticArrayFixedUnsafe(
+  ; CHECK: %unsafe_stack_ptr_init = load i8*, i8** @__safestack_unsafe_stack_ptr_init
+  ; CHECK-NEXT: [[REG0:%[0-9]+]] = icmp ne i8* %unsafe_stack_ptr_init, null
+  ; CHECK-NEXT: br i1 [[REG0]]
+  ; CHECK: %unsafe_stack_ptr = load i8*, i8** @__safestack_unsafe_stack_ptr
+  ; CHECK: store i8* [[REG1:.*]], i8** @__safestack_unsafe_stack_ptr_init
+  ; CHECK: store i8* [[REG1]], i8** @__safestack_unsafe_stack_ptr
+
+  %buf = alloca i8, i32 4, align 1
+  %gep = getelementptr inbounds i8, i8* %buf, i32 5
+  %x = load i8, i8* %gep, align 1
+  ret i8 %x
+}