[safestack] Experimental mode where stack pointer is accessed with a function call.

[eugenis]

Summary

This adds an -mllvm flag that forces the use of a runtime function call to
get the unsafe stack pointer, the same that is currently used on non-x86, non-aarch64 android.

[eugenis]

The function should really have coldcc or preserve_allcc but those are either not implemented (i.e. crash the compiler) or just not do anything useful on arm and aarch64.
As is, generated code is atrocious.

[eugenis]

Attempt to use a custom calling convention for the safestack pointer accessor function.

Not ready for review.

[eugenis]

An attempt to inline the safestack accessor function.

To use, add -mllvm -safestack-pointer-address-inline=1.
Function must be called safestack_pointer_address_inline.
"attribute__((used)) static" is recommended so the function is not discarded early.

[eugenis]

Remove -customcc code. Added inlining logic.

Please review.

[eugenis]

removed one flag, fixed a condition, improved the test

[eugenis]

fix one test

[pcc]

LGTM

clang/lib/Driver/ToolChains/CommonArgs.cpp

if (SanArgs.requiresMinimalRuntime()) {
       StaticRuntimes.push_back("ubsan_minimal");
     } else {
       StaticRuntimes.push_back("ubsan_standalone");
       if (SanArgs.linkCXXRuntimes())
         StaticRuntimes.push_back("ubsan_standalone_cxx");
     }
   }
   if (SanArgs.needsSafeStackRt()) {
-    NonWholeStaticRuntimes.push_back("safestack");
+    StaticRuntimes.push_back("safestack");
     RequiredSymbols.push_back("__safestack_init");
   }
   if (SanArgs.needsCfiRt())
     StaticRuntimes.push_back("cfi");
   if (SanArgs.needsCfiDiagRt()) {
     StaticRuntimes.push_back("cfi_diag");
     if (SanArgs.linkCXXRuntimes())
       StaticRuntimes.push_back("ubsan_standalone_cxx");

compiler-rt/lib/safestack/CMakeLists.txt

 add_compiler_rt_component(safestack)
 
-set(SAFESTACK_SOURCES safestack.cc)
+set(SAFESTACK_SOURCES safestack.cc safestack_customcc.S)
 
 include_directories(..)
 
 set(SAFESTACK_CFLAGS ${SANITIZER_COMMON_CFLAGS})
 
 if(APPLE)
   # Build universal binary on APPLE.
   add_compiler_rt_runtime(clang_rt.safestack

compiler-rt/lib/safestack/safestack.cc

 // the same mechanism as currently being used by the traditional stack
 // protector pass to store the stack guard (see getStackCookieLocation()
 // function above). Doing so requires changing the tcbhead_t struct in glibc
 // on Linux and tcb struct in libc on FreeBSD.
 //
 // For now, store it in a thread-local variable.
 extern "C" {
 __attribute__((visibility(
-    "default"))) __thread void *__safestack_unsafe_stack_ptr = nullptr;
+    "default"))) __attribute__((tls_model("initial-exec"))) __thread void *__safestack_unsafe_stack_ptr = nullptr;
 }
 
+
 // Per-thread unsafe stack information. It's not frequently accessed, so there
 // it can be kept out of the tcb in normal thread-local variables.
 static __thread void *unsafe_stack_start = nullptr;
 static __thread size_t unsafe_stack_size = 0;
 static __thread size_t unsafe_stack_guard = 0;
 
 using namespace __sanitizer;
 
extern "C"
     __attribute__((visibility("default"))) void *__get_unsafe_stack_start() {
   return unsafe_stack_start;
 }
 
 extern "C"
     __attribute__((visibility("default"))) void *__get_unsafe_stack_ptr() {
   return __safestack_unsafe_stack_ptr;
 }
+
+extern "C" __attribute__((visibility("default"))) void **
+__safestack_pointer_address() {
+  return &__safestack_unsafe_stack_ptr;
+}

compiler-rt/lib/safestack/safestack_customcc.S

+#if defined(__x86_64__)
+        .text
+        .globl  __safestack_pointer_address_customcc
+        .p2align        4, 0x90
+        .type   __safestack_pointer_address_customcc,@function
+__safestack_pointer_address_customcc: 
+        movq    %fs:0, %r14
+        addq    __safestack_unsafe_stack_ptr@GOTTPOFF(%rip), %r14
+        retq
+#endif

llvm/lib/CodeGen/SafeStack.cpp

 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <string>
 #include <utility>
 
 using namespace llvm;

 STATISTIC(NumAllocas, "Total number of allocas");
 STATISTIC(NumUnsafeStaticAllocas, "Number of unsafe static allocas");
 STATISTIC(NumUnsafeDynamicAllocas, "Number of unsafe dynamic allocas");
 STATISTIC(NumUnsafeByValArguments, "Number of unsafe byval arguments");
 STATISTIC(NumUnsafeStackRestorePoints, "Number of setjmps and landingpads");
 
 } // namespace llvm
 
+static cl::opt<bool>
+    SafeStackUsePointerAddressCustomCC("safestack-use-pointer-address-customcc",
+                               cl::init(false), cl::Hidden);
+
+static cl::opt<bool>
+    SafeStackUsePointerAddressInline("safestack-use-pointer-address-inline",
+                                     cl::init(false), cl::Hidden);
+
 namespace {
 
 /// Rewrite an SCEV expression for a memory access address to an expression that
 /// represents offset from the given alloca.
 ///
 /// The implementation simply replaces all mentions of the alloca with zero.
 class AllocaOffsetRewriter : public SCEVRewriteVisitor<AllocaOffsetRewriter> {
   const Value *AllocaPtr;
class SafeStack {
 
   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);
 
+  Value *callCustomCC(IRBuilder<> &IRB, StringRef Name);
+
 public:
   SafeStack(Function &F, const TargetLoweringBase &TL, const DataLayout &DL,
             ScalarEvolution &SE)
       : F(F), TL(TL), DL(DL), SE(SE),
         StackPtrTy(Type::getInt8PtrTy(F.getContext())),
         IntPtrTy(DL.getIntPtrType(F.getContext())),
         Int32Ty(Type::getInt32Ty(F.getContext())),
         Int8Ty(Type::getInt8Ty(F.getContext())) {}
for (inst_iterator It = inst_begin(&F), Ie = inst_end(&F); It != Ie;) {
         SI->takeName(II);
         assert(II->use_empty());
         II->eraseFromParent();
       }
     }
   }
 }
 
+Value *SafeStack::callCustomCC(IRBuilder<> &IRB, StringRef Name) {
+  Triple TargetTriple{F.getParent()->getTargetTriple()};
+  Value *Fn =
+      F.getParent()->getOrInsertFunction(Name, StackPtrTy->getPointerTo(0));
+  InlineAsm *Asm;
+  const char *AsmStr;
+  const char *ConstraintStr;
+  if (TargetTriple.getArch() == Triple::aarch64) {
+    AsmStr = "bl $1";
+    ConstraintStr = "={x19},s,~{memory},~{lr}";
+  } else if (TargetTriple.isARM() || TargetTriple.isThumb()) {
+    AsmStr = "bl $1";
+    ConstraintStr = "={r4},s,~{memory},~{lr}";
+  } else if (TargetTriple.getArch() == Triple::x86_64) {
+    AsmStr = "callq ${1:c}";
+    ConstraintStr = "={r14},s,~{memory}";
+  } else if (TargetTriple.getArch() == Triple::x86) {
+    AsmStr = "call ${1:c}";
+    ConstraintStr = "={esi},s,~{memory}";
+  } else {
+    report_fatal_error("Unsupported architecture for safestack-use-pointer-address");
+  }
+  Asm = InlineAsm::get(
+      FunctionType::get(StackPtrTy->getPointerTo(0), {Fn->getType()}, false),
+      AsmStr, ConstraintStr, false);
+  return IRB.CreateCall(Asm, Fn);
+}
+
 bool SafeStack::run() {
   assert(F.hasFnAttribute(Attribute::SafeStack) &&
          "Can't run SafeStack on a function without the attribute");
   assert(!F.isDeclaration() && "Can't run SafeStack on a function declaration");
 
   ++NumFunctions;
 
   SmallVector<AllocaInst *, 16> StaticAllocas;
bool SafeStack::run() {
   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());
+  if (SafeStackUsePointerAddressCustomCC) {
+    UnsafeStackPtr = callCustomCC(IRB, "__safestack_pointer_address_customcc");
+  } else if (SafeStackUsePointerAddressInline) {
+    Value *Fn = F.getParent()->getOrInsertFunction(
+        "__safestack_pointer_address_inline", StackPtrTy->getPointerTo(0));
+    UnsafeStackPtr = IRB.CreateCall(Fn);
+  } else {
     UnsafeStackPtr = TL.getSafeStackPointerLocation(IRB);
+  }
 
   // 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");
   assert(BasePointer->getType() == StackPtrTy);
 
   AllocaInst *StackGuardSlot = nullptr;
moveDynamicAllocasToUnsafeStack(F, UnsafeStackPtr, DynamicTop,
                                   DynamicAllocas);
 
   // Restore the unsafe stack pointer before each return.
   for (ReturnInst *RI : Returns) {
     IRB.SetInsertPoint(RI);
     IRB.CreateStore(BasePointer, UnsafeStackPtr);
   }
 
+  if (SafeStackUsePointerAddressInline) {
+    CallSite CS(UnsafeStackPtr);
+    InlineFunctionInfo IFI;
+    if (!InlineFunction(CS, IFI))
+      report_fatal_error(
+          "SafeStack failed to inline __safestack_pointer_address_inline in " +
+          F.getName());
+  }
+
   DEBUG(dbgs() << "[SafeStack]     safestack applied\n");
   return true;
 }
 
 class SafeStackLegacyPass : public FunctionPass {
   const TargetMachine *TM = nullptr;
 
 public:

llvm/test/CodeGen/X86/safestack.ll

 ; RUN: llc -mtriple=i386-linux < %s -o - | FileCheck --check-prefix=LINUX-I386 %s
 ; RUN: llc -mtriple=x86_64-linux < %s -o - | FileCheck --check-prefix=LINUX-X64 %s
 ; RUN: llc -mtriple=i386-linux-android < %s -o - | FileCheck --check-prefix=ANDROID-I386 %s
 ; RUN: llc -mtriple=x86_64-linux-android < %s -o - | FileCheck --check-prefix=ANDROID-X64 %s
 ; RUN: llc -mtriple=x86_64-fuchsia < %s -o - | FileCheck --check-prefix=FUCHSIA-X64 %s
 
+; RUN: llc -mtriple=i386-linux -safestack-use-pointer-address < %s -o - | FileCheck --check-prefix=LINUX-I386-PA %s
+
 define void @_Z1fv() safestack {
 entry:
   %x = alloca i32, align 4
   %0 = bitcast i32* %x to i8*
   call void @_Z7CapturePi(i32* nonnull %x)
   ret void
 }
 

 
 ; ANDROID-X64: movq %fs:72, %[[A:.*]]
 ; ANDROID-X64: leaq -16(%[[A]]), %[[B:.*]]
 ; ANDROID-X64: movq %[[B]], %fs:72
 
 ; FUCHSIA-X64: movq %fs:24, %[[A:.*]]
 ; FUCHSIA-X64: leaq -16(%[[A]]), %[[B:.*]]
 ; FUCHSIA-X64: movq %[[B]], %fs:24
+
+; LINUX-I386-PA: calll __safestack_pointer_address
+; LINUX-I386-PA: movl %eax, %[[A:.*]]
+; LINUX-I386-PA: movl (%eax), %[[B:.*]]
+; LINUX-I386-PA: leal -16(%[[B]]), %[[C:.*]]
+; LINUX-I386-PA: movl %[[C]], (%[[A]])