[PATCH] D76900: Enable IBT(Indirect Branch Tracking) in JIT with CET(Control-flow Enforcement Technology)

[Xiang Zhang via Phabricator via llvm-commits]

xiangzhangllvm updated this revision to Diff 253035.

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D76900/new/

https://reviews.llvm.org/D76900

Files:
  llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
  llvm/lib/Target/X86/X86TargetMachine.cpp
  llvm/lib/Target/X86/X86TargetMachine.h


Index: llvm/lib/Target/X86/X86TargetMachine.h
===================================================================
--- llvm/lib/Target/X86/X86TargetMachine.h
+++ llvm/lib/Target/X86/X86TargetMachine.h
@@ -30,6 +30,8 @@
 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 @@
   TargetLoweringObjectFile *getObjFileLowering() const override {
     return TLOF.get();
   }
+
+  bool isJIT() const { return IsJIT; }
 };
 
 } // end namespace llvm
Index: llvm/lib/Target/X86/X86TargetMachine.cpp
===================================================================
--- llvm/lib/Target/X86/X86TargetMachine.cpp
+++ llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -214,7 +214,7 @@
           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()) {
Index: llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
===================================================================
--- llvm/lib/Target/X86/X86IndirectBranchTracking.cpp
+++ 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 @@
   // 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.
+#ifdef __CET__
+  const X86TargetMachine *TM =
+      static_cast<const X86TargetMachine *>(&MF.getTarget());
+  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 @@
   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
+  // Jitted code generated by CET-enabled JIT compiler, 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() ||
+  if ((isJITwithCET || MF.getFunction().hasAddressTaken() ||
        !MF.getFunction().hasLocalLinkage()) &&
       !MF.getFunction().doesNoCfCheck()) {
     auto MBB = MF.begin();
@@ -136,8 +147,8 @@
         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;
       }
     }
   }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D76900.253035.patch
Type: text/x-patch
Size: 3612 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200327/ba8b826a/attachment.bin>