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