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

Summary:
This patch comes from H.J.'s hjl-tools@2bd54ce

**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

Subscribers: tstellar, efriedma, hiraditya, llvm-commits

Tags: #llvm

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

Author: xiangzh1

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

llvm/lib/Target/X86/X86TargetMachine.cpp

@@ -214,7 +214,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()) {

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 @@ class X86TargetMachine final : public LLVMTargetMachine {
   TargetLoweringObjectFile *getObjFileLowering() const override {
     return TLOF.get();
   }
+
+  bool isJIT() const { return IsJIT; }
 };
 
 } // end namespace llvm

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}

llvm/test/ExecutionEngine/MCJIT/cet-code-model-lager.ll

@@ -0,0 +1,33 @@
+; RUN: %lli -code-model=large %s > /dev/null
+target triple = "x86_64-unknown-linux-gnu"
+
+@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}