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
27using namespace llvm;
28
29#define DEBUG_TYPE "x86-indirect-branch-tracking"
30
31cl::opt<bool> IndirectBranchTracking(
32 "x86-indirect-branch-tracking", cl::init(Val: false), cl::Hidden,
33 cl::desc("Enable X86 indirect branch tracking pass."));
34
35STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
36
37namespace {
38class X86IndirectBranchTrackingPass : public MachineFunctionPass {
39public:
40 X86IndirectBranchTrackingPass() : MachineFunctionPass(ID) {}
41
42 StringRef getPassName() const override {
43 return "X86 Indirect Branch Tracking";
44 }
45
46 bool runOnMachineFunction(MachineFunction &MF) override;
47
48private:
49 static char ID;
50
51 /// Machine instruction info used throughout the class.
52 const X86InstrInfo *TII = nullptr;
53
54 /// Endbr opcode for the current machine function.
55 unsigned int EndbrOpcode = 0;
56
57 /// Adds a new ENDBR instruction to the beginning of the MBB.
58 /// The function will not add it if already exists.
59 /// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
60 /// \returns true if the ENDBR was added and false otherwise.
61 bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const;
62};
63
64} // end anonymous namespace
65
66char X86IndirectBranchTrackingPass::ID = 0;
67
68FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
69 return new X86IndirectBranchTrackingPass();
70}
71
72bool X86IndirectBranchTrackingPass::addENDBR(
73 MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const {
74 assert(TII && "Target instruction info was not initialized");
75 assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
76 "Unexpected Endbr opcode");
77
78 // If the MBB/I is empty or the current instruction is not ENDBR,
79 // insert ENDBR instruction to the location of I.
80 if (I == MBB.end() || I->getOpcode() != EndbrOpcode) {
81 BuildMI(MBB, I, MBB.findDebugLoc(MBBI: I), TII->get(EndbrOpcode));
82 ++NumEndBranchAdded;
83 return true;
84 }
85 return false;
86}
87
88static bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
89 if (!MOp.isGlobal())
90 return false;
91 auto *CalleeFn = dyn_cast<Function>(Val: MOp.getGlobal());
92 if (!CalleeFn)
93 return false;
94 AttributeList Attrs = CalleeFn->getAttributes();
95 return Attrs.hasFnAttr(Attribute::ReturnsTwice);
96}
97
98// Checks if function should have an ENDBR in its prologue
99static bool needsPrologueENDBR(MachineFunction &MF, const Module *M) {
100 Function &F = MF.getFunction();
101
102 if (F.doesNoCfCheck())
103 return false;
104
105 switch (MF.getTarget().getCodeModel()) {
106 // Large code model functions always reachable through indirect calls.
107 case CodeModel::Large:
108 return true;
109 // Address taken or externally linked functions may be reachable.
110 default:
111 return (F.hasAddressTaken() || !F.hasLocalLinkage());
112 }
113}
114
115bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
116 const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
117
118 const Module *M = MF.getMMI().getModule();
119 // Check that the cf-protection-branch is enabled.
120 Metadata *isCFProtectionSupported = M->getModuleFlag(Key: "cf-protection-branch");
121
122 // NB: We need to enable IBT in jitted code if JIT compiler is CET
123 // enabled.
124 const X86TargetMachine *TM =
125 static_cast<const X86TargetMachine *>(&MF.getTarget());
126#ifdef __CET__
127 bool isJITwithCET = TM->isJIT();
128#else
129 bool isJITwithCET = false;
130#endif
131 if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET)
132 return false;
133
134 // True if the current MF was changed and false otherwise.
135 bool Changed = false;
136
137 TII = SubTarget.getInstrInfo();
138 EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
139
140 // If function is reachable indirectly, mark the first BB with ENDBR.
141 if (needsPrologueENDBR(MF, M)) {
142 auto MBB = MF.begin();
143 Changed |= addENDBR(MBB&: *MBB, I: MBB->begin());
144 }
145
146 for (auto &MBB : MF) {
147 // Find all basic blocks that their address was taken (for example
148 // in the case of indirect jump) and add ENDBR instruction.
149 if (MBB.hasAddressTaken())
150 Changed |= addENDBR(MBB, I: MBB.begin());
151
152 for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
153 if (I->isCall() && I->getNumOperands() > 0 &&
154 IsCallReturnTwice(MOp&: I->getOperand(i: 0))) {
155 Changed |= addENDBR(MBB, I: std::next(x: I));
156 }
157 }
158
159 // Exception handle may indirectly jump to catch pad, So we should add
160 // ENDBR before catch pad instructions. For SjLj exception model, it will
161 // create a new BB(new landingpad) indirectly jump to the old landingpad.
162 if (TM->Options.ExceptionModel == ExceptionHandling::SjLj) {
163 for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
164 // New Landingpad BB without EHLabel.
165 if (MBB.isEHPad()) {
166 if (I->isDebugInstr())
167 continue;
168 Changed |= addENDBR(MBB, I);
169 break;
170 } else if (I->isEHLabel()) {
171 // Old Landingpad BB (is not Landingpad now) with
172 // the old "callee" EHLabel.
173 MCSymbol *Sym = I->getOperand(i: 0).getMCSymbol();
174 if (!MF.hasCallSiteLandingPad(Sym))
175 continue;
176 Changed |= addENDBR(MBB, I: std::next(x: I));
177 break;
178 }
179 }
180 } else if (MBB.isEHPad()){
181 for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
182 if (!I->isEHLabel())
183 continue;
184 Changed |= addENDBR(MBB, I: std::next(x: I));
185 break;
186 }
187 }
188 }
189 return Changed;
190}
191

source code of llvm/lib/Target/X86/X86IndirectBranchTracking.cpp