Index: clang/include/clang/Basic/CodeGenOptions.def =================================================================== --- clang/include/clang/Basic/CodeGenOptions.def +++ clang/include/clang/Basic/CodeGenOptions.def @@ -109,6 +109,14 @@ ///< set to full or branch. CODEGENOPT(IBTSeal, 1, 0) ///< set to optimize CFProtectionBranch. +CODEGENOPT(IBTFixDirectAggressive, 1, 0)///< Set to skip ENDBRs in targets of + ///< direct calls. Optimizes declaration + ///< targets. + +CODEGENOPT(IBTFixDirectConservative, 1, 0)///< Set to skip ENDBRs in targets of + ///< direct calls. Only optimizes + ///< materializable targets. + CODEGENOPT(XRayInstrumentFunctions , 1, 0) ///< Set when -fxray-instrument is ///< enabled. CODEGENOPT(StackSizeSection , 1, 0) ///< Set when -fstack-size-section is enabled. Index: clang/include/clang/Driver/Options.td =================================================================== --- clang/include/clang/Driver/Options.td +++ clang/include/clang/Driver/Options.td @@ -1932,6 +1932,13 @@ def mibt_seal : Flag<["-"], "mibt-seal">, Group, Flags<[CoreOption, CC1Option]>, HelpText<"Optimize fcf-protection=branch/full (requires LTO).">; +def mibt_fix_direct_EQ : Joined<["-"], "mibt-fix-direct=">, Flags<[CoreOption, CC1Option]>, Group, + HelpText<"Fix direct calls to skip ENDBRs in targets whenever possible. Options: conservative, aggressive, none.">, Values<"conservative,aggressive,none">; + +def mibt_fix_direct : Flag<["-"], "mibt-fix-direct">, Group, Flags<[CoreOption, CC1Option]>, + Alias, AliasArgs<["conservative"]>, + HelpText<"Enable ibt-fix-direct in 'conservative' mode">; + defm xray_instrument : BoolFOption<"xray-instrument", LangOpts<"XRayInstrument">, DefaultFalse, PosFlag, Index: clang/lib/CodeGen/CodeGenModule.cpp =================================================================== --- clang/lib/CodeGen/CodeGenModule.cpp +++ clang/lib/CodeGen/CodeGenModule.cpp @@ -712,8 +712,18 @@ 1); } - if (CodeGenOpts.IBTSeal) + // CET/IBT optimizations + if (CodeGenOpts.IBTSeal) { getModule().addModuleFlag(llvm::Module::Override, "ibt-seal", 1); + } + + if (CodeGenOpts.IBTFixDirectAggressive) { + getModule().addModuleFlag(llvm::Module::Override, "ibt-fix-direct-aggr", 1); + } + + if (CodeGenOpts.IBTFixDirectConservative) { + getModule().addModuleFlag(llvm::Module::Override, "ibt-fix-direct-cons", 1); + } // Add module metadata for return address signing (ignoring // non-leaf/all) and stack tagging. These are actually turned on by function Index: clang/lib/Driver/ToolChains/Clang.cpp =================================================================== --- clang/lib/Driver/ToolChains/Clang.cpp +++ clang/lib/Driver/ToolChains/Clang.cpp @@ -6169,9 +6169,17 @@ Args.MakeArgString(Twine("-fcf-protection=") + A->getValue())); } + // CET/IBT optimizations if (IsUsingLTO) Args.AddLastArg(CmdArgs, options::OPT_mibt_seal); + if (Arg *A = Args.getLastArg(options::OPT_mibt_fix_direct_EQ)) { + CmdArgs.push_back( + Args.MakeArgString(Twine("-mibt-fix-direct=") + A->getValue())); + } + + Args.AddLastArg(CmdArgs, options::OPT_mibt_fix_direct); + // 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. Index: clang/lib/Frontend/CompilerInvocation.cpp =================================================================== --- clang/lib/Frontend/CompilerInvocation.cpp +++ clang/lib/Frontend/CompilerInvocation.cpp @@ -1473,6 +1473,11 @@ else if (Opts.CFProtectionBranch) GenerateArg(Args, OPT_fcf_protection_EQ, "branch", SA); + if (Opts.IBTFixDirectAggressive) + GenerateArg(Args, OPT_mibt_fix_direct_EQ, "aggressive", SA); + else if (Opts.IBTFixDirectConservative) + GenerateArg(Args, OPT_mibt_fix_direct_EQ, "conservative", SA); + for (const auto &F : Opts.LinkBitcodeFiles) { bool Builtint = F.LinkFlags == llvm::Linker::Flags::LinkOnlyNeeded && F.PropagateAttrs && F.Internalize; @@ -1814,6 +1819,20 @@ Diags.Report(diag::err_drv_invalid_value) << A->getAsString(Args) << Name; } + if (const Arg *A = Args.getLastArg(OPT_mibt_fix_direct_EQ)) { + if (!Opts.CFProtectionBranch) { + Diags.Report(diag::err_drv_unsupported_opt) + << "-mibt-fix-direct without -fcf-protection="; + } + StringRef Name = A->getValue(); + if (Name == "conservative") + Opts.IBTFixDirectConservative = 1; + else if (Name == "aggressive") + Opts.IBTFixDirectAggressive = 1; + else if (Name != "none") + Diags.Report(diag::err_drv_invalid_value) << A->getAsString(Args) << Name; + } + if (Opts.PrepareForLTO && Args.hasArg(OPT_mibt_seal)) Opts.IBTSeal = 1; Index: llvm/lib/Target/X86/X86IndirectBranchTracking.cpp =================================================================== --- llvm/lib/Target/X86/X86IndirectBranchTracking.cpp +++ llvm/lib/Target/X86/X86IndirectBranchTracking.cpp @@ -27,6 +27,10 @@ using namespace llvm; #define DEBUG_TYPE "x86-indirect-branch-tracking" +#define ENDBR_LEN 4 +#define WARN \ + StringRef name = MF.getName(); \ + WithColor::warning() cl::opt IndirectBranchTracking( "x86-indirect-branch-tracking", cl::init(false), cl::Hidden, @@ -59,6 +63,15 @@ /// It will add ENDBR32 or ENDBR64 opcode, depending on the target. /// \returns true if the ENDBR was added and false otherwise. bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const; + + /// Checks if the function should get an ENDBR instruction in its prologue. + bool needsPrologueENDBR(const Function *F, const Module *M, + const X86TargetMachine *TM) const; + + /// Adds +4 offset to direct calls that are targeting an ENDBR instruction, + /// preventing control-flow from decoding superfluous instructions. + bool fixDirectCalls(MachineFunction &MF, const Module *M, + bool Aggressive) const; }; } // end anonymous namespace @@ -69,6 +82,16 @@ return new X86IndirectBranchTrackingPass(); } +static bool IsCallReturnTwice(llvm::MachineOperand &MOp) { + if (!MOp.isGlobal()) + return false; + auto *CalleeFn = dyn_cast(MOp.getGlobal()); + if (!CalleeFn) + return false; + AttributeList Attrs = CalleeFn->getAttributes(); + return Attrs.hasFnAttr(Attribute::ReturnsTwice); +} + bool X86IndirectBranchTrackingPass::addENDBR( MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const { assert(TII && "Target instruction info was not initialized"); @@ -85,25 +108,12 @@ return false; } -static bool IsCallReturnTwice(llvm::MachineOperand &MOp) { - if (!MOp.isGlobal()) - return false; - auto *CalleeFn = dyn_cast(MOp.getGlobal()); - if (!CalleeFn) - return false; - AttributeList Attrs = CalleeFn->getAttributes(); - 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()) +bool X86IndirectBranchTrackingPass::needsPrologueENDBR( + const Function *F, const Module *M, const X86TargetMachine *TM) const { + if (F->doesNoCfCheck()) return false; - const X86TargetMachine *TM = - static_cast(&MF.getTarget()); Metadata *IBTSeal = M->getModuleFlag("ibt-seal"); switch (TM->getCodeModel()) { @@ -115,16 +125,69 @@ case CodeModel::Kernel: // Check if ibt-seal was enabled (implies LTO is being used). if (IBTSeal) { - return F.hasAddressTaken(); + 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()); + return (F->hasAddressTaken() || !F->hasLocalLinkage()); } } +bool X86IndirectBranchTrackingPass::fixDirectCalls(MachineFunction &MF, + const Module *M, + bool Aggressive) const { + bool Changed = false; + for (auto &BB : MF) { + for (auto &I : BB) { + unsigned Opcode = I.getOpcode(); + if (Opcode == X86::CALL64pcrel32 || Opcode == X86::TAILJMPd || + Opcode == X86::TAILJMPd64_CC || Opcode == X86::TAILJMPd_CC || + Opcode == X86::TAILJMPd64) { + + auto &O = I.getOperand(0); + if (O.getOffset()) { + LLVM_DEBUG(WARN << "X86/ibt-fix-direct: call to " << name + << " already has an offset\n";); + continue; + } + + // Direct calls to targets that will receive and ENDBR instruction + // on their prologues are incremented with ENDBR_LEN offset. If + // -mibt-fix-direct is set to aggressive, calls to declarations are + // also fixed. If set to conservative, only calls to materializable + // functions are fixed. + if (O.isGlobal()) { + const Function *F = NULL; + const Value *Target = O.getGlobal(); + Target = Target->stripPointerCastsAndAliases(); + if (Target) + F = dyn_cast_or_null(Target); + if (!F) { + LLVM_DEBUG(WARN << "X86/ibt-fix-direct: Unknown call target in " + << name << "\n"); + continue; + } + + if (!Aggressive && F->isDeclaration()) + continue; + + const X86TargetMachine *TM = + static_cast(&MF.getTarget()); + if (needsPrologueENDBR(F, M, TM)) { + LLVM_DEBUG(WARN << "X86/ibt-fix-direct: Direct call in " << name + << " was fixed.\n";); + O.setOffset(ENDBR_LEN); + Changed = true; + } + } + } + } + } + return Changed; +} + bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) { const X86Subtarget &SubTarget = MF.getSubtarget(); @@ -151,11 +214,17 @@ EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32; // If function is reachable indirectly, mark the first BB with ENDBR. - if (needsPrologueENDBR(MF, M)) { + if (needsPrologueENDBR(&MF.getFunction(), M, TM)) { auto MBB = MF.begin(); Changed |= addENDBR(*MBB, MBB->begin()); } + if (M->getModuleFlag("ibt-fix-direct-aggr")) { + Changed |= fixDirectCalls(MF, M, true); + } else if (M->getModuleFlag("ibt-fix-direct-cons")) { + Changed |= fixDirectCalls(MF, M, false); + } + for (auto &MBB : MF) { // Find all basic blocks that their address was taken (for example // in the case of indirect jump) and add ENDBR instruction. Index: llvm/test/CodeGen/X86/ibtfix-direct-aggr.ll =================================================================== --- /dev/null +++ llvm/test/CodeGen/X86/ibtfix-direct-aggr.ll @@ -0,0 +1,44 @@ +; RUN: llc < %s -O2 -mtriple=x86_64-unknown-linux-gnu -x86-indirect-branch-tracking | FileCheck %s --check-prefix=CHECK-IBT-FIX-DIRECT + +; CHECK-IBT-FIX-DIRECT: foo: +; CHECK-IBT-FIX-DIRECT: endbr +; CHECK-IBT-FIX-DIRECT: tester: +; CHECK-IBT-FIX-DIRECT-NOT: endbr +; CHECK-IBT-FIX-DIRECT: callq * +; CHECK-IBT-FIX-DIRECT: callq foo+4 +; CHECK-IBT-FIX-DIRECT: callq bar+4 +; CHECK-IBT-FIX-DIRECT-NOT: call foo_nocf_check+4 +; CHECK-IBT-FIX-DIRECT-NOT: call bar_nocf_check+4 +; CHECK-IBT-FIX-DIRECT-NOT: callq tester+4 + +target triple = "x86_64-unknown-linux-gnu" + +define dso_local void @foo() { + ret void +} + +declare dso_local void @bar() + +define dso_local void @foo_nocf_check() #0 { + ret void +} + +declare dso_local void @bar_nocf_check() #0 + +attributes #0 = { nocf_check } + +define internal void @tester() { + %1 = alloca void (...)*, align 8 + store void (...)* bitcast (void ()* @foo to void (...)*), void (...)** %1, align 8 + %2 = load void (...)*, void (...)** %1, align 8 + call void (...) %2() + call void @foo() + call void @bar() + call void @foo_nocf_check() + call void @bar_nocf_check() + call void @tester() + ret void +} + +!llvm.module.flags = !{!1} +!1 = !{i32 4, !"ibt-fix-direct-aggr", i32 1} Index: llvm/test/CodeGen/X86/ibtfix-direct-cons.ll =================================================================== --- /dev/null +++ llvm/test/CodeGen/X86/ibtfix-direct-cons.ll @@ -0,0 +1,44 @@ +; RUN: llc < %s -O2 -mtriple=x86_64-unknown-linux-gnu -x86-indirect-branch-tracking | FileCheck %s --check-prefix=CHECK-IBT-FIX-DIRECT + +; CHECK-IBT-FIX-DIRECT: foo: +; CHECK-IBT-FIX-DIRECT: endbr +; CHECK-IBT-FIX-DIRECT: tester: +; CHECK-IBT-FIX-DIRECT-NOT: endbr +; CHECK-IBT-FIX-DIRECT: callq * +; CHECK-IBT-FIX-DIRECT: callq foo+4 +; CHECK-IBT-FIX-DIRECT-NOT: callq bar+4 +; CHECK-IBT-FIX-DIRECT-NOT: call foo_nocf_check+4 +; CHECK-IBT-FIX-DIRECT-NOT: call bar_nocf_check+4 +; CHECK-IBT-FIX-DIRECT-NOT: callq tester+4 + +target triple = "x86_64-unknown-linux-gnu" + +define dso_local void @foo() { + ret void +} + +declare dso_local void @bar() + +define dso_local void @foo_nocf_check() #0 { + ret void +} + +declare dso_local void @bar_nocf_check() #0 + +attributes #0 = { nocf_check } + +define internal void @tester() { + %1 = alloca void (...)*, align 8 + store void (...)* bitcast (void ()* @foo to void (...)*), void (...)** %1, align 8 + %2 = load void (...)*, void (...)** %1, align 8 + call void (...) %2() + call void @foo() + call void @bar() + call void @foo_nocf_check() + call void @bar_nocf_check() + call void @tester() + ret void +} + +!llvm.module.flags = !{!1} +!1 = !{i32 4, !"ibt-fix-direct-cons", i32 1}