3 files changed, 81 insertions, 9 deletions

diff --git a/lib/CodeGen/PrologEpilogInserter.cpp b/lib/CodeGen/PrologEpilogInserter.cpp
index d3b0b11c70..e118dd2449 100644
--- a/lib/CodeGen/PrologEpilogInserter.cpp
+++ b/lib/CodeGen/PrologEpilogInserter.cpp
@@ -406,6 +406,33 @@ void PEI::calculateFrameObjectOffsets(MachineFunction &Fn) {
     }
   }
 
+  // Make sure that the stack protector comes before the local variables on the
+  // stack.
+  if (FFI->hasStackProtector()) {
+    int FI = FFI->getStackProtectorIndex();
+
+    // If stack grows down, we need to add size of find the lowest
+    // address of the object.
+    if (StackGrowsDown)
+      Offset += FFI->getObjectSize(FI);
+
+    unsigned Align = FFI->getObjectAlignment(FI);
+
+    // If the alignment of this object is greater than that of the stack, then
+    // increase the stack alignment to match.
+    MaxAlign = std::max(MaxAlign, Align);
+
+    // Adjust to alignment boundary.
+    Offset = (Offset + Align - 1) / Align * Align;
+
+    if (StackGrowsDown) {
+      FFI->setObjectOffset(FI, -Offset); // Set the computed offset
+    } else {
+      FFI->setObjectOffset(FI, Offset);
+      Offset += FFI->getObjectSize(FI);
+    }
+  }
+
   // Then assign frame offsets to stack objects that are not used to spill
   // callee saved registers.
   for (unsigned i = 0, e = FFI->getObjectIndexEnd(); i != e; ++i) {
diff --git a/lib/CodeGen/SelectionDAG/SelectionDAGBuild.cpp b/lib/CodeGen/SelectionDAG/SelectionDAGBuild.cpp
index 6a2f2760ee..267ae36908 100644
--- a/lib/CodeGen/SelectionDAG/SelectionDAGBuild.cpp
+++ b/lib/CodeGen/SelectionDAG/SelectionDAGBuild.cpp
@@ -25,6 +25,7 @@
 #include "llvm/Instructions.h"
 #include "llvm/Intrinsics.h"
 #include "llvm/IntrinsicInst.h"
+#include "llvm/Module.h"
 #include "llvm/CodeGen/FastISel.h"
 #include "llvm/CodeGen/GCStrategy.h"
 #include "llvm/CodeGen/GCMetadata.h"
@@ -34,6 +35,7 @@
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/PseudoSourceValue.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/Target/TargetRegisterInfo.h"
 #include "llvm/Target/TargetData.h"
@@ -3793,6 +3795,47 @@ SelectionDAGLowering::visitIntrinsicCall(CallInst &I, unsigned Intrinsic) {
     DAG.setRoot(DAG.getNode(ISD::STACKRESTORE, MVT::Other, getRoot(), Tmp));
     return 0;
   }
+  case Intrinsic::stackprotector_prologue: {
+    // Emit code into the DAG to store the stack guard onto the stack.
+    MachineFunction &MF = DAG.getMachineFunction();
+    MachineFrameInfo *MFI = MF.getFrameInfo();
+    MVT PtrTy = TLI.getPointerTy();
+
+    // Retrieve the stack protector guard's value.
+    SDValue Src = getValue(I.getOperand(1));
+
+    // Create a slot on the stack for the stack protector. It should go first
+    // before local variables are allocated.
+    unsigned Align =
+      TLI.getTargetData()->getPrefTypeAlignment(PtrTy.getTypeForMVT());
+    int FI = MFI->CreateStackObject(PtrTy.getSizeInBits() / 8, Align);
+
+    MFI->setStackProtector(true);
+    MFI->setStackProtectorIndex(FI);
+
+    SDValue FIN = DAG.getFrameIndex(FI, PtrTy);
+
+    // Store the stack protector onto the stack.
+    SDValue Result = DAG.getStore(getRoot(), Src, FIN,
+                                  PseudoSourceValue::getFixedStack(FI),
+                                  0, true);
+    setValue(&I, Result);
+    DAG.setRoot(Result);
+    return 0;
+  }
+  case Intrinsic::stackprotector_epilogue: {
+    // Emit code into the DAG to retrieve the stack guard off of the stack.
+    MachineFunction &MF = DAG.getMachineFunction();
+    MachineFrameInfo *MFI = MF.getFrameInfo();
+    MVT PtrTy = TLI.getPointerTy();
+
+    // Load the value stored on the stack.
+    int FI = MFI->getStackProtectorIndex();
+    SDValue FIN = DAG.getFrameIndex(MFI->getStackProtectorIndex(), PtrTy);
+    setValue(&I, DAG.getLoad(PtrTy, getRoot(), FIN,
+                             PseudoSourceValue::getFixedStack(FI), 0, true));
+    return 0;
+  }
   case Intrinsic::var_annotation:
     // Discard annotate attributes
     return 0;
diff --git a/lib/CodeGen/StackProtector.cpp b/lib/CodeGen/StackProtector.cpp
index 5cd4c67b5f..659f8a01ab 100644
--- a/lib/CodeGen/StackProtector.cpp
+++ b/lib/CodeGen/StackProtector.cpp
@@ -20,6 +20,7 @@
 #include "llvm/DerivedTypes.h"
 #include "llvm/Function.h"
 #include "llvm/Instructions.h"
+#include "llvm/Intrinsics.h"
 #include "llvm/Module.h"
 #include "llvm/Pass.h"
 #include "llvm/ADT/APInt.h"
@@ -110,16 +111,15 @@ bool StackProtector::InsertStackProtectors() {
   // onto the stack.
   BasicBlock &Entry = F->getEntryBlock();
   Instruction *InsertPt = &Entry.front();
+
   const PointerType *GuardTy = PointerType::getUnqual(Type::Int8Ty);
 
   // The global variable for the stack guard.
   Constant *StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", GuardTy);
-
-  // The place on the stack that the stack protector guard is kept.
-  AllocaInst *StackProtFrameSlot =
-    new AllocaInst(GuardTy, "StackProt_Frame", InsertPt);
   LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsertPt);
-  new StoreInst(LI, StackProtFrameSlot, false, InsertPt);
+  CallInst::
+    Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_prologue),
+           LI, "", InsertPt);
 
   // Create the basic block to jump to when the guard check fails.
   BasicBlock *FailBB = CreateFailBB();
@@ -137,7 +137,7 @@ bool StackProtector::InsertStackProtectors() {
   //     %1 = load __stack_chk_guard
   //     %2 = load <stored stack guard>
   //     %3 = cmp i1 %1, %2
-  //     br i1 %3, label %SPRet, label %CallStackCheckFailBlk
+  //     br i1 %3, label %SP_return, label %CallStackCheckFailBlk
   //
   //   SP_return:
   //     ret ...
@@ -161,9 +161,11 @@ bool StackProtector::InsertStackProtectors() {
     F->getBasicBlockList().insert(InsPt, NewBB);
 
     // Generate the stack protector instructions in the old basic block.
-    LoadInst *LI2 = new LoadInst(StackGuardVar, "", false, BB);
-    LoadInst *LI1 = new LoadInst(StackProtFrameSlot, "", true, BB);
-    ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, LI1, LI2, "", BB);
+    LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
+    CallInst *CI = CallInst::
+      Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_epilogue),
+             "", BB);
+    ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, CI, LI1, "", BB);
     BranchInst::Create(NewBB, FailBB, Cmp, BB);
   }