1 //===---- X86IndirectBranchTracking.cpp - Enables CET IBT mechanism -------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file defines a pass that enables Indirect Branch Tracking (IBT) as part
11 // of Control-Flow Enforcement Technology (CET).
12 // The pass adds ENDBR (End Branch) machine instructions at the beginning of
13 // each basic block or function that is referenced by an indrect jump/call
14 // instruction.
15 // The ENDBR instructions have a NOP encoding and as such are ignored in
16 // targets that do not support CET IBT mechanism.
17 //===----------------------------------------------------------------------===//
18 
19 #include "X86.h"
20 #include "X86InstrInfo.h"
21 #include "X86Subtarget.h"
22 #include "llvm/ADT/Statistic.h"
23 #include "llvm/CodeGen/MachineFunctionPass.h"
24 #include "llvm/CodeGen/MachineInstrBuilder.h"
25 #include "llvm/CodeGen/MachineJumpTableInfo.h"
26 #include "llvm/CodeGen/MachineModuleInfo.h"
27 
28 using namespace llvm;
29 
30 #define DEBUG_TYPE "x86-indirect-branch-tracking"
31 
32 static cl::opt<bool> IndirectBranchTracking(
33     "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
34     cl::desc("Enable X86 indirect branch tracking pass."));
35 
36 STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
37 
38 namespace {
39 class X86IndirectBranchTrackingPass : public MachineFunctionPass {
40 public:
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 
49 private:
50   static char ID;
51 
52   /// Machine instruction info used throughout the class.
53   const X86InstrInfo *TII;
54 
55   /// Endbr opcode for the current machine function.
56   unsigned int EndbrOpcode;
57 
58   /// The function looks for an indirect jump terminator in MBB predecessors.
59   ///
60   /// Jump tables are generated when lowering switch-case statements or
61   /// setjmp/longjump functions.
62   /// As a result only indirect jumps use jump tables.
63   /// The function verifies this assumption.
64   ///
65   /// \return true if the input \p MBB has a predecessor MBB with indirect
66   /// branch terminator or false otherwise.
67   bool verifyIndirectJump(const MachineBasicBlock *MBB) const;
68 
69   /// Adds a new ENDBR instruction to the begining of the MBB.
70   /// The function will not add it if already exists.
71   /// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
72   void addENDBR(MachineBasicBlock &MBB) const;
73 };
74 
75 } // end anonymous namespace
76 
77 char X86IndirectBranchTrackingPass::ID = 0;
78 
79 FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
80   return new X86IndirectBranchTrackingPass();
81 }
82 
83 bool X86IndirectBranchTrackingPass::verifyIndirectJump(
84     const MachineBasicBlock *MBB) const {
85   for (auto &PredMBB : MBB->predecessors())
86     for (auto &TermI : PredMBB->terminators())
87       if (TermI.isIndirectBranch())
88         return true;
89 
90   return false;
91 }
92 
93 void X86IndirectBranchTrackingPass::addENDBR(MachineBasicBlock &MBB) const {
94   assert(TII && "Target instruction info was not initialized");
95   assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
96          "Unexpected Endbr opcode");
97 
98   auto MI = MBB.begin();
99   // If the MBB is empty or the first instruction is not ENDBR,
100   // add the ENDBR instruction to the beginning of the MBB.
101   if (MI == MBB.end() || EndbrOpcode != MI->getOpcode()) {
102     BuildMI(MBB, MI, MBB.findDebugLoc(MI), TII->get(EndbrOpcode));
103     NumEndBranchAdded++;
104   }
105 }
106 
107 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
108   const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
109 
110   // Make sure that the target supports ENDBR instruction.
111   if (!SubTarget.hasIBT())
112     return false;
113 
114   // Check that the cf-protection-branch is enabled.
115   Metadata *isCFProtectionSupported =
116       MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
117   if (!isCFProtectionSupported && !IndirectBranchTracking)
118     return false;
119 
120   // True if the current MF was changed and false otherwise.
121   bool Changed = false;
122 
123   TII = SubTarget.getInstrInfo();
124   EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
125 
126   // Non-internal function or function whose address was taken, can be
127   // invoked through indirect calls. Mark the first BB with ENDBR instruction.
128   // TODO: Do not add ENDBR instruction in case notrack attribute is used.
129   if (MF.getFunction().hasAddressTaken() ||
130       !MF.getFunction().hasLocalLinkage()) {
131     auto MBB = MF.begin();
132     addENDBR(*MBB);
133     Changed = true;
134   }
135 
136   for (auto &MBB : MF) {
137     // Find all basic blocks that thier address was taken (for example
138     // in the case of indirect jump) and add ENDBR instruction.
139     if (MBB.hasAddressTaken()) {
140       addENDBR(MBB);
141       Changed = true;
142     }
143   }
144 
145   // Adds ENDBR instructions to MBB destinations of the jump table.
146   // TODO: In case of more than 50 destinations, do not add ENDBR and
147   // instead add DS_PREFIX.
148   if (MachineJumpTableInfo *JTI = MF.getJumpTableInfo()) {
149     for (const auto &JT : JTI->getJumpTables()) {
150       for (auto *MBB : JT.MBBs) {
151         // This assert verifies the assumption that this MBB has an indirect
152         // jump terminator in one of its predecessor.
153         assert(verifyIndirectJump(MBB) &&
154                "The MBB is not the destination of an indirect jump");
155 
156         addENDBR(*MBB);
157         Changed = true;
158       }
159     }
160   }
161 
162   return Changed;
163 }
164