Enable IBT(Indirect Branch Tracking) in JIT with CET(Control-flow Enforcement Technology)

Do not commit the llvm/test/ExecutionEngine/MCJIT/cet-code-model-lager.ll because it will
cause build bot fail(not suitable for window 32 target).

Summary:
This patch comes from H.J.'s https://github.com/hjl-tools/llvm-project/commit/2bd54ce7fa9e94fcd1118b948e14d1b6fc54dfd2

**This patch fix the failed llvm unit tests which running on CET machine. **(e.g. ExecutionEngine/MCJIT/MCJITTests)

The reason we enable IBT at "JIT compiled with CET" is mainly that:  the JIT don't know the its caller program is CET enable or not.
If JIT's caller program is non-CET, it is no problem JIT generate CET code or not.
But if JIT's caller program is CET enabled,  JIT must generate CET code or it will cause Control protection exceptions.

I have test the patch at llvm-unit-test and llvm-test-suite at CET machine. It passed.
and H.J. also test it at building and running VNCserver(Virtual Network Console), it works too.
(if not apply this patch, VNCserver will crash at CET machine.)

Reviewers: hjl.tools, craig.topper, LuoYuanke, annita.zhang, pengfei

Reviewed By: LuoYuanke

Subscribers: tstellar, efriedma, hiraditya, llvm-commits

Tags: #llvm

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

(cherry picked from commit 01a32f2bd3fad4331cebe9f4faa270d7c082d281)

4 files changed, 59 insertions, 8 deletions

diff --git a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
index 743df41a00f..0a79b793a98 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"
@@ -102,7 +103,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 +121,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();
@@ -136,8 +147,8 @@ bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
         Changed |= addENDBR(MBB, std::next(I));
 
       if (EHPadIBTNeeded && I->isEHLabel()) {
-          Changed |= addENDBR(MBB, std::next(I));
-          EHPadIBTNeeded = false;
+        Changed |= addENDBR(MBB, std::next(I));
+        EHPadIBTNeeded = false;
       }
     }
   }
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index 8c696e9adbe..7176e46f07b 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -222,7 +222,7 @@ X86TargetMachine::X86TargetMachine(const Target &T, const Triple &TT,
           getEffectiveRelocModel(TT, JIT, RM),
           getEffectiveX86CodeModel(CM, JIT, TT.getArch() == Triple::x86_64),
           OL),
-      TLOF(createTLOF(getTargetTriple())) {
+      TLOF(createTLOF(getTargetTriple())), IsJIT(JIT) {
   // On PS4, the "return address" of a 'noreturn' call must still be within
   // the calling function, and TrapUnreachable is an easy way to get that.
   if (TT.isPS4() || TT.isOSBinFormatMachO()) {
diff --git a/llvm/lib/Target/X86/X86TargetMachine.h b/llvm/lib/Target/X86/X86TargetMachine.h
index ec3db7b1e9e..757ce8bc5c7 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.h
+++ b/llvm/lib/Target/X86/X86TargetMachine.h
@@ -30,6 +30,8 @@ class TargetTransformInfo;
 class X86TargetMachine final : public LLVMTargetMachine {
   std::unique_ptr<TargetLoweringObjectFile> TLOF;
   mutable StringMap<std::unique_ptr<X86Subtarget>> SubtargetMap;
+  // True if this is used in JIT.
+  bool IsJIT;
 
 public:
   X86TargetMachine(const Target &T, const Triple &TT, StringRef CPU,
@@ -52,6 +54,8 @@ public:
   TargetLoweringObjectFile *getObjFileLowering() const override {
     return TLOF.get();
   }
+
+  bool isJIT() const { return IsJIT; }
 };
 
 } // end namespace llvm
diff --git a/llvm/test/CodeGen/X86/indirect-branch-tracking-cm-lager.ll b/llvm/test/CodeGen/X86/indirect-branch-tracking-cm-lager.ll
new file mode 100644
index 00000000000..f515c1a34c7
--- /dev/null
+++ b/llvm/test/CodeGen/X86/indirect-branch-tracking-cm-lager.ll
@@ -0,0 +1,36 @@
+; RUN: llc -mtriple=x86_64-unknown-unknown -code-model=large < %s | FileCheck %s
+
+; In large code model indirect branches are needed when branching to addresses
+; whose offset from the current instruction pointer is unknown.
+;CHECK-COUNT-3: endbr
+
+@a = dso_local local_unnamed_addr global i32 1, align 4
+
+; Function Attrs: nofree noinline norecurse nounwind uwtable writeonly
+define dso_local void @ext() local_unnamed_addr #0 {
+entry:
+  store i32 0, i32* @a, align 4
+  ret void
+}
+
+; Function Attrs: nofree norecurse nounwind uwtable
+define dso_local i32 @main() local_unnamed_addr #1 {
+entry:
+  tail call fastcc void @foo()
+  %0 = load i32, i32* @a, align 4
+  ret i32 %0
+}
+
+; Function Attrs: nofree noinline norecurse nounwind uwtable writeonly
+define internal fastcc void @foo() unnamed_addr #0 {
+entry:
+  tail call void @ext()
+  ret void
+}
+
+!llvm.module.flags = !{!0, !1, !2, !3}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 4, !"cf-protection-return", i32 1}
+!2 = !{i32 4, !"cf-protection-branch", i32 1}
+!3 = !{i32 1, !"Code Model", i32 4}