[safestack] Fast access to the unsafe stack pointer on AArch64/Android (2nd try).

[eugenis]

Summary

Android libc provides a fixed TLS slot for the unsafe stack pointer,
and this change implements direct access to that slot on AArch64 via
__builtin_thread_pointer() + offset.

This change also moves more code into TargetLowering and its
target-specific subclasses to get rid of target-specific codegen
in SafeStackPass.

This change does not touch the ARM backend because ARM lowers
builting_thread_pointer as aeabi_read_tp, which is not available
on Android.

This is a second attempt which leave the generic, compiler-rt-based implementation in SafeStack.cpp. This way things still work even when there is no TargetMachine.

[eugenis]

To clarify, Eric, Renato, this goes against some of your comments in the original review. Namely, the TargetLowering methods are, again, Android-specific because I had to move the generic implementation to SafeStack pass to have it working without TargetMachine. This sounds like the right thing to do, because the compiler-rt based implementation does not require target knowledge to emit the IR, and this allows target-independent testing.

I'd like a confirmation that you are OK with this.

[rengolin]

Overall, this is less android specific in that only x86_64 and AArch64 knows about it, but the generic target lowering is independent. In that sense, it is better this way than previously.

lib/Transforms/Instrumentation/SafeStack.cpp

513

Moving this around makes it harder to review. Can you move it back where it was?

[eugenis]

lib/Transforms/Instrumentation/SafeStack.cpp

513

Sorry, did not noticed the method was moved. Fixed.

[echristo]

LGTM.

-eric

[eugenis]

Thank you.
r251324

include/llvm/Target/TargetLowering.h

public:
   /// Return true if the target stores stack protector cookies at a fixed offset
   /// in some non-standard address space, and populates the address space and
   /// offset as appropriate.
   virtual bool getStackCookieLocation(unsigned &/*AddressSpace*/,
                                       unsigned &/*Offset*/) const {
     return false;
   }
 
-  /// Return true if the target stores SafeStack pointer at a fixed offset in
-  /// some non-standard address space, and populates the address space and
-  /// offset as appropriate.
-  virtual bool getSafeStackPointerLocation(unsigned & /*AddressSpace*/,
-                                           unsigned & /*Offset*/) const {
-    return false;
-  }
+  /// If the target has a standard location for the unsafe stack pointer,
+  /// returns the address of that location. Otherwise, returns nullptr.
+  virtual Value *getSafeStackPointerLocation(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

lib/CodeGen/TargetLoweringBase.cpp

while (true) {
     if (LK.first == TypeSplitVector || LK.first == TypeExpandInteger)
       Cost *= 2;
 
     // Keep legalizing the type.
     MTy = LK.second;
   }
 }
 
+Value *TargetLoweringBase::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
+  if (!TM.getTargetTriple().isAndroid())
+    return nullptr;
+
+  // 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);
+}
+
 //===----------------------------------------------------------------------===//
 //  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/Target/AArch64/AArch64ISelLowering.h

public:
   shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const override;
 
   bool shouldExpandAtomicCmpXchgInIR(AtomicCmpXchgInst *AI) const override;
 
   bool useLoadStackGuardNode() const override;
   TargetLoweringBase::LegalizeTypeAction
   getPreferredVectorAction(EVT VT) const override;
 
+  /// If the target has a standard location for the unsafe stack pointer,
+  /// returns the address of that location. Otherwise, returns nullptr.
+  Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const override;
+
 private:
   bool isExtFreeImpl(const Instruction *Ext) const override;
 
   /// Subtarget - Keep a pointer to the AArch64Subtarget around so that we can
   /// make the right decision when generating code for different targets.
   const AArch64Subtarget *Subtarget;
 
   void addTypeForNEON(EVT VT, EVT PromotedBitwiseVT);

lib/Target/AArch64/AArch64ISelLowering.cpp

bool AArch64TargetLowering::functionArgumentNeedsConsecutiveRegisters(
     Type *Ty, CallingConv::ID CallConv, bool isVarArg) const {
   return Ty->isArrayTy();
 }
 
 bool AArch64TargetLowering::shouldNormalizeToSelectSequence(LLVMContext &,
                                                             EVT) const {
   return false;
 }
+
+Value *AArch64TargetLowering::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
+  if (!Subtarget->isTargetAndroid())
+    return TargetLowering::getSafeStackPointerLocation(IRB);
+
+  // Android provides a fixed TLS slot for the SafeStack pointer. See the
+  // definition of TLS_SLOT_SAFESTACK in
+  // https://android.googlesource.com/platform/bionic/+/master/libc/private/bionic_tls.h
+  const unsigned TlsOffset = 0x48;
+  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
+  Function *ThreadPointerFunc =
+      Intrinsic::getDeclaration(M, Intrinsic::aarch64_thread_pointer);
+  return IRB.CreatePointerCast(
+      IRB.CreateConstGEP1_32(IRB.CreateCall(ThreadPointerFunc), TlsOffset),
+      Type::getInt8PtrTy(IRB.getContext())->getPointerTo(0));
+}

lib/Target/X86/X86ISelLowering.h

public:
     /// offset in some non-standard address space, and populates the address
     /// space and offset as appropriate.
     bool getStackCookieLocation(unsigned &AddressSpace,
                                 unsigned &Offset) const override;
 
     /// Return true if the target stores SafeStack pointer at a fixed offset in
     /// some non-standard address space, and populates the address space and
     /// offset as appropriate.
-    bool getSafeStackPointerLocation(unsigned &AddressSpace,
-                                     unsigned &Offset) const override;
+    Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const override;
 
     SDValue BuildFILD(SDValue Op, EVT SrcVT, SDValue Chain, SDValue StackSlot,
                       SelectionDAG &DAG) const;
 
     bool isNoopAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const override;
 
     bool useLoadStackGuardNode() const override;
     /// \brief Customize the preferred legalization strategy for certain types.

lib/Target/X86/X86ISelLowering.cpp

bool X86TargetLowering::getStackCookieLocation(unsigned &AddressSpace,
   } else {
     // %gs:0x14 on i386
     Offset = 0x14;
     AddressSpace = 256;
   }
   return true;
 }
 
-/// Android provides a fixed TLS slot for the SafeStack pointer.
-/// See the definition of TLS_SLOT_SAFESTACK in
-/// https://android.googlesource.com/platform/bionic/+/master/libc/private/bionic_tls.h
-bool X86TargetLowering::getSafeStackPointerLocation(unsigned &AddressSpace,
-                                                    unsigned &Offset) const {
+Value *X86TargetLowering::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
   if (!Subtarget->isTargetAndroid())
-    return false;
+    return TargetLowering::getSafeStackPointerLocation(IRB);
 
+  // Android provides a fixed TLS slot for the SafeStack pointer. See the
+  // definition of TLS_SLOT_SAFESTACK in
+  // https://android.googlesource.com/platform/bionic/+/master/libc/private/bionic_tls.h
+  unsigned AddressSpace, Offset;
   if (Subtarget->is64Bit()) {
     // %fs:0x48, unless we're using a Kernel code model, in which case it's %gs:
     Offset = 0x48;
     if (getTargetMachine().getCodeModel() == CodeModel::Kernel)
       AddressSpace = 256;
     else
       AddressSpace = 257;
   } else {
     // %gs:0x24 on i386
     Offset = 0x24;
     AddressSpace = 256;
   }
-  return true;
+
+  return ConstantExpr::getIntToPtr(
+      ConstantInt::get(Type::getInt32Ty(IRB.getContext()), Offset),
+      Type::getInt8PtrTy(IRB.getContext())->getPointerTo(AddressSpace));
 }
 
 bool X86TargetLowering::isNoopAddrSpaceCast(unsigned SrcAS,
                                             unsigned DestAS) const {
   assert(SrcAS != DestAS && "Expected different address spaces!");
 
   return SrcAS < 256 && DestAS < 256;
 }

lib/Transforms/Instrumentation/SafeStack.cpp

 #include "llvm/Target/TargetSubtargetInfo.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "safestack"
 
-static const char *const kUnsafeStackPtrVar = "__safestack_unsafe_stack_ptr";
-static const char *const kUnsafeStackPtrAddrFn = "__safestack_pointer_address";
-
 namespace llvm {
 
 STATISTIC(NumFunctions, "Total number of functions");
 STATISTIC(NumUnsafeStackFunctions, "Number of functions with unsafe stack");
 STATISTIC(NumUnsafeStackRestorePointsFunctions,
           "Number of functions that use setjmp or exceptions");
 
 STATISTIC(NumAllocas, "Total number of allocas");
class SafeStack : public FunctionPass {
   /// 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 Build a constant representing a pointer to the unsafe stack
-  /// pointer.
+  /// \brief Build a value representing a pointer to the unsafe stack pointer.
   Value *getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F);
 
   /// \brief Find all static allocas, dynamic allocas, return instructions and
   /// stack restore points (exception unwind blocks and setjmp calls) in the
   /// given function and append them to the respective vectors.
   void findInsts(Function &F, SmallVectorImpl<AllocaInst *> &StaticAllocas,
                  SmallVectorImpl<AllocaInst *> &DynamicAllocas,
                  SmallVectorImpl<ReturnInst *> &Returns,
bool doInitialization(Module &M) override {
 
     return false;
   }
 
   bool runOnFunction(Function &F) override;
 }; // class SafeStack
 
 Value *SafeStack::getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F) {
-  Module &M = *F.getParent();
-  Triple TargetTriple(M.getTargetTriple());
+  // Check if there is a target-specific location for the unsafe stack pointer.
+  if (TLI)
+    if (Value *V = TLI->getSafeStackPointerLocation(IRB))
+      return V;
 
-  unsigned Offset;
-  unsigned AddressSpace;
-  // Check if the target keeps the unsafe stack pointer at a fixed offset.
-  if (TLI && TLI->getSafeStackPointerLocation(AddressSpace, Offset)) {
-    Constant *OffsetVal =
-        ConstantInt::get(Type::getInt32Ty(F.getContext()), Offset);
-    return ConstantExpr::getIntToPtr(OffsetVal,
-                                     StackPtrTy->getPointerTo(AddressSpace));
-  }
-
-  // Android provides a libc function that returns the stack pointer address.
-  if (TargetTriple.isAndroid()) {
-    Value *Fn = M.getOrInsertFunction(kUnsafeStackPtrAddrFn,
-                                      StackPtrTy->getPointerTo(0), nullptr);
-    return IRB.CreateCall(Fn);
-  } else {
-    // Otherwise, declare a thread-local variable with a magic name.
-    auto UnsafeStackPtr =
-        dyn_cast_or_null<GlobalVariable>(M.getNamedValue(kUnsafeStackPtrVar));
+  // Otherwise, assume the target links with compiler-rt, which provides a
+  // thread-local variable with a magic name.
+  Module &M = *F.getParent();
+  const char *UnsafeStackPtrVar = "__safestack_unsafe_stack_ptr";
+      dyn_cast_or_null<GlobalVariable>(M.getNamedValue(UnsafeStackPtrVar));
 
-    if (!UnsafeStackPtr) {
-      // The global variable is not defined yet, define it ourselves.
-      // 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(
-          /*Module=*/M, /*Type=*/StackPtrTy,
-          /*isConstant=*/false, /*Linkage=*/GlobalValue::ExternalLinkage,
-          /*Initializer=*/nullptr, /*Name=*/kUnsafeStackPtrVar,
-          /*InsertBefore=*/nullptr,
-          /*ThreadLocalMode=*/GlobalValue::InitialExecTLSModel);
-    } else {
-      // The variable exists, check its type and attributes.
-      if (UnsafeStackPtr->getValueType() != StackPtrTy) {
-        report_fatal_error(Twine(kUnsafeStackPtrVar) + " must have void* type");
-      }
-
-      if (!UnsafeStackPtr->isThreadLocal()) {
-        report_fatal_error(Twine(kUnsafeStackPtrVar) + " must be thread-local");
-      }
-    }
-    return UnsafeStackPtr;
-  }
-}
+        M, StackPtrTy, false, GlobalValue::ExternalLinkage, 0,
+        UnsafeStackPtrVar, nullptr, GlobalValue::InitialExecTLSModel);
+    if (UnsafeStackPtr->getValueType() != StackPtrTy)
+      report_fatal_error(Twine(UnsafeStackPtrVar) + " must have void* type");
+    if (!UnsafeStackPtr->isThreadLocal())
+      report_fatal_error(Twine(UnsafeStackPtrVar) + " must be thread-local");
 
 void SafeStack::findInsts(Function &F,
                           SmallVectorImpl<AllocaInst *> &StaticAllocas,
                           SmallVectorImpl<AllocaInst *> &DynamicAllocas,
                           SmallVectorImpl<ReturnInst *> &Returns,
                           SmallVectorImpl<Instruction *> &StackRestorePoints) {
   for (Instruction &I : instructions(&F)) {
     if (auto AI = dyn_cast<AllocaInst>(&I)) {
for (inst_iterator It = inst_begin(&F), Ie = inst_end(&F); It != Ie;) {
         SI->takeName(II);
         assert(II->use_empty());
         II->eraseFromParent();
       }
     }
   }
 }
 
 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;
   }
 

test/Transforms/SafeStack/AArch64/abi.ll

 ; RUN: opt -safe-stack -S -mtriple=aarch64-linux-android < %s -o - | FileCheck %s
 
 
 define void @foo() nounwind uwtable safestack {
 entry:
-; CHECK: %[[SPA:.*]] = call i8** @__safestack_pointer_address()
+; CHECK: %[[TP:.*]] = call i8* @llvm.aarch64.thread.pointer()
+; CHECK: %[[SPA0:.*]] = getelementptr i8, i8* %[[TP]], i32 72
+; CHECK: %[[SPA:.*]] = bitcast i8* %[[SPA0]] to i8**
 ; CHECK: %[[USP:.*]] = load i8*, i8** %[[SPA]]
 ; CHECK: %[[USST:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
 ; CHECK: store i8* %[[USST]], i8** %[[SPA]]
 
   %a = alloca i8, align 8
   call void @Capture(i8* %a)
 
 ; CHECK: store i8* %[[USP]], i8** %[[SPA]]
   ret void
 }
 
 declare void @Capture(i8*)