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(&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; } } } 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 TLOF; mutable StringMap> 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()) {