[safestack] Add canary to unsafe stack frames

Add StackProtector to SafeStack. This adds limited protection against
data corruption in the caller frame. Current implementation treats
all stack protector levels as -fstack-protector-all.

llvm-svn: 266004

lib/CodeGen/SafeStack.cpp

@@ -17,6 +17,7 @@
 
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/CodeGen/Passes.h"
@@ -31,6 +32,7 @@
 #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/Pass.h"
 #include "llvm/Support/CommandLine.h"
@@ -40,6 +42,7 @@
 #include "llvm/Support/raw_os_ostream.h"
 #include "llvm/Target/TargetLowering.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
@@ -122,6 +125,13 @@ class SafeStack : public FunctionPass {
   /// \brief Build a value representing a pointer to the unsafe stack pointer.
   Value *getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F);
 
+  /// \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
   /// stack restore points (exception unwind blocks and setjmp calls) in the
   /// given function and append them to the respective vectors.
@@ -145,7 +155,8 @@ class SafeStack : public FunctionPass {
                                         ArrayRef<AllocaInst *> StaticAllocas,
                                         ArrayRef<Argument *> ByValArguments,
                                         ArrayRef<ReturnInst *> Returns,
-                                        Instruction *BasePointer);
+                                        Instruction *BasePointer,
+                                        AllocaInst *StackGuardSlot);
 
   /// \brief Generate code to restore the stack after all stack restore points
   /// in \p StackRestorePoints.
@@ -379,6 +390,16 @@ Value *SafeStack::getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F) {
   return UnsafeStackPtr;
 }
 
+Value *SafeStack::getStackGuard(IRBuilder<> &IRB, Function &F) {
+  Value *StackGuardVar = nullptr;
+  if (TL)
+    StackGuardVar = TL->getIRStackGuard(IRB);
+  if (!StackGuardVar)
+    StackGuardVar =
+        F.getParent()->getOrInsertGlobal("__stack_chk_guard", StackPtrTy);
+  return IRB.CreateLoad(StackGuardVar, "StackGuard");
+}
+
 void SafeStack::findInsts(Function &F,
                           SmallVectorImpl<AllocaInst *> &StaticAllocas,
                           SmallVectorImpl<AllocaInst *> &DynamicAllocas,
@@ -464,13 +485,33 @@ SafeStack::createStackRestorePoints(IRBuilder<> &IRB, Function &F,
   return DynamicTop;
 }
 
+void SafeStack::checkStackGuard(IRBuilder<> &IRB, Function &F, ReturnInst &RI,
+                                AllocaInst *StackGuardSlot, Value *StackGuard) {
+  Value *V = IRB.CreateLoad(StackGuardSlot);
+  Value *Cmp = IRB.CreateICmpNE(StackGuard, V);
+
+  auto SuccessProb = BranchProbabilityInfo::getBranchProbStackProtector(true);
+  auto FailureProb = BranchProbabilityInfo::getBranchProbStackProtector(false);
+  MDNode *Weights = MDBuilder(F.getContext())
+                        .createBranchWeights(SuccessProb.getNumerator(),
+                                             FailureProb.getNumerator());
+  Instruction *CheckTerm =
+      SplitBlockAndInsertIfThen(Cmp, &RI,
+                                /* Unreachable */ true, Weights);
+  IRBuilder<> IRBFail(CheckTerm);
+  // FIXME: respect -fsanitize-trap / -ftrap-function here?
+  Constant *StackChkFail = F.getParent()->getOrInsertFunction(
+      "__stack_chk_fail", IRB.getVoidTy(), nullptr);
+  IRBFail.CreateCall(StackChkFail, {});
+}
+
 /// We explicitly compute and set the unsafe stack layout for all unsafe
 /// static alloca instructions. We save the unsafe "base pointer" in the
 /// prologue into a local variable and restore it in the epilogue.
 Value *SafeStack::moveStaticAllocasToUnsafeStack(
     IRBuilder<> &IRB, Function &F, ArrayRef<AllocaInst *> StaticAllocas,
     ArrayRef<Argument *> ByValArguments, ArrayRef<ReturnInst *> Returns,
-    Instruction *BasePointer) {
+    Instruction *BasePointer, AllocaInst *StackGuardSlot) {
   if (StaticAllocas.empty() && ByValArguments.empty())
     return BasePointer;
 
@@ -506,6 +547,18 @@ Value *SafeStack::moveStaticAllocasToUnsafeStack(
   int64_t StaticOffset = 0; // Current stack top.
   IRB.SetInsertPoint(BasePointer->getNextNode());
 
+  if (StackGuardSlot) {
+    StaticOffset += getStaticAllocaAllocationSize(StackGuardSlot);
+    Value *Off = IRB.CreateGEP(BasePointer, // BasePointer is i8*
+                               ConstantInt::get(Int32Ty, -StaticOffset));
+    Value *NewAI =
+        IRB.CreateBitCast(Off, StackGuardSlot->getType(), "StackGuardSlot");
+
+    // Replace alloc with the new location.
+    StackGuardSlot->replaceAllUsesWith(NewAI);
+    StackGuardSlot->eraseFromParent();
+  }
+
   for (Argument *Arg : ByValArguments) {
     Type *Ty = Arg->getType()->getPointerElementType();
 
@@ -667,18 +720,6 @@ bool SafeStack::runOnFunction(Function &F) {
   TL = TM ? TM->getSubtargetImpl(F)->getTargetLowering() : nullptr;
   SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
 
-  {
-    // Make sure the regular stack protector won't run on this function
-    // (safestack attribute takes precedence).
-    AttrBuilder B;
-    B.addAttribute(Attribute::StackProtect)
-        .addAttribute(Attribute::StackProtectReq)
-        .addAttribute(Attribute::StackProtectStrong);
-    F.removeAttributes(
-        AttributeSet::FunctionIndex,
-        AttributeSet::get(F.getContext(), AttributeSet::FunctionIndex, B));
-  }
-
   ++NumFunctions;
 
   SmallVector<AllocaInst *, 16> StaticAllocas;
@@ -715,13 +756,29 @@ bool SafeStack::runOnFunction(Function &F) {
   // 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");
+      IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
   assert(BasePointer->getType() == StackPtrTy);
 
-  // The top of the unsafe stack after all unsafe static allocas are allocated.
-  Value *StaticTop = moveStaticAllocasToUnsafeStack(IRB, F, StaticAllocas,
-                                                    ByValArguments, Returns,
-                                                    BasePointer);
+  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);
+    StackGuardSlot = IRB.CreateAlloca(StackPtrTy, nullptr);
+    IRB.CreateStore(StackGuard, StackGuardSlot);
+
+    for (ReturnInst *RI : Returns) {
+      IRBuilder<> IRBRet(RI);
+      checkStackGuard(IRBRet, F, *RI, StackGuardSlot, StackGuard);
+    }
+  }
+
+  // The top of the unsafe stack after all unsafe static allocas are
+  // allocated.
+  Value *StaticTop =
+      moveStaticAllocasToUnsafeStack(IRB, F, StaticAllocas, ByValArguments,
+                                     Returns, BasePointer, StackGuardSlot);
 
   // Safe stack object that stores the current unsafe stack top. It is updated
   // as unsafe dynamic (non-constant-sized) allocas are allocated and freed.

lib/CodeGen/StackProtector.cpp

@@ -210,6 +210,9 @@ bool StackProtector::RequiresStackProtector() {
                                       Intrinsic::stackprotector))
           HasPrologue = true;
 
+  if (F->hasFnAttribute(Attribute::SafeStack))
+    return false;
+
   if (F->hasFnAttribute(Attribute::StackProtectReq)) {
     NeedsProtector = true;
     Strong = true; // Use the same heuristic as strong to determine SSPLayout

lib/IR/Attributes.cpp

@@ -1477,20 +1477,14 @@ static void adjustCallerSSPLevel(Function &Caller, const Function &Callee) {
                                               AttributeSet::FunctionIndex,
                                               B);
 
-  if (Callee.hasFnAttribute(Attribute::SafeStack)) {
-    Caller.removeAttributes(AttributeSet::FunctionIndex, OldSSPAttr);
-    Caller.addFnAttr(Attribute::SafeStack);
-  } else if (Callee.hasFnAttribute(Attribute::StackProtectReq) &&
-             !Caller.hasFnAttribute(Attribute::SafeStack)) {
+  if (Callee.hasFnAttribute(Attribute::StackProtectReq)) {
     Caller.removeAttributes(AttributeSet::FunctionIndex, OldSSPAttr);
     Caller.addFnAttr(Attribute::StackProtectReq);
   } else if (Callee.hasFnAttribute(Attribute::StackProtectStrong) &&
-             !Caller.hasFnAttribute(Attribute::SafeStack) &&
              !Caller.hasFnAttribute(Attribute::StackProtectReq)) {
     Caller.removeAttributes(AttributeSet::FunctionIndex, OldSSPAttr);
     Caller.addFnAttr(Attribute::StackProtectStrong);
   } else if (Callee.hasFnAttribute(Attribute::StackProtect) &&
-             !Caller.hasFnAttribute(Attribute::SafeStack) &&
              !Caller.hasFnAttribute(Attribute::StackProtectReq) &&
              !Caller.hasFnAttribute(Attribute::StackProtectStrong))
     Caller.addFnAttr(Attribute::StackProtect);

test/CodeGen/X86/safestack_ssp.ll

@@ -0,0 +1,27 @@
+; Test codegen pipeline for SafeStack + StackProtector combination.
+; 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
+
+define void @_Z1fv() safestack sspreq {
+entry:
+  %x = alloca i32, align 4
+  %0 = bitcast i32* %x to i8*
+  call void @_Z7CapturePi(i32* nonnull %x)
+  ret void
+}
+
+declare void @_Z7CapturePi(i32*)
+
+; LINUX-X64-DAG: movq __safestack_unsafe_stack_ptr@GOTTPOFF(%rip), %[[A:.*]]
+; LINUX-X64-DAG: movq %fs:(%[[A]]), %[[B:.*]]
+; LINUX-X64-DAG: movq %fs:40, %[[COOKIE:.*]]
+; LINUX-X64-DAG: leaq -16(%[[B]]), %[[C:.*]]
+; LINUX-X64-DAG: movq %[[C]], %fs:(%[[A]])
+; LINUX-X64-DAG: movq %[[COOKIE]], -8(%[[B]])
+
+; LINUX-I386-DAG: movl __safestack_unsafe_stack_ptr@INDNTPOFF, %[[A:.*]]
+; LINUX-I386-DAG: movl %gs:(%[[A]]), %[[B:.*]]
+; LINUX-I386-DAG: movl %gs:20, %[[COOKIE:.*]]
+; LINUX-I386-DAG: leal -16(%[[B]]), %[[C:.*]]
+; LINUX-I386-DAG: movl %[[C]], %gs:(%[[A]])
+; LINUX-I386-DAG: movl %[[COOKIE]], -4(%[[B]])

test/Transforms/SafeStack/AArch64/abi_ssp.ll

@@ -0,0 +1,22 @@
+; RUN: opt -safe-stack -S -mtriple=aarch64-linux-android < %s -o - | FileCheck --check-prefix=TLS %s
+
+
+define void @foo() nounwind uwtable safestack sspreq {
+entry:
+; The first @llvm.aarch64.thread.pointer is for the unsafe stack pointer, skip it.
+; TLS: call i8* @llvm.aarch64.thread.pointer()
+
+; TLS: %[[TP2:.*]] = call i8* @llvm.aarch64.thread.pointer()
+; TLS: %[[B:.*]] = getelementptr i8, i8* %[[TP2]], i32 40
+; TLS: %[[C:.*]] = bitcast i8* %[[B]] to i8**
+; TLS: %[[StackGuard:.*]] = load i8*, i8** %[[C]]
+; TLS: store i8* %[[StackGuard]], i8** %[[StackGuardSlot:.*]]
+  %a = alloca i128, align 16
+  call void @Capture(i128* %a)
+
+; TLS: %[[A:.*]] = load i8*, i8** %[[StackGuardSlot]]
+; TLS: icmp ne i8* %[[StackGuard]], %[[A]]
+  ret void
+}
+
+declare void @Capture(i128*)

test/Transforms/SafeStack/X86/abi_ssp.ll

@@ -0,0 +1,19 @@
+; RUN: opt -safe-stack -S -mtriple=i686-pc-linux-gnu < %s -o - | FileCheck --check-prefix=TLS --check-prefix=TLS32 %s
+; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck --check-prefix=TLS --check-prefix=TLS64 %s
+; RUN: opt -safe-stack -S -mtriple=i686-linux-android < %s -o - | FileCheck --check-prefix=TLS --check-prefix=TLS32 %s
+; RUN: opt -safe-stack -S -mtriple=x86_64-linux-android < %s -o - | FileCheck --check-prefix=TLS --check-prefix=TLS64 %s
+
+define void @foo() safestack sspreq {
+entry:
+; TLS32: %[[StackGuard:.*]] = load i8*, i8* addrspace(256)* inttoptr (i32 20 to i8* addrspace(256)*)
+; TLS64: %[[StackGuard:.*]] = load i8*, i8* addrspace(257)* inttoptr (i32 40 to i8* addrspace(257)*)
+; TLS:   store i8* %[[StackGuard]], i8** %[[StackGuardSlot:.*]]
+  %a = alloca i8, align 1
+  call void @Capture(i8* %a)
+
+; TLS: %[[A:.*]] = load i8*, i8** %[[StackGuardSlot]]
+; TLS: icmp ne i8* %[[StackGuard]], %[[A]]
+  ret void
+}
+
+declare void @Capture(i8*)

test/Transforms/SafeStack/X86/ssp.ll

@@ -0,0 +1,30 @@
+; RUN: opt -safe-stack -S -mtriple=x86_64-unknown < %s -o - | FileCheck %s
+
+define void @foo() safestack sspreq {
+entry:
+; CHECK: %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK: %[[USST:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
+; CHECK: store i8* %[[USST]], i8** @__safestack_unsafe_stack_ptr
+
+; CHECK: %[[A:.*]] = getelementptr i8, i8* %[[USP]], i32 -8
+; CHECK: %[[StackGuardSlot:.*]] = bitcast i8* %[[A]] to i8**
+; CHECK: %[[StackGuard:.*]] = load i8*, i8** @__stack_chk_guard
+; CHECK: store i8* %[[StackGuard]], i8** %[[StackGuardSlot]]
+  %a = alloca i8, align 1
+
+; CHECK: call void @Capture
+  call void @Capture(i8* %a)
+
+; CHECK: %[[B:.*]] = load i8*, i8** %[[StackGuardSlot]]
+; CHECK: %[[COND:.*]] = icmp ne i8* %[[StackGuard]], %[[B]]
+; CHECK: br i1 %[[COND]], {{.*}} !prof
+
+; CHECK:      call void @__stack_chk_fail()
+; CHECK-NEXT: unreachable
+
+; CHECK:      store i8* %[[USP]], i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT: ret void
+  ret void
+}
+
+declare void @Capture(i8*)