1//===---- X86IndirectBranchTracking.cpp - Enables CET IBT mechanism -------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file defines a pass that enables Indirect Branch Tracking (IBT) as part
10// of Control-Flow Enforcement Technology (CET).
11// The pass adds ENDBR (End Branch) machine instructions at the beginning of
12// each basic block or function that is referenced by an indrect jump/call
13// instruction.
14// The ENDBR instructions have a NOP encoding and as such are ignored in
15// targets that do not support CET IBT mechanism.
16//===----------------------------------------------------------------------===//
17
18#include "X86.h"
19#include "X86InstrInfo.h"
20#include "X86Subtarget.h"
21#include "X86TargetMachine.h"
22#include "llvm/ADT/Statistic.h"
23#include "llvm/CodeGen/MachineFunctionPass.h"
24#include "llvm/CodeGen/MachineInstrBuilder.h"
25#include "llvm/CodeGen/MachineModuleInfo.h"
26#include "llvm/IR/Module.h"
27
28using namespace llvm;
29
30#define DEBUG_TYPE "x86-indirect-branch-tracking"
31
32cl::opt<bool> IndirectBranchTracking(
33 "x86-indirect-branch-tracking", cl::init(Val: false), cl::Hidden,
34 cl::desc("Enable X86 indirect branch tracking pass."));
35
36STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
37
38namespace {
39class X86IndirectBranchTrackingLegacy : public MachineFunctionPass {
40public:
41 static char ID;
42
43 X86IndirectBranchTrackingLegacy() : MachineFunctionPass(ID) {}
44
45 StringRef getPassName() const override {
46 return "X86 Indirect Branch Tracking";
47 }
48
49 bool runOnMachineFunction(MachineFunction &MF) override;
50};
51
52/// Adds a new ENDBR instruction to the beginning of the MBB.
53/// The function will not add it if already exists.
54/// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
55/// \returns true if the ENDBR was added and false otherwise.
56static bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) {
57 MachineFunction &MF = *MBB.getParent();
58 const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
59 const X86InstrInfo *TII = SubTarget.getInstrInfo();
60 unsigned EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
61
62 assert(TII && "Target instruction info was not initialized");
63 assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
64 "Unexpected Endbr opcode");
65
66 // If the MBB/I is empty or the current instruction is not ENDBR,
67 // insert ENDBR instruction to the location of I.
68 if (I == MBB.end() || I->getOpcode() != EndbrOpcode) {
69 BuildMI(BB&: MBB, I, MIMD: MBB.findDebugLoc(MBBI: I), MCID: TII->get(Opcode: EndbrOpcode));
70 ++NumEndBranchAdded;
71 return true;
72 }
73 return false;
74}
75
76} // end anonymous namespace
77
78char X86IndirectBranchTrackingLegacy::ID = 0;
79
80INITIALIZE_PASS(X86IndirectBranchTrackingLegacy, DEBUG_TYPE,
81 "X86 Indirect Branch Tracking", false, false)
82
83FunctionPass *llvm::createX86IndirectBranchTrackingLegacyPass() {
84 return new X86IndirectBranchTrackingLegacy();
85}
86
87static bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
88 if (!MOp.isGlobal())
89 return false;
90 auto *CalleeFn = dyn_cast<Function>(Val: MOp.getGlobal());
91 if (!CalleeFn)
92 return false;
93 AttributeList Attrs = CalleeFn->getAttributes();
94 return Attrs.hasFnAttr(Kind: Attribute::ReturnsTwice);
95}
96
97// Checks if function should have an ENDBR in its prologue
98static bool needsPrologueENDBR(MachineFunction &MF, const Module *M) {
99 Function &F = MF.getFunction();
100
101 if (F.doesNoCfCheck())
102 return false;
103
104 switch (MF.getTarget().getCodeModel()) {
105 // Large code model functions always reachable through indirect calls.
106 case CodeModel::Large:
107 return true;
108 // Address taken or externally linked functions may be reachable.
109 default:
110 return (F.hasAddressTaken() || !F.hasLocalLinkage());
111 }
112}
113
114static bool runIndirectBranchTracking(MachineFunction &MF) {
115 const Module *M = MF.getFunction().getParent();
116 // Check that the cf-protection-branch is enabled.
117 Metadata *isCFProtectionSupported = M->getModuleFlag(Key: "cf-protection-branch");
118
119 // NB: We need to enable IBT in jitted code if JIT compiler is CET
120 // enabled.
121 const X86TargetMachine *TM =
122 static_cast<const X86TargetMachine *>(&MF.getTarget());
123#ifdef __CET__
124 bool isJITwithCET = TM->isJIT();
125#else
126 bool isJITwithCET = false;
127#endif
128 if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET)
129 return false;
130
131 // True if the current MF was changed and false otherwise.
132 bool Changed = false;
133
134 // If function is reachable indirectly, mark the first BB with ENDBR.
135 if (needsPrologueENDBR(MF, M)) {
136 auto MBB = MF.begin();
137 Changed |= addENDBR(MBB&: *MBB, I: MBB->begin());
138 }
139
140 for (auto &MBB : MF) {
141 // Find all basic blocks that their address was taken (for example
142 // in the case of indirect jump) and add ENDBR instruction.
143 if (MBB.isMachineBlockAddressTaken() || MBB.isIRBlockAddressTaken())
144 Changed |= addENDBR(MBB, I: MBB.begin());
145
146 for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
147 if (I->isCall() && I->getNumOperands() > 0 &&
148 IsCallReturnTwice(MOp&: I->getOperand(i: 0))) {
149 Changed |= addENDBR(MBB, I: std::next(x: I));
150 }
151 }
152
153 // Exception handle may indirectly jump to catch pad, So we should add
154 // ENDBR before catch pad instructions. For SjLj exception model, it will
155 // create a new BB(new landingpad) indirectly jump to the old landingpad.
156 if (TM->Options.ExceptionModel == ExceptionHandling::SjLj) {
157 for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
158 // New Landingpad BB without EHLabel.
159 if (MBB.isEHPad()) {
160 if (I->isDebugInstr())
161 continue;
162 Changed |= addENDBR(MBB, I);
163 break;
164 } else if (I->isEHLabel()) {
165 // Old Landingpad BB (is not Landingpad now) with
166 // the old "callee" EHLabel.
167 MCSymbol *Sym = I->getOperand(i: 0).getMCSymbol();
168 if (!MF.hasCallSiteLandingPad(Sym))
169 continue;
170 Changed |= addENDBR(MBB, I: std::next(x: I));
171 break;
172 }
173 }
174 } else if (MBB.isEHPad()){
175 for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
176 if (!I->isEHLabel())
177 continue;
178 Changed |= addENDBR(MBB, I: std::next(x: I));
179 break;
180 }
181 }
182 }
183 return Changed;
184}
185
186bool X86IndirectBranchTrackingLegacy::runOnMachineFunction(
187 MachineFunction &MF) {
188 return runIndirectBranchTracking(MF);
189}
190
191PreservedAnalyses
192X86IndirectBranchTrackingPass::run(MachineFunction &MF,
193 MachineFunctionAnalysisManager &MFAM) {
194 return runIndirectBranchTracking(MF)
195 ? getMachineFunctionPassPreservedAnalyses()
196 .preserveSet<CFGAnalyses>()
197 : PreservedAnalyses::all();
198}
199