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 "llvm/ADT/Statistic.h"
22 #include "llvm/CodeGen/MachineFunctionPass.h"
23 #include "llvm/CodeGen/MachineInstrBuilder.h"
24 #include "llvm/CodeGen/MachineModuleInfo.h"
25
26 using namespace llvm;
27
28 #define DEBUG_TYPE "x86-indirect-branch-tracking"
29
30 static cl::opt<bool> IndirectBranchTracking(
31 "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
32 cl::desc("Enable X86 indirect branch tracking pass."));
33
34 STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
35
36 namespace {
37 class X86IndirectBranchTrackingPass : public MachineFunctionPass {
38 public:
X86IndirectBranchTrackingPass()39 X86IndirectBranchTrackingPass() : MachineFunctionPass(ID) {}
40
getPassName() const41 StringRef getPassName() const override {
42 return "X86 Indirect Branch Tracking";
43 }
44
45 bool runOnMachineFunction(MachineFunction &MF) override;
46
47 private:
48 static char ID;
49
50 /// Machine instruction info used throughout the class.
51 const X86InstrInfo *TII = nullptr;
52
53 /// Endbr opcode for the current machine function.
54 unsigned int EndbrOpcode = 0;
55
56 /// Adds a new ENDBR instruction to the beginning of the MBB.
57 /// The function will not add it if already exists.
58 /// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
59 /// \returns true if the ENDBR was added and false otherwise.
60 bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const;
61 };
62
63 } // end anonymous namespace
64
65 char X86IndirectBranchTrackingPass::ID = 0;
66
createX86IndirectBranchTrackingPass()67 FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
68 return new X86IndirectBranchTrackingPass();
69 }
70
addENDBR(MachineBasicBlock & MBB,MachineBasicBlock::iterator I) const71 bool X86IndirectBranchTrackingPass::addENDBR(
72 MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const {
73 assert(TII && "Target instruction info was not initialized");
74 assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
75 "Unexpected Endbr opcode");
76
77 // If the MBB/I is empty or the current instruction is not ENDBR,
78 // insert ENDBR instruction to the location of I.
79 if (I == MBB.end() || I->getOpcode() != EndbrOpcode) {
80 BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode));
81 ++NumEndBranchAdded;
82 return true;
83 }
84 return false;
85 }
86
IsCallReturnTwice(llvm::MachineOperand & MOp)87 static bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
88 if (!MOp.isGlobal())
89 return false;
90 auto *CalleeFn = dyn_cast<Function>(MOp.getGlobal());
91 if (!CalleeFn)
92 return false;
93 AttributeList Attrs = CalleeFn->getAttributes();
94 if (Attrs.hasAttribute(AttributeList::FunctionIndex, Attribute::ReturnsTwice))
95 return true;
96 return false;
97 }
98
runOnMachineFunction(MachineFunction & MF)99 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
100 const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
101
102 // Check that the cf-protection-branch is enabled.
103 Metadata *isCFProtectionSupported =
104 MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
105 if (!isCFProtectionSupported && !IndirectBranchTracking)
106 return false;
107
108 // True if the current MF was changed and false otherwise.
109 bool Changed = false;
110
111 TII = SubTarget.getInstrInfo();
112 EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
113
114 // Non-internal function or function whose address was taken, can be
115 // accessed through indirect calls. Mark the first BB with ENDBR instruction
116 // unless nocf_check attribute is used.
117 if ((MF.getFunction().hasAddressTaken() ||
118 !MF.getFunction().hasLocalLinkage()) &&
119 !MF.getFunction().doesNoCfCheck()) {
120 auto MBB = MF.begin();
121 Changed |= addENDBR(*MBB, MBB->begin());
122 }
123
124 for (auto &MBB : MF) {
125 // Find all basic blocks that their address was taken (for example
126 // in the case of indirect jump) and add ENDBR instruction.
127 if (MBB.hasAddressTaken())
128 Changed |= addENDBR(MBB, MBB.begin());
129
130 for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
131 if (!I->isCall())
132 continue;
133 if (IsCallReturnTwice(I->getOperand(0)))
134 Changed |= addENDBR(MBB, std::next(I));
135 }
136 }
137 return Changed;
138 }
139