1 files changed, 48 insertions, 11 deletions

diff --git a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
index 48d0d8a35704..1628f85da808 100644
--- a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
+++ b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
@@ -18,6 +18,7 @@
 #include "X86.h"
 #include "X86InstrInfo.h"
 #include "X86Subtarget.h"
+#include "X86TargetMachine.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
@@ -91,9 +92,7 @@ static bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
   if (!CalleeFn)
     return false;
   AttributeList Attrs = CalleeFn->getAttributes();
-  if (Attrs.hasAttribute(AttributeList::FunctionIndex, Attribute::ReturnsTwice))
-    return true;
-  return false;
+  return Attrs.hasFnAttribute(Attribute::ReturnsTwice);
 }
 
 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
@@ -102,7 +101,16 @@ bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
   // Check that the cf-protection-branch is enabled.
   Metadata *isCFProtectionSupported =
       MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
-  if (!isCFProtectionSupported && !IndirectBranchTracking)
+  // NB: We need to enable IBT in jitted code if JIT compiler is CET
+  // enabled.
+  const X86TargetMachine *TM =
+      static_cast<const X86TargetMachine *>(&MF.getTarget());
+#ifdef __CET__
+  bool isJITwithCET = TM->isJIT();
+#else
+  bool isJITwithCET = false;
+#endif
+  if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET)
     return false;
 
   // True if the current MF was changed and false otherwise.
@@ -111,10 +119,11 @@ bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
   TII = SubTarget.getInstrInfo();
   EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
 
-  // 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() ||
+  // Large code model, 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 ((TM->getCodeModel() == CodeModel::Large ||
+       MF.getFunction().hasAddressTaken() ||
        !MF.getFunction().hasLocalLinkage()) &&
       !MF.getFunction().doesNoCfCheck()) {
     auto MBB = MF.begin();
@@ -128,10 +137,38 @@ bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
       Changed |= addENDBR(MBB, MBB.begin());
 
     for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
-      if (!I->isCall())
-        continue;
-      if (IsCallReturnTwice(I->getOperand(0)))
+      if (I->isCall() && IsCallReturnTwice(I->getOperand(0)))
+        Changed |= addENDBR(MBB, std::next(I));
+    }
+
+    // Exception handle may indirectly jump to catch pad, So we should add
+    // ENDBR before catch pad instructions. For SjLj exception model, it will
+    // create a new BB(new landingpad) indirectly jump to the old landingpad.
+    if (TM->Options.ExceptionModel == ExceptionHandling::SjLj) {
+      for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
+        // New Landingpad BB without EHLabel.
+        if (MBB.isEHPad()) {
+          if (I->isDebugInstr())
+            continue;
+          Changed |= addENDBR(MBB, I);
+          break;
+        } else if (I->isEHLabel()) {
+          // Old Landingpad BB (is not Landingpad now) with
+          // the the old "callee" EHLabel.
+          MCSymbol *Sym = I->getOperand(0).getMCSymbol();
+          if (!MF.hasCallSiteLandingPad(Sym))
+            continue;
+          Changed |= addENDBR(MBB, std::next(I));
+          break;
+        }
+      }
+    } else if (MBB.isEHPad()){
+      for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
+        if (!I->isEHLabel())
+          continue;
         Changed |= addENDBR(MBB, std::next(I));
+        break;
+      }
     }
   }
   return Changed;