[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

``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
     function will have a ``safestack`` attribute.
+``safestack_call_for_usp``
+    This attribute causes the `SafeStack
+    <http://clang.llvm.org/docs/SafeStack.html>`_ pass to call the
+    ``__safestack_pointer_address`` function to obtain the location of the
+    unsafe stack pointer. The ``__safestack_pointer_address`` function may be
+    provided by a runtime library. This attribute may be useful for functions
+    that may be used both during initialization of a dynamic linker or language
+    runtime and later, after the default (e.g. thread-local) unsafe stack
+    pointer is configured. For example, prior to initialization of thread-local
+    storage, the ``__safestack_pointer_address`` function may return a pointer
+    to a non-thread-local variable.
 ``sanitize_address``
     This attribute indicates that AddressSanitizer checks
     (dynamic address safety analysis) are enabled for this function.
 ``sanitize_memory``
     This attribute indicates that MemorySanitizer checks (dynamic detection
     of accesses to uninitialized memory) are enabled for this function.
 ``sanitize_thread``
     This attribute indicates that ThreadSanitizer checks

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_SAFESTACK_CALL_FOR_USP = 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 Returned : EnumAttr<"returned">;
 
 /// Function can return twice.
 def ReturnsTwice : EnumAttr<"returns_twice">;
 
 /// Safe Stack protection.
 def SafeStack : EnumAttr<"safestack">;
 
+/// Call __safestack_pointer_address to locate the unsafe stack pointer.
+def SafeStackCallForUSP : EnumAttr<"safestack_call_for_usp">;
+
 /// 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)).
 def StackAlignment : EnumAttr<"alignstack">;
 
 /// Stack protection.

lib/AsmParser/LLLexer.cpp

#define KEYWORD(STR)                                                           \
   KEYWORD(returned);
   KEYWORD(returns_twice);
   KEYWORD(signext);
   KEYWORD(sret);
   KEYWORD(ssp);
   KEYWORD(sspreq);
   KEYWORD(sspstrong);
   KEYWORD(safestack);
+  KEYWORD(safestack_call_for_usp);
   KEYWORD(sanitize_address);
   KEYWORD(sanitize_thread);
   KEYWORD(sanitize_memory);
   KEYWORD(swifterror);
   KEYWORD(swiftself);
   KEYWORD(uwtable);
   KEYWORD(writeonly);
   KEYWORD(zeroext);

lib/AsmParser/LLParser.cpp

while (true) {
     case lltok::kw_readonly: B.addAttribute(Attribute::ReadOnly); break;
     case lltok::kw_returns_twice:
       B.addAttribute(Attribute::ReturnsTwice); 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_safestack_call_for_usp:
+      B.addAttribute(Attribute::SafeStackCallForUSP); break;
     case lltok::kw_sanitize_address:
       B.addAttribute(Attribute::SanitizeAddress); break;
     case lltok::kw_sanitize_thread:
       B.addAttribute(Attribute::SanitizeThread); break;
     case lltok::kw_sanitize_memory:
       B.addAttribute(Attribute::SanitizeMemory); break;
     case lltok::kw_uwtable: B.addAttribute(Attribute::UWTable); break;
     case lltok::kw_writeonly: B.addAttribute(Attribute::WriteOnly); break;
while (true) {
     case lltok::kw_optsize:
     case lltok::kw_returns_twice:
     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_call_for_usp:
     case lltok::kw_safestack:
     case lltok::kw_uwtable:
       HaveError |= Error(Lex.getLoc(), "invalid use of function-only attribute");
       break;
     }
 
     Lex.Lex();
   }
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_safestack_call_for_usp:
     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_readonly,
   kw_returned,
   kw_returns_twice,
   kw_signext,
   kw_ssp,
   kw_sspreq,
   kw_sspstrong,
   kw_safestack,
+  kw_safestack_call_for_usp,
   kw_sret,
   kw_sanitize_thread,
   kw_sanitize_memory,
   kw_swifterror,
   kw_swiftself,
   kw_uwtable,
   kw_writeonly,
   kw_zeroext,

lib/Bitcode/Reader/BitcodeReader.cpp

static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
   case bitc::ATTR_KIND_STACK_PROTECT:
     return Attribute::StackProtect;
   case bitc::ATTR_KIND_STACK_PROTECT_REQ:
     return Attribute::StackProtectReq;
   case bitc::ATTR_KIND_STACK_PROTECT_STRONG:
     return Attribute::StackProtectStrong;
   case bitc::ATTR_KIND_SAFESTACK:
     return Attribute::SafeStack;
+  case bitc::ATTR_KIND_SAFESTACK_CALL_FOR_USP:
+    return Attribute::SafeStackCallForUSP;
   case bitc::ATTR_KIND_STRUCT_RET:
     return Attribute::StructRet;
   case bitc::ATTR_KIND_SANITIZE_ADDRESS:
     return Attribute::SanitizeAddress;
   case bitc::ATTR_KIND_SANITIZE_THREAD:
     return Attribute::SanitizeThread;
   case bitc::ATTR_KIND_SANITIZE_MEMORY:
     return Attribute::SanitizeMemory;

lib/Bitcode/Writer/BitcodeWriter.cpp

static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
   case Attribute::StackProtect:
     return bitc::ATTR_KIND_STACK_PROTECT;
   case Attribute::StackProtectReq:
     return bitc::ATTR_KIND_STACK_PROTECT_REQ;
   case Attribute::StackProtectStrong:
     return bitc::ATTR_KIND_STACK_PROTECT_STRONG;
   case Attribute::SafeStack:
     return bitc::ATTR_KIND_SAFESTACK;
+  case Attribute::SafeStackCallForUSP:
+    return bitc::ATTR_KIND_SAFESTACK_CALL_FOR_USP;
   case Attribute::StructRet:
     return bitc::ATTR_KIND_STRUCT_RET;
   case Attribute::SanitizeAddress:
     return bitc::ATTR_KIND_SANITIZE_ADDRESS;
   case Attribute::SanitizeThread:
     return bitc::ATTR_KIND_SANITIZE_THREAD;
   case Attribute::SanitizeMemory:
     return bitc::ATTR_KIND_SANITIZE_MEMORY;

lib/CodeGen/TargetLoweringBase.cpp

if (!UnsafeStackPtr) {
     if (UseTLS != UnsafeStackPtr->isThreadLocal())
       report_fatal_error(Twine(UnsafeStackPtrVar) + " must " +
                          (UseTLS ? "" : "not ") + "be thread-local");
   }
   return UnsafeStackPtr;
 }
 
 Value *TargetLoweringBase::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
-  if (!TM.getTargetTriple().isAndroid())
+  Function *F = IRB.GetInsertBlock()->getParent();
+  if (!(TM.getTargetTriple().isAndroid() ||
+        F->hasFnAttribute(Attribute::SafeStackCallForUSP)))
     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();
+  // thread's unsafe stack pointer. The safestack_call_for_usp attribute can be
+  // used on non-Android targets to force usage of a similarly-named function.
+  Module *M = F->getParent();
   Type *StackPtrTy = Type::getInt8PtrTy(M->getContext());
   Value *Fn = M->getOrInsertFunction("__safestack_pointer_address",
                                      StackPtrTy->getPointerTo(0), nullptr);
   return IRB.CreateCall(Fn);
 }
 
 //===----------------------------------------------------------------------===//
 //  Loop Strength Reduction hooks

lib/IR/Attributes.cpp

std::string Attribute::getAsString(bool InAttrGrp) const {
   if (hasAttribute(Attribute::StackProtect))
     return "ssp";
   if (hasAttribute(Attribute::StackProtectReq))
     return "sspreq";
   if (hasAttribute(Attribute::StackProtectStrong))
     return "sspstrong";
   if (hasAttribute(Attribute::SafeStack))
     return "safestack";
+  if (hasAttribute(Attribute::SafeStackCallForUSP))
+    return "safestack_call_for_usp";
   if (hasAttribute(Attribute::StructRet))
     return "sret";
   if (hasAttribute(Attribute::SanitizeThread))
     return "sanitize_thread";
   if (hasAttribute(Attribute::SanitizeMemory))
     return "sanitize_memory";
   if (hasAttribute(Attribute::UWTable))
     return "uwtable";
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::SafeStackCallForUSP: 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::NoUnwind ||
         I->getKindAsEnum() == Attribute::NoInline ||
         I->getKindAsEnum() == Attribute::AlwaysInline ||
         I->getKindAsEnum() == Attribute::OptimizeForSize ||
         I->getKindAsEnum() == Attribute::StackProtect ||
         I->getKindAsEnum() == Attribute::StackProtectReq ||
         I->getKindAsEnum() == Attribute::StackProtectStrong ||
         I->getKindAsEnum() == Attribute::SafeStack ||
+        I->getKindAsEnum() == Attribute::SafeStackCallForUSP ||
         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 ||

lib/Transforms/IPO/ForceFunctionAttrs.cpp

return StringSwitch<Attribute::AttrKind>(Kind)
       .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("safestack", Attribute::SafeStack)
+      .Case("safestack_call_for_usp", Attribute::SafeStackCallForUSP)
       .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)
       .Default(Attribute::None);

test/Transforms/SafeStack/call-for-usp.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
+
+; Load from an array at a fixed offset with overflow.
+define i8 @StaticArrayFixedUnsafe() nounwind uwtable safestack safestack_call_for_usp {
+entry:
+  ; CHECK-LABEL: define i8 @StaticArrayFixedUnsafe(
+  ; CHECK: %0 = call i8** @__safestack_pointer_address()
+  ; CHECK-NEXT: %unsafe_stack_ptr = load i8*, i8** %0
+  ; CHECK-NEXT: %unsafe_stack_static_top = getelementptr i8, i8* %unsafe_stack_ptr, i32 -16
+  ; CHECK-NEXT: store i8* %unsafe_stack_static_top, i8** %0
+  ; CHECK: store i8* %unsafe_stack_ptr, i8** %0
+
+  %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
+}
+
+; CHECK: declare i8** @__safestack_pointer_address()