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 | |
27 | using namespace llvm; |
28 | |
29 | #define DEBUG_TYPE "x86-indirect-branch-tracking" |
30 | |
31 | cl::opt<bool> IndirectBranchTracking( |
32 | "x86-indirect-branch-tracking" , cl::init(Val: false), cl::Hidden, |
33 | cl::desc("Enable X86 indirect branch tracking pass." )); |
34 | |
35 | STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added" ); |
36 | |
37 | namespace { |
38 | class X86IndirectBranchTrackingPass : public MachineFunctionPass { |
39 | public: |
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 | |
48 | private: |
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 | |
66 | char X86IndirectBranchTrackingPass::ID = 0; |
67 | |
68 | FunctionPass *llvm::createX86IndirectBranchTrackingPass() { |
69 | return new X86IndirectBranchTrackingPass(); |
70 | } |
71 | |
72 | bool 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 | |
88 | static 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 |
99 | static 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 | |
115 | bool 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 | |