[safestack] Make sure the unsafe stack pointer is popped in all cases

The unsafe stack pointer is only popped in moveStaticAllocasToUnsafeStack so it won't happen if there are no static allocas.

Fixes https://llvm.org/bugs/show_bug.cgi?id=26122

Differential Revision: http://reviews.llvm.org/D16339

llvm-svn: 259447

llvm/lib/CodeGen/SafeStack.cpp

@@ -144,7 +144,8 @@ class SafeStack : public FunctionPass {
   Value *moveStaticAllocasToUnsafeStack(IRBuilder<> &IRB, Function &F,
                                         ArrayRef<AllocaInst *> StaticAllocas,
                                         ArrayRef<Argument *> ByValArguments,
-                                        ArrayRef<ReturnInst *> Returns);
+                                        ArrayRef<ReturnInst *> Returns,
+                                        Instruction *BasePointer);
 
   /// \brief Generate code to restore the stack after all stack restore points
   /// in \p StackRestorePoints.
@@ -431,6 +432,8 @@ AllocaInst *
 SafeStack::createStackRestorePoints(IRBuilder<> &IRB, Function &F,
                                     ArrayRef<Instruction *> StackRestorePoints,
                                     Value *StaticTop, bool NeedDynamicTop) {
+  assert(StaticTop && "The stack top isn't set.");
+
   if (StackRestorePoints.empty())
     return nullptr;
 
@@ -441,19 +444,13 @@ SafeStack::createStackRestorePoints(IRBuilder<> &IRB, Function &F,
   // runtime itself.
 
   AllocaInst *DynamicTop = nullptr;
-  if (NeedDynamicTop)
+  if (NeedDynamicTop) {
     // If we also have dynamic alloca's, the stack pointer value changes
     // throughout the function. For now we store it in an alloca.
     DynamicTop = IRB.CreateAlloca(StackPtrTy, /*ArraySize=*/nullptr,
                                   "unsafe_stack_dynamic_ptr");
-
-  if (!StaticTop)
-    // We need the original unsafe stack pointer value, even if there are
-    // no unsafe static allocas.
-    StaticTop = IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
-
-  if (NeedDynamicTop)
     IRB.CreateStore(StaticTop, DynamicTop);
+  }
 
   // Restore current stack pointer after longjmp/exception catch.
   for (Instruction *I : StackRestorePoints) {
@@ -467,29 +464,18 @@ SafeStack::createStackRestorePoints(IRBuilder<> &IRB, Function &F,
   return DynamicTop;
 }
 
+/// 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) {
+    ArrayRef<Argument *> ByValArguments, ArrayRef<ReturnInst *> Returns,
+    Instruction *BasePointer) {
   if (StaticAllocas.empty() && ByValArguments.empty())
-    return nullptr;
+    return BasePointer;
 
   DIBuilder DIB(*F.getParent());
 
-  // 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.
-
-  // 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);
-
-  for (ReturnInst *RI : Returns) {
-    IRB.SetInsertPoint(RI);
-    IRB.CreateStore(BasePointer, UnsafeStackPtr);
-  }
-
   // Compute maximum alignment among static objects on the unsafe stack.
   unsigned MaxAlignment = 0;
   for (Argument *Arg : ByValArguments) {
@@ -726,9 +712,16 @@ bool SafeStack::runOnFunction(Function &F) {
   IRBuilder<> IRB(&F.front(), F.begin()->getFirstInsertionPt());
   UnsafeStackPtr = getOrCreateUnsafeStackPtr(IRB, 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");
+  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);
+                                                    ByValArguments, Returns,
+                                                    BasePointer);
 
   // Safe stack object that stores the current unsafe stack top. It is updated
   // as unsafe dynamic (non-constant-sized) allocas are allocated and freed.
@@ -743,6 +736,12 @@ bool SafeStack::runOnFunction(Function &F) {
   moveDynamicAllocasToUnsafeStack(F, UnsafeStackPtr, DynamicTop,
                                   DynamicAllocas);
 
+  // Restore the unsafe stack pointer before each return.
+  for (ReturnInst *RI : Returns) {
+    IRB.SetInsertPoint(RI);
+    IRB.CreateStore(BasePointer, UnsafeStackPtr);
+  }
+
   DEBUG(dbgs() << "[SafeStack]     safestack applied\n");
   return true;
 }

llvm/test/Transforms/SafeStack/ARM/setjmp.ll

@@ -6,8 +6,8 @@
 define void @f(i32 %b) safestack {
 entry:
 ; CHECK: %[[SPA:.*]] = call i8** @__safestack_pointer_address()
-; CHECK: %[[USDP:.*]] = alloca i8*
 ; CHECK: %[[USP:.*]] = load i8*, i8** %[[SPA]]
+; CHECK: %[[USDP:.*]] = alloca i8*
 ; CHECK: store i8* %[[USP]], i8** %[[USDP]]
 ; CHECK: call i32 @setjmp
 
@@ -26,6 +26,8 @@ if.then:
   br label %if.end
 
 if.end:
+; CHECK: store i8* %[[USP:.*]], i8** %[[SPA:.*]]
+
   ret void
 }
 

llvm/test/Transforms/SafeStack/dynamic-alloca.ll

@@ -8,7 +8,7 @@
 ; Requires protector.
 define void @foo(i32 %n) nounwind uwtable safestack {
 entry:
-  ; CHECK: __safestack_unsafe_stack_ptr
+  ; CHECK: %[[SP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
   %n.addr = alloca i32, align 4
   %a = alloca i32*, align 8
   store i32 %n, i32* %n.addr, align 4
@@ -17,5 +17,6 @@ entry:
   %1 = alloca i8, i64 %conv
   %2 = bitcast i8* %1 to i32*
   store i32* %2, i32** %a, align 8
+  ; CHECK: store i8* %[[SP:.*]], i8** @__safestack_unsafe_stack_ptr
   ret void
 }

llvm/test/Transforms/SafeStack/setjmp2.ll

@@ -12,8 +12,8 @@
 ; CHECK: @foo(i32 %[[ARG:.*]])
 define i32 @foo(i32 %size) nounwind uwtable safestack {
 entry:
-  ; CHECK: %[[DYNPTR:.*]] = alloca i8*
-  ; CHECK-NEXT: %[[SP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+  ; CHECK: %[[SP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+  ; CHECK-NEXT: %[[DYNPTR:.*]] = alloca i8*
   ; CHECK-NEXT: store i8* %[[SP]], i8** %[[DYNPTR]]
 
   ; CHECK-NEXT: %[[ZEXT:.*]] = zext i32 %[[ARG]] to i64
@@ -35,6 +35,7 @@ entry:
 
   ; CHECK: call void @funcall(i32* %[[ALLOCA]])
   call void @funcall(i32* %a)
+  ; CHECK-NEXT: store i8* %[[SP:.*]], i8** @__safestack_unsafe_stack_ptr
   ret i32 0
 }