[X86] [CET] Deal with return-twice function such as vfork, setjmp when

CET-IBT enabled

Return-twice functions will indirectly jump after the caller's position.
So when CET-IBT is enable, we should make sure these is endbr*
instructions follow these Return-twice function caller. Like GCC does.

Patch by Xiang Zhang (xiangzhangllvm)

Differential Revision: https://reviews.llvm.org/D61881

llvm-svn: 361342

lib/Target/X86/X86IndirectBranchTracking.cpp

@@ -57,7 +57,7 @@ private:
   /// The function will not add it if already exists.
   /// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
   /// \returns true if the ENDBR was added and false otherwise.
-  bool addENDBR(MachineBasicBlock &MBB) const;
+  bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const;
 };
 
 } // end anonymous namespace
@@ -68,20 +68,31 @@ FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
   return new X86IndirectBranchTrackingPass();
 }
 
-bool X86IndirectBranchTrackingPass::addENDBR(MachineBasicBlock &MBB) const {
+bool X86IndirectBranchTrackingPass::addENDBR(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const {
   assert(TII && "Target instruction info was not initialized");
   assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
          "Unexpected Endbr opcode");
 
-  auto MI = MBB.begin();
-  // If the MBB is empty or the first instruction is not ENDBR,
-  // add the ENDBR instruction to the beginning of the MBB.
-  if (MI == MBB.end() || EndbrOpcode != MI->getOpcode()) {
-    BuildMI(MBB, MI, MBB.findDebugLoc(MI), TII->get(EndbrOpcode));
-    NumEndBranchAdded++;
+  // If the MBB/I is empty or the current instruction is not ENDBR,
+  // insert ENDBR instruction to the location of I.
+  if (I == MBB.end() || I->getOpcode() != EndbrOpcode) {
+    BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode));
+    ++NumEndBranchAdded;
     return true;
   }
+  return false;
+}
 
+bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
+  if (!MOp.isGlobal())
+    return false;
+  auto *CalleeFn = dyn_cast<Function>(MOp.getGlobal());
+  if (!CalleeFn)
+    return false;
+  AttributeList Attrs = CalleeFn->getAttributes();
+  if (Attrs.hasAttribute(AttributeList::FunctionIndex, Attribute::ReturnsTwice))
+    return true;
   return false;
 }
 
@@ -107,14 +118,21 @@ bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
        !MF.getFunction().hasLocalLinkage()) &&
       !MF.getFunction().doesNoCfCheck()) {
     auto MBB = MF.begin();
-    Changed |= addENDBR(*MBB);
+    Changed |= addENDBR(*MBB, MBB->begin());
   }
 
-  for (auto &MBB : MF)
+  for (auto &MBB : MF) {
     // Find all basic blocks that their address was taken (for example
     // in the case of indirect jump) and add ENDBR instruction.
     if (MBB.hasAddressTaken())
-      Changed |= addENDBR(MBB);
+      Changed |= addENDBR(MBB, MBB.begin());
 
+    for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
+      if (!I->isCall())
+        continue;
+      if (IsCallReturnTwice(I->getOperand(0)))
+        Changed |= addENDBR(MBB, std::next(I));
+    }
+  }
   return Changed;
 }

test/CodeGen/X86/indirect-branch-tracking-r2.ll

@@ -0,0 +1,76 @@
+; RUN: llc -mtriple=x86_64-unknown-unknown < %s | FileCheck %s --check-prefix=X64
+; RUN: llc -mtriple=i386-unknown-unknown < %s | FileCheck %s   --check-prefix=X86
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; foo
+;; -----
+;; Checks ENDBR insertion after return twice functions.
+;; setjmp, sigsetjmp, savectx, vfork, getcontext
+;; setzx is not return twice function, should not followed by endbr.
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+; X64:           callq   setjmp
+; X64-NEXT:      endbr64
+; X64:           callq   setzx
+; X64-NEXT:      xorl
+; X64:           callq   sigsetjmp
+; X64-NEXT:      endbr64
+; X64:           callq   savectx
+; X64-NEXT:      endbr64
+; X64:           callq   vfork
+; X64-NEXT:      endbr64
+; X64:           callq   getcontext
+; X64-NEXT:      endbr64
+
+; X86:           calll   setjmp
+; X86-NEXT:      endbr32
+; X86:           calll   setzx
+; X86-NEXT:      xorl
+; X86:           calll   sigsetjmp
+; X86-NEXT:      endbr32
+; X86:           calll   savectx
+; X86-NEXT:      endbr32
+; X86:           calll   vfork
+; X86-NEXT:      endbr32
+; X86:           calll   getcontext
+; X86-NEXT:      endbr32
+
+; Function Attrs: noinline nounwind optnone uwtable
+define dso_local void @foo() #0 {
+entry:
+  %call = call i32 (i32, ...) bitcast (i32 (...)* @setjmp to i32 (i32, ...)*)(i32 0) #1
+  %call1 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0)
+  %call2 = call i32 (i32, ...) bitcast (i32 (...)* @sigsetjmp to i32 (i32, ...)*)(i32 0) #1
+  %call3 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0)
+  %call4 = call i32 (i32, ...) bitcast (i32 (...)* @savectx to i32 (i32, ...)*)(i32 0) #1
+  %call5 = call i32 @vfork() #1
+  %call6 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0)
+  %call7 = call i32 (i32, ...) bitcast (i32 (...)* @getcontext to i32 (i32, ...)*)(i32 0) #1
+  ret void
+}
+
+; Function Attrs: returns_twice
+declare dso_local i32 @setjmp(...) #1
+
+declare dso_local i32 @setzx(...)
+
+; Function Attrs: returns_twice
+declare dso_local i32 @sigsetjmp(...) #1
+
+; Function Attrs: returns_twice
+declare dso_local i32 @savectx(...) #1
+
+; Function Attrs: returns_twice
+declare dso_local i32 @vfork() #1
+
+; Function Attrs: returns_twice
+declare dso_local i32 @getcontext(...) #1
+
+attributes #0 = { noinline nounwind optnone uwtable }
+attributes #1 = { returns_twice }
+
+!llvm.module.flags = !{!0, !1, !2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 4, !"cf-protection-return", i32 1}
+!2 = !{i32 4, !"cf-protection-branch", i32 1}