StackProtector: insert check before `musttail` call.

[t.p.northover]

Summary

musttail calls have to directly precede the return instruction (of their value), so any stack checks must be inserted before the call instead of in between it and the return. Fortunately this same restriction means we don't have to look too hard to find out if we're about to break a musttail call.

[t.p.northover]

We've just discovered there are problems on the SDAG path too for guaranteed tail calls. I'll update the patch once I've got a fix ready.

[t.p.northover]

Fixed the issue in SelectionDAG too. Not quite identical there, but still fundamentally hinged on exactly where the check should be placed.

[aemerson]

LGTM with comment nit.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

1713 ↗

(On Diff #331239)

An example of what code you're looking for in this comment would be helpful.

[t.p.northover]

Thanks Amara. Committed with a beefed up comment as 5e3d9fcc3a88.

llvm/lib/CodeGen/StackProtector.cpp

if (!AI) {
       AI = cast<AllocaInst>(SPCall->getArgOperand(1));
     }
 
     // Set HasIRCheck to true, so that SelectionDAG will not generate its own
     // version. SelectionDAG called 'shouldEmitSDCheck' to check whether
     // 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:
       //
       //   return:
       //     ...
       //     ret ...
       //
       // into this:
if (Function *GuardCheck = TLI->getSSPStackGuardCheck(*M)) {
       //     unreachable
 
       // Create the FailBB. We duplicate the BB every time since the MI tail
       // merge pass will merge together all of the various BB into one including
       // fail BB generated by the stack protector pseudo instruction.
       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)) {
         DT->addNewBlock(NewBB, BB);
         DT->addNewBlock(FailBB, BB);
       }
 
       // Remove default branch instruction to the new BB.

llvm/test/CodeGen/AArch64/stack-protector-musttail.ll

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

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