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(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(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>(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   const X86TargetMachine *TM =
106       static_cast<const X86TargetMachine *>(&MF.getTarget());
107   Metadata *IBTSeal = M->getModuleFlag("ibt-seal");
108 
109   switch (TM->getCodeModel()) {
110   // Large code model functions always reachable through indirect calls.
111   case CodeModel::Large:
112     return true;
113   // Only address taken functions in LTO'ed kernel are reachable indirectly.
114   // IBTSeal implies LTO, thus only check if function is address taken.
115   case CodeModel::Kernel:
116     // Check if ibt-seal was enabled (implies LTO is being used).
117     if (IBTSeal) {
118       return F.hasAddressTaken();
119     }
120     // if !IBTSeal, fall into default case.
121     LLVM_FALLTHROUGH;
122   // Address taken or externally linked functions may be reachable.
123   default:
124     return (F.hasAddressTaken() || !F.hasLocalLinkage());
125   }
126 }
127 
128 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
129   const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
130 
131   const Module *M = MF.getMMI().getModule();
132   // Check that the cf-protection-branch is enabled.
133   Metadata *isCFProtectionSupported = M->getModuleFlag("cf-protection-branch");
134 
135   //  NB: We need to enable IBT in jitted code if JIT compiler is CET
136   //  enabled.
137   const X86TargetMachine *TM =
138       static_cast<const X86TargetMachine *>(&MF.getTarget());
139 #ifdef __CET__
140   bool isJITwithCET = TM->isJIT();
141 #else
142   bool isJITwithCET = false;
143 #endif
144   if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET)
145     return false;
146 
147   // True if the current MF was changed and false otherwise.
148   bool Changed = false;
149 
150   TII = SubTarget.getInstrInfo();
151   EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
152 
153   // If function is reachable indirectly, mark the first BB with ENDBR.
154   if (needsPrologueENDBR(MF, M)) {
155     auto MBB = MF.begin();
156     Changed |= addENDBR(*MBB, MBB->begin());
157   }
158 
159   for (auto &MBB : MF) {
160     // Find all basic blocks that their address was taken (for example
161     // in the case of indirect jump) and add ENDBR instruction.
162     if (MBB.hasAddressTaken())
163       Changed |= addENDBR(MBB, MBB.begin());
164 
165     for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
166       if (I->isCall() && I->getNumOperands() > 0 &&
167           IsCallReturnTwice(I->getOperand(0))) {
168         Changed |= addENDBR(MBB, std::next(I));
169       }
170     }
171 
172     // Exception handle may indirectly jump to catch pad, So we should add
173     // ENDBR before catch pad instructions. For SjLj exception model, it will
174     // create a new BB(new landingpad) indirectly jump to the old landingpad.
175     if (TM->Options.ExceptionModel == ExceptionHandling::SjLj) {
176       for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
177         // New Landingpad BB without EHLabel.
178         if (MBB.isEHPad()) {
179           if (I->isDebugInstr())
180             continue;
181           Changed |= addENDBR(MBB, I);
182           break;
183         } else if (I->isEHLabel()) {
184           // Old Landingpad BB (is not Landingpad now) with
185           // the the old "callee" EHLabel.
186           MCSymbol *Sym = I->getOperand(0).getMCSymbol();
187           if (!MF.hasCallSiteLandingPad(Sym))
188             continue;
189           Changed |= addENDBR(MBB, std::next(I));
190           break;
191         }
192       }
193     } else if (MBB.isEHPad()){
194       for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
195         if (!I->isEHLabel())
196           continue;
197         Changed |= addENDBR(MBB, std::next(I));
198         break;
199       }
200     }
201   }
202   return Changed;
203 }
204