Deal with return-twice function such as vfork, setjmp when CET-IBT enabled

[xiangzhangllvm]

Summary

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.

[MaskRay]

lib/Target/X86/X86IndirectBranchTracking.cpp

72

indent with clang-format.

78

++I;

139

for (const MachineInstr &MI : MBB)

[xiangzhangllvm]

lib/Target/X86/X86IndirectBranchTracking.cpp

72

Hi! Fang rui, Nice to see you here!
It will more than 80 chars if I aligned the "Machine..", Is it suitable to write like this: ?

bool
X86IndirectBranchTrackingPass::addENDBR(MachineBasicBlock &MBB,
                                        MachineBasicBlock::iterator I) const {

78

It seems same with ++I ? More performance ?

139

I used iterator I, because I want to pass I into BuildMI in addENDBR()

[MaskRay]

lib/Target/X86/X86IndirectBranchTracking.cpp

72

You may use clang-format (git diff -U0 --no-color 'HEAD^' | ~/llvm/tools/clang/tools/clang-format/clang-format-diff.py -i -p1). It reformats this block as:

bool X86IndirectBranchTrackingPass::addENDBR(
    MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const {

78

Its https://llvm.org/docs/CodingStandards.html#prefer-preincrement

[xiangzhangllvm]

Thank You very much! Fangrui

[craig.topper]

lib/Target/X86/X86IndirectBranchTracking.cpp

76–84

Is this really "I != MBB.end()"? MachineBasicBlock::iterator is an iterator into a linked list and I'm not sure what it means for that to be nullptr.

141

Variable names should be capitalized and use camel case. But I think this variable is unnecessary and you could just call I->isCall from the if.

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

10

funtion->function

70

Please reduce this to only the attributes necessary. target-cpu, no-trapping-math, etc. are all unnecessary.

[xiangzhangllvm]

lib/Target/X86/X86IndirectBranchTracking.cpp

76–84

Here we just want to use a special "value" to let the addENDBR() ignore the argument "I".
Yes, "I" should "point to" an MachineInstr, it should not be MBB.end() too.
I also feel it a little ugly here, but compile with no error. I'll rewrite to if (I != MBB.end()).
Thank you!

141

Yes! make sense!

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

10

Yes!

70

OK! I'll simply them. Thank you!

[craig.topper]

lib/Target/X86/X86IndirectBranchTracking.cpp

144

What if we passed I++ here and passed MBB.begin() on the other callsites. Could we then avoid two different cases in addENDBR?

[xiangzhangllvm]

lib/Target/X86/X86IndirectBranchTracking.cpp

139

I find here line 139 maybe problem:
I should change to

for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end();  //because the MBB may update by insert endbr.

144

Hi Craig, we can not avoid.
because the 2 cases in addENDBR(MBB, I) are different.

one case is insert endbr after I
and another case is insert endbr before I

[xiangzhangllvm]

lib/Target/X86/X86IndirectBranchTracking.cpp

144

Sorry, I miss understand before. Your suggestion is right!
Here can pass "I+1" not "I++" to addENDBR()

[craig.topper]

LGTM

lib/Target/X86/X86IndirectBranchTracking.cpp

private:
 
   /// Endbr opcode for the current machine function.
   unsigned int EndbrOpcode;
 
   /// Adds a new ENDBR instruction to the begining of the MBB.
   /// 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
 
 char X86IndirectBranchTrackingPass::ID = 0;
 
 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");
-
+  if (I != MBB.end()) {
+    if (EndbrOpcode != I->getOpcode()) {
+      ++I;
+      BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode));
+      ++NumEndBranchAdded;
+      return true;
+    }
+    return false;
+  }
   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++;
+    ++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;
 }
 
 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
   const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
 
   // Check that the cf-protection-branch is enabled.
   Metadata *isCFProtectionSupported =
bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
 
   // Non-internal function or function whose address was taken, can be
   // accessed through indirect calls. Mark the first BB with ENDBR instruction
   // unless nocf_check attribute is used.
   if ((MF.getFunction().hasAddressTaken() ||
        !MF.getFunction().hasLocalLinkage()) &&
       !MF.getFunction().doesNoCfCheck()) {
     auto MBB = MF.begin();
-    Changed |= addENDBR(*MBB);
+    Changed |= addENDBR(*MBB, MBB->end());
   }
 
-  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.end());
 
+    for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); I != E;
+         ++I) {
+      if (!I->isCall())
+        continue;
+      if (IsCallReturnTwice(I->getOperand(0)))
+        Changed |= addENDBR(MBB, I);
+    }
+  }
   return Changed;
 }

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

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