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