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

[zaks.anna]

Summary

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

[eugenis]

LGTM

lib/Transforms/Instrumentation/SafeStack.cpp

class SafeStack : public FunctionPass {
   /// replace allocas with pointers into the unsafe stack and generate code to
   /// restore the stack pointer before all return instructions in \p Returns.
   ///
   /// \returns A pointer to the top of the unsafe stack after all unsafe static
   /// allocas are allocated.
   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.
   ///
   /// \returns A local variable in which to maintain the dynamic top of the
   /// unsafe stack if needed.
   AllocaInst *
   createStackRestorePoints(IRBuilder<> &IRB, Function &F,
for (Argument &Arg : F.args()) {
     ByValArguments.push_back(&Arg);
   }
 }
 
 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;
 
   // We need the current value of the shadow stack pointer to restore
   // after longjmp or exception catching.
 
   // FIXME: On some platforms this could be handled by the longjmp/exception
   // 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) {
     ++NumUnsafeStackRestorePoints;
 
     IRB.SetInsertPoint(I->getNextNode());
     Value *CurrentTop = DynamicTop ? IRB.CreateLoad(DynamicTop) : StaticTop;
     IRB.CreateStore(CurrentTop, UnsafeStackPtr);
   }
 
   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) {
     Type *Ty = Arg->getType()->getPointerElementType();
     unsigned Align = std::max((unsigned)DL->getPrefTypeAlignment(Ty),
                               Arg->getParamAlignment());
     if (Align > MaxAlignment)
       MaxAlignment = Align;
if (!StaticAllocas.empty() || !DynamicAllocas.empty() ||
     ++NumUnsafeStackFunctions; // This function has the unsafe stack.
 
   if (!StackRestorePoints.empty())
     ++NumUnsafeStackRestorePointsFunctions;
 
   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.
   // This is only needed if we need to restore stack pointer after longjmp
   // or exceptions, and we have dynamic allocations.
   // FIXME: a better alternative might be to store the unsafe stack pointer
   // before setjmp / invoke instructions.
   AllocaInst *DynamicTop = createStackRestorePoints(
       IRB, F, StackRestorePoints, StaticTop, !DynamicAllocas.empty());
 
   // Handle dynamic allocas.
   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;
 }
 
 } // anonymous namespace
 
 char SafeStack::ID = 0;
 INITIALIZE_TM_PASS_BEGIN(SafeStack, "safe-stack",
                          "Safe Stack instrumentation pass", false, false)
 INITIALIZE_TM_PASS_END(SafeStack, "safe-stack",
                        "Safe Stack instrumentation pass", false, false)
 
 FunctionPass *llvm::createSafeStackPass(const llvm::TargetMachine *TM) {
   return new SafeStack(TM);
 }

test/Transforms/SafeStack/ARM/setjmp.ll

 ; Test stack pointer restore after setjmp() with the function-call safestack ABI.
 ; RUN: opt -safe-stack -S -mtriple=arm-linux-androideabi < %s -o - | FileCheck %s
 
 @env = global [64 x i32] zeroinitializer, align 4
 
 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
 
   %call = call i32 @setjmp(i32* getelementptr inbounds ([64 x i32], [64 x i32]* @env, i32 0, i32 0)) returns_twice
 
 ; CHECK: %[[USP2:.*]] = load i8*, i8** %[[USDP]]
 ; CHECK: store i8* %[[USP2]], i8** %[[SPA]]
 
   %tobool = icmp eq i32 %b, 0
   br i1 %tobool, label %if.end, label %if.then
 
 if.then:
   %0 = alloca [42 x i8], align 1
   %.sub = getelementptr inbounds [42 x i8], [42 x i8]* %0, i32 0, i32 0
   call void @_Z7CapturePv(i8* %.sub)
   br label %if.end
 
 if.end:
+; CHECK: store i8* %[[USP:.*]], i8** %[[SPA:.*]]
+
   ret void
 }
 
 declare i32 @setjmp(i32*) returns_twice
 
 declare void @_Z7CapturePv(i8*)

test/Transforms/SafeStack/dynamic-alloca.ll

 ; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
 ; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
 
 @.str = private unnamed_addr constant [4 x i8] c"%s\0A\00", align 1
 
 ; Variable sized alloca
 ;  safestack attribute
 ; 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
   %0 = load i32, i32* %n.addr, align 4
   %conv = sext i32 %0 to i64
   %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
 }

test/Transforms/SafeStack/setjmp2.ll

 ; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
 ; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
 
 %struct.__jmp_buf_tag = type { [8 x i64], i32, %struct.__sigset_t }
 %struct.__sigset_t = type { [16 x i64] }
 
 @.str = private unnamed_addr constant [4 x i8] c"%s\0A\00", align 1
 @buf = internal global [1 x %struct.__jmp_buf_tag] zeroinitializer, align 16
 
 ; setjmp/longjmp test with dynamically sized array.
 ; Requires protector.
 ; 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
   ; CHECK-NEXT: %[[MUL:.*]] = mul i64 %[[ZEXT]], 4
   ; CHECK-NEXT: %[[SP2:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
   ; CHECK-NEXT: %[[PTRTOINT:.*]] = ptrtoint i8* %[[SP2]] to i64
   ; CHECK-NEXT: %[[SUB:.*]] = sub i64 %[[PTRTOINT]], %[[MUL]]
   ; CHECK-NEXT: %[[AND:.*]] = and i64 %[[SUB]], -16
   ; CHECK-NEXT: %[[INTTOPTR:.*]] = inttoptr i64 %[[AND]] to i8*
   ; CHECK-NEXT: store i8* %[[INTTOPTR]], i8** @__safestack_unsafe_stack_ptr
   ; CHECK-NEXT: store i8* %[[INTTOPTR]], i8** %unsafe_stack_dynamic_ptr
   ; CHECK-NEXT: %[[ALLOCA:.*]] = bitcast i8* %[[INTTOPTR]] to i32*
   %a = alloca i32, i32 %size
 
   ; CHECK: setjmp
   ; CHECK-NEXT: %[[LOAD:.*]] = load i8*, i8** %[[DYNPTR]]
   ; CHECK-NEXT: store i8* %[[LOAD]], i8** @__safestack_unsafe_stack_ptr
   %call = call i32 @_setjmp(%struct.__jmp_buf_tag* getelementptr inbounds ([1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* @buf, i32 0, i32 0)) returns_twice
 
   ; CHECK: call void @funcall(i32* %[[ALLOCA]])
   call void @funcall(i32* %a)
+  ; CHECK-NEXT: store i8* %[[SP:.*]], i8** @__safestack_unsafe_stack_ptr
   ret i32 0
 }
 
 declare i32 @_setjmp(%struct.__jmp_buf_tag*)
 declare void @funcall(i32*)