StackProtector: ensure protection does not interfere with tail call frame.

The IR stack protector pass must insert stack checks before the call instead of
between it and the return.

Similarly, SDAG one should recognize that ADJCALLFRAME instructions could be
part of the terminal sequence of a tail call. In this case because such call
frames cannot be nested in LLVM the stack protection code must skip over the
whole sequence (or risk clobbering argument registers).

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

@@ -1694,9 +1694,9 @@ static bool MIIsInTerminatorSequence(const MachineInstr &MI) {
 /// terminator, but additionally the copies that move the vregs into the
 /// physical registers.
 static MachineBasicBlock::iterator
-FindSplitPointForStackProtector(MachineBasicBlock *BB) {
+FindSplitPointForStackProtector(MachineBasicBlock *BB,
+                                const TargetInstrInfo &TII) {
   MachineBasicBlock::iterator SplitPoint = BB->getFirstTerminator();
-  //
   if (SplitPoint == BB->begin())
     return SplitPoint;
 
@@ -1704,6 +1704,31 @@ FindSplitPointForStackProtector(MachineBasicBlock *BB) {
   MachineBasicBlock::iterator Previous = SplitPoint;
   --Previous;
 
+  if (TII.isTailCall(*SplitPoint) &&
+      Previous->getOpcode() == TII.getCallFrameDestroyOpcode()) {
+    // call itself, then we must insert before the sequence even starts. For
+    // example:
+    //     <split point>
+    //     ADJCALLSTACKDOWN ...
+    //     <Moves>
+    //     ADJCALLSTACKUP ...
+    //     TAILJMP somewhere
+    // On the other hand, it could be an unrelated call in which case this tail call
+    // has to register moves of its own and should be the split point. For example:
+    //     ADJCALLSTACKDOWN
+    //     CALL something_else
+    //     ADJCALLSTACKUP
+    //     <split point>
+    //     TAILJMP somewhere
+    do {
+      --Previous;
+      if (Previous->isCall())
+        return SplitPoint;
+    } while(Previous->getOpcode() != TII.getCallFrameSetupOpcode());
+
+    return Previous;
+  }
+
   while (MIIsInTerminatorSequence(*Previous)) {
     SplitPoint = Previous;
     if (Previous == Start)
@@ -1743,7 +1768,7 @@ SelectionDAGISel::FinishBasicBlock() {
     // Add load and check to the basicblock.
     FuncInfo->MBB = ParentMBB;
     FuncInfo->InsertPt =
-        FindSplitPointForStackProtector(ParentMBB);
+        FindSplitPointForStackProtector(ParentMBB, *TII);
     SDB->visitSPDescriptorParent(SDB->SPDescriptor, ParentMBB);
     CurDAG->setRoot(SDB->getRoot());
     SDB->clear();
@@ -1762,7 +1787,7 @@ SelectionDAGISel::FinishBasicBlock() {
     // register allocation issues caused by us splitting the parent mbb. The
     // register allocator will clean up said virtual copies later on.
     MachineBasicBlock::iterator SplitPoint =
-        FindSplitPointForStackProtector(ParentMBB);
+        FindSplitPointForStackProtector(ParentMBB, *TII);
 
     // Splice the terminator of ParentMBB into SuccessMBB.
     SuccessMBB->splice(SuccessMBB->end(), ParentMBB,

llvm/lib/CodeGen/StackProtector.cpp

@@ -471,21 +471,36 @@ bool StackProtector::InsertStackProtectors() {
     // 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<CallInst>(Prev) && cast<CallInst>(Prev)->isMustTailCall())
+      CheckLoc = Prev;
+    else if (Prev) {
+      Prev = Prev->getPrevNonDebugInstruction();
+      if (Prev && isa<CallInst>(Prev) && cast<CallInst>(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 @@ bool StackProtector::InsertStackProtectors() {
       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)) {

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
+}

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
+}

llvm/test/CodeGen/X86/tailcc-ssp.ll

@@ -0,0 +1,26 @@
+; RUN: llc -mtriple=x86_64-windows-msvc %s -o - -verify-machineinstrs | FileCheck %s
+
+declare void @h(i8*, i64, i8*)
+
+define tailcc void @tailcall_frame(i8* %0, i64 %1) sspreq {
+; CHECK-LABEL: tailcall_frame:
+; CHECK: callq __security_check_cookie
+; CHECK: xorl %ecx, %ecx
+; CHECK: jmp h
+
+   tail call tailcc void @h(i8* null, i64 0, i8* null)
+   ret void
+}
+
+declare void @bar()
+define void @tailcall_unrelated_frame() sspreq {
+; CHECK-LABEL: tailcall_unrelated_frame:
+; CHECK: subq [[STACK:\$.*]], %rsp
+; CHECK: callq bar
+; CHECK: callq __security_check_cookie
+; CHECK: addq [[STACK]], %rsp
+; CHECK: jmp bar
+  call void @bar()
+  tail call void @bar()
+  ret void
+}