Index: llvm/lib/CodeGen/StackProtector.cpp =================================================================== --- llvm/lib/CodeGen/StackProtector.cpp +++ llvm/lib/CodeGen/StackProtector.cpp @@ -471,21 +471,36 @@ // instrumentation has already been generated. HasIRCheck = true; + // If we're instrumenting a block with a musttail call, the check has to be + // inserted before the call rather than between it and the return. The + // verifier guarantees that a musttail call is either directly before the + // return or with a single correct bitcast of the return value in between so + // we don't need to worry about many situations here. + Instruction *CheckLoc = RI; + Instruction *Prev = RI->getPrevNonDebugInstruction(); + if (Prev && isa(Prev) && cast(Prev)->isMustTailCall()) + CheckLoc = Prev; + else if (Prev) { + Prev = Prev->getPrevNonDebugInstruction(); + if (Prev && isa(Prev) && cast(Prev)->isMustTailCall()) + CheckLoc = Prev; + } + // Generate epilogue instrumentation. The epilogue intrumentation can be // function-based or inlined depending on which mechanism the target is // providing. if (Function *GuardCheck = TLI->getSSPStackGuardCheck(*M)) { // Generate the function-based epilogue instrumentation. // The target provides a guard check function, generate a call to it. - IRBuilder<> B(RI); + IRBuilder<> B(CheckLoc); LoadInst *Guard = B.CreateLoad(B.getInt8PtrTy(), AI, true, "Guard"); CallInst *Call = B.CreateCall(GuardCheck, {Guard}); Call->setAttributes(GuardCheck->getAttributes()); Call->setCallingConv(GuardCheck->getCallingConv()); } else { // Generate the epilogue with inline instrumentation. - // If we do not support SelectionDAG based tail calls, generate IR level - // tail calls. + // If we do not support SelectionDAG based calls, generate IR level + // calls. // // For each block with a return instruction, convert this: // @@ -515,7 +530,8 @@ BasicBlock *FailBB = CreateFailBB(); // Split the basic block before the return instruction. - BasicBlock *NewBB = BB->splitBasicBlock(RI->getIterator(), "SP_return"); + BasicBlock *NewBB = + BB->splitBasicBlock(CheckLoc->getIterator(), "SP_return"); // Update the dominator tree if we need to. if (DT && DT->isReachableFromEntry(BB)) { Index: llvm/test/CodeGen/AArch64/stack-protector-musttail.ll =================================================================== --- /dev/null +++ llvm/test/CodeGen/AArch64/stack-protector-musttail.ll @@ -0,0 +1,66 @@ +; RUN: llc -mtriple=arm64-apple-macosx -fast-isel %s -o - -start-before=stack-protector -stop-after=stack-protector | FileCheck %s + +@var = global [2 x i64]* null + +declare void @callee() + +define void @caller1() ssp { +; CHECK-LABEL: define void @caller1() +; Prologue: +; CHECK: @llvm.stackguard + +; CHECK: [[GUARD:%.*]] = call i8* @llvm.stackguard() +; CHECK: [[TOKEN:%.*]] = load volatile i8*, i8** {{%.*}} +; CHECK: [[TST:%.*]] = icmp eq i8* [[GUARD]], [[TOKEN]] +; CHECK: br i1 [[TST]] + +; CHECK: musttail call void @callee() +; CHECK-NEXT: ret void + %var = alloca [2 x i64] + store [2 x i64]* %var, [2 x i64]** @var + musttail call void @callee() + ret void +} + +define void @justret() ssp { +; CHECK-LABEL: define void @justret() +; Prologue: +; CHECK: @llvm.stackguard + +; CHECK: [[GUARD:%.*]] = call i8* @llvm.stackguard() +; CHECK: [[TOKEN:%.*]] = load volatile i8*, i8** {{%.*}} +; CHECK: [[TST:%.*]] = icmp eq i8* [[GUARD]], [[TOKEN]] +; CHECK: br i1 [[TST]] + +; CHECK: ret void + %var = alloca [2 x i64] + store [2 x i64]* %var, [2 x i64]** @var + br label %retblock + +retblock: + ret void +} + + +declare i64* @callee2() + +define i8* @caller2() ssp { +; CHECK-LABEL: define i8* @caller2() +; Prologue: +; CHECK: @llvm.stackguard + +; CHECK: [[GUARD:%.*]] = call i8* @llvm.stackguard() +; CHECK: [[TOKEN:%.*]] = load volatile i8*, i8** {{%.*}} +; CHECK: [[TST:%.*]] = icmp eq i8* [[GUARD]], [[TOKEN]] +; CHECK: br i1 [[TST]] + +; CHECK: [[TMP:%.*]] = musttail call i64* @callee2() +; CHECK-NEXT: [[RES:%.*]] = bitcast i64* [[TMP]] to i8* +; CHECK-NEXT: ret i8* [[RES]] + + %var = alloca [2 x i64] + store [2 x i64]* %var, [2 x i64]** @var + %tmp = musttail call i64* @callee2() + %res = bitcast i64* %tmp to i8* + ret i8* %res +} Index: llvm/test/CodeGen/ARM/Windows/stack-protector-musttail.ll =================================================================== --- /dev/null +++ llvm/test/CodeGen/ARM/Windows/stack-protector-musttail.ll @@ -0,0 +1,56 @@ +; RUN: llc -mtriple=thumbv7-windows-msvc -fast-isel %s -o - -start-before=stack-protector -stop-after=stack-protector | FileCheck %s + +@var = global [2 x i64]* null + +declare void @callee() + +define void @caller1() sspreq { +; CHECK-LABEL: define void @caller1() +; Prologue: + +; CHECK: call void @__security_check_cookie + +; CHECK: musttail call void @callee() +; CHECK-NEXT: ret void + %var = alloca [2 x i64] + store [2 x i64]* %var, [2 x i64]** @var + musttail call void @callee() + ret void +} + +define void @justret() sspreq { +; CHECK-LABEL: define void @justret() +; Prologue: +; CHECK: @llvm.stackguard + +; CHECK: call void @__security_check_cookie + +; CHECK: ret void + %var = alloca [2 x i64] + store [2 x i64]* %var, [2 x i64]** @var + br label %retblock + +retblock: + ret void +} + + +declare i64* @callee2() + +define i8* @caller2() sspreq { +; CHECK-LABEL: define i8* @caller2() +; Prologue: +; CHECK: @llvm.stackguard + +; CHECK: call void @__security_check_cookie + +; CHECK: [[TMP:%.*]] = musttail call i64* @callee2() +; CHECK-NEXT: [[RES:%.*]] = bitcast i64* [[TMP]] to i8* +; CHECK-NEXT: ret i8* [[RES]] + + %var = alloca [2 x i64] + store [2 x i64]* %var, [2 x i64]** @var + %tmp = musttail call i64* @callee2() + %res = bitcast i64* %tmp to i8* + ret i8* %res +}