From cad7994c3b4f023767fbc5cf66ed4eeee940a149 Mon Sep 17 00:00:00 2001 From: Anna Zaks Date: Tue, 2 Feb 2016 01:03:11 +0000 Subject: [PATCH] [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 | 53 +++++++++---------- llvm/test/Transforms/SafeStack/ARM/setjmp.ll | 4 +- .../Transforms/SafeStack/dynamic-alloca.ll | 3 +- llvm/test/Transforms/SafeStack/setjmp2.ll | 5 +- 4 files changed, 34 insertions(+), 31 deletions(-) diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp index 7adcdea8d9ea..c1c6ab6941e5 100644 --- a/llvm/lib/CodeGen/SafeStack.cpp +++ b/llvm/lib/CodeGen/SafeStack.cpp @@ -144,7 +144,8 @@ class SafeStack : public FunctionPass { Value *moveStaticAllocasToUnsafeStack(IRBuilder<> &IRB, Function &F, ArrayRef StaticAllocas, ArrayRef ByValArguments, - ArrayRef Returns); + ArrayRef 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 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 StaticAllocas, - ArrayRef ByValArguments, ArrayRef Returns) { + ArrayRef ByValArguments, ArrayRef 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; } diff --git a/llvm/test/Transforms/SafeStack/ARM/setjmp.ll b/llvm/test/Transforms/SafeStack/ARM/setjmp.ll index 8c57908bbe4b..20e46f8f0e21 100644 --- a/llvm/test/Transforms/SafeStack/ARM/setjmp.ll +++ b/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 } diff --git a/llvm/test/Transforms/SafeStack/dynamic-alloca.ll b/llvm/test/Transforms/SafeStack/dynamic-alloca.ll index bfec66f82a2f..b0571f72f1aa 100644 --- a/llvm/test/Transforms/SafeStack/dynamic-alloca.ll +++ b/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 } diff --git a/llvm/test/Transforms/SafeStack/setjmp2.ll b/llvm/test/Transforms/SafeStack/setjmp2.ll index bb15d7e03ace..dc83c4824207 100644 --- a/llvm/test/Transforms/SafeStack/setjmp2.ll +++ b/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 }