diff --git a/clang/include/clang/Basic/CodeGenOptions.def b/clang/include/clang/Basic/CodeGenOptions.def --- a/clang/include/clang/Basic/CodeGenOptions.def +++ b/clang/include/clang/Basic/CodeGenOptions.def @@ -107,6 +107,8 @@ ///< set to full or return. CODEGENOPT(CFProtectionBranch , 1, 0) ///< if -fcf-protection is ///< set to full or branch. +CODEGENOPT(IBTSeal, 1, 0) ///< set to optimize CFProtectionBranch. + CODEGENOPT(XRayInstrumentFunctions , 1, 0) ///< Set when -fxray-instrument is ///< enabled. CODEGENOPT(StackSizeSection , 1, 0) ///< Set when -fstack-size-section is enabled. diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td --- a/clang/include/clang/Driver/Options.td +++ b/clang/include/clang/Driver/Options.td @@ -1927,6 +1927,8 @@ def fcf_protection : Flag<["-"], "fcf-protection">, Group, Flags<[CoreOption, CC1Option]>, Alias, AliasArgs<["full"]>, HelpText<"Enable cf-protection in 'full' mode">; +def mibt_seal : Flag<["-"], "mibt-seal">, Group, Flags<[CoreOption, CC1Option]>, + HelpText<"Optimize fcf-protection=branch/full (requires LTO).">; defm xray_instrument : BoolFOption<"xray-instrument", LangOpts<"XRayInstrument">, DefaultFalse, diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp --- a/clang/lib/CodeGen/CodeGenModule.cpp +++ b/clang/lib/CodeGen/CodeGenModule.cpp @@ -712,6 +712,9 @@ 1); } + if (CodeGenOpts.IBTSeal) + getModule().addModuleFlag(llvm::Module::Override, "ibt-seal", 1); + // Add module metadata for return address signing (ignoring // non-leaf/all) and stack tagging. These are actually turned on by function // attributes, but we use module metadata to emit build attributes. This is diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp --- a/clang/lib/Driver/ToolChains/Clang.cpp +++ b/clang/lib/Driver/ToolChains/Clang.cpp @@ -6166,6 +6166,9 @@ Args.MakeArgString(Twine("-fcf-protection=") + A->getValue())); } + if (IsUsingLTO) + Args.AddLastArg(CmdArgs, options::OPT_mibt_seal); + // Forward -f options with positive and negative forms; we translate these by // hand. Do not propagate PGO options to the GPU-side compilations as the // profile info is for the host-side compilation only. diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp --- a/clang/lib/Frontend/CompilerInvocation.cpp +++ b/clang/lib/Frontend/CompilerInvocation.cpp @@ -1814,6 +1814,9 @@ Diags.Report(diag::err_drv_invalid_value) << A->getAsString(Args) << Name; } + if (Opts.PrepareForLTO && Args.hasArg(OPT_mibt_seal)) + Opts.IBTSeal = 1; + for (auto *A : Args.filtered(OPT_mlink_bitcode_file, OPT_mlink_builtin_bitcode)) { CodeGenOptions::BitcodeFileToLink F; diff --git a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp --- a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp +++ b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp @@ -95,14 +95,45 @@ return Attrs.hasFnAttr(Attribute::ReturnsTwice); } +// Checks if function should have an ENDBR in its prologue +static bool needsPrologueENDBR(MachineFunction &MF, const Module *M) { + Function &F = MF.getFunction(); + + if (F.doesNoCfCheck()) + return false; + + const X86TargetMachine *TM = + static_cast(&MF.getTarget()); + Metadata *IBTSeal = M->getModuleFlag("ibt-seal"); + + switch (TM->getCodeModel()) { + // Large code model functions always reachable through indirect calls. + case CodeModel::Large: + return true; + // Only address taken functions in LTO'ed kernel are reachable indirectly. + // IBTSeal implies LTO, thus only check if function is address taken. + case CodeModel::Kernel: + // Check if ibt-seal was enabled (implies LTO is being used). + if (IBTSeal) { + return F.hasAddressTaken(); + } + // if !IBTSeal, fall into default case. + LLVM_FALLTHROUGH; + // Address taken or externally linked functions may be reachable. + default: + return (F.hasAddressTaken() || !F.hasLocalLinkage()); + } +} + bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) { const X86Subtarget &SubTarget = MF.getSubtarget(); + const Module *M = MF.getMMI().getModule(); // Check that the cf-protection-branch is enabled. - Metadata *isCFProtectionSupported = - MF.getMMI().getModule()->getModuleFlag("cf-protection-branch"); - // NB: We need to enable IBT in jitted code if JIT compiler is CET - // enabled. + Metadata *isCFProtectionSupported = M->getModuleFlag("cf-protection-branch"); + + // NB: We need to enable IBT in jitted code if JIT compiler is CET + // enabled. const X86TargetMachine *TM = static_cast(&MF.getTarget()); #ifdef __CET__ @@ -119,13 +150,8 @@ TII = SubTarget.getInstrInfo(); EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32; - // 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()) { + // If function is reachable indirectly, mark the first BB with ENDBR. + if (needsPrologueENDBR(MF, M)) { auto MBB = MF.begin(); Changed |= addENDBR(*MBB, MBB->begin()); } diff --git a/llvm/test/CodeGen/X86/ibtseal-kernel.ll b/llvm/test/CodeGen/X86/ibtseal-kernel.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/X86/ibtseal-kernel.ll @@ -0,0 +1,19 @@ +; RUN: llc < %s -O2 -mtriple=x86_64-unknown-linux-gnu -x86-indirect-branch-tracking --code-model=kernel | FileCheck %s --check-prefix=CHECK-KERNEL-IBTSEAL + +; CHECK-KERNEL-IBTSEAL: foo: +; CHECK-KERNEL-IBTSEAL: endbr +; CHECK-KERNEL-IBTSEAL: bar: +; CHECK-KERNEL-IBTSEAL-NOT: endbr + +target triple = "x86_64-unknown-linux-gnu" + +define dso_local void @foo() { + ret void +} + +define dso_local i8* @bar() { + ret i8* bitcast (void ()* @foo to i8*) +} + +!llvm.module.flags = !{!1} +!1 = !{i32 4, !"ibt-seal", i32 1} diff --git a/llvm/test/CodeGen/X86/ibtseal-large.ll b/llvm/test/CodeGen/X86/ibtseal-large.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/X86/ibtseal-large.ll @@ -0,0 +1,19 @@ +; RUN: llc < %s -O2 -mtriple=x86_64-unknown-linux-gnu -x86-indirect-branch-tracking --code-model=large | FileCheck %s --check-prefix=CHECK-LARGE-IBTSEAL + +; CHECK-LARGE-IBTSEAL: foo: +; CHECK-LARGE-IBTSEAL: endbr +; CHECK-LARGE-IBTSEAL: bar: +; CHECK-LARGE-IBTSEAL: endbr + +target triple = "x86_64-unknown-linux-gnu" + +define dso_local void @foo() { + ret void +} + +define dso_local i8* @bar() { + ret i8* bitcast (void ()* @foo to i8*) +} + +!llvm.module.flags = !{!1} +!1 = !{i32 4, !"ibt-seal", i32 1} diff --git a/llvm/test/CodeGen/X86/ibtseal-small.ll b/llvm/test/CodeGen/X86/ibtseal-small.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/X86/ibtseal-small.ll @@ -0,0 +1,19 @@ +; RUN: llc < %s -O2 -mtriple=x86_64-unknown-linux-gnu -x86-indirect-branch-tracking --code-model=small | FileCheck %s --check-prefix=CHECK-SMALL-IBTSEAL + +; CHECK-SMALL-IBTSEAL: foo: +; CHECK-SMALL-IBTSEAL: endbr +; CHECK-SMALL-IBTSEAL: bar: +; CHECK-SMALL-IBTSEAL: endbr + +target triple = "x86_64-unknown-linux-gnu" + +define dso_local void @foo() { + ret void +} + +define dso_local i8* @bar() { + ret i8* bitcast (void ()* @foo to i8*) +} + +!llvm.module.flags = !{!1} +!1 = !{i32 4, !"ibt-seal", i32 1}