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/MachineModuleInfo.h"
26 
27 using namespace llvm;
28 
29 #define DEBUG_TYPE "x86-indirect-branch-tracking"
30 
31 static 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;
53 
54   /// Endbr opcode for the current machine function.
55   unsigned int EndbrOpcode;
56 
57   /// Adds a new ENDBR instruction to the begining 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) 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(MachineBasicBlock &MBB) const {
73   assert(TII && "Target instruction info was not initialized");
74   assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
75          "Unexpected Endbr opcode");
76 
77   auto MI = MBB.begin();
78   // If the MBB is empty or the first instruction is not ENDBR,
79   // add the ENDBR instruction to the beginning of the MBB.
80   if (MI == MBB.end() || EndbrOpcode != MI->getOpcode()) {
81     BuildMI(MBB, MI, MBB.findDebugLoc(MI), TII->get(EndbrOpcode));
82     NumEndBranchAdded++;
83     return true;
84   }
85 
86   return false;
87 }
88 
89 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
90   const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
91 
92   // Check that the cf-protection-branch is enabled.
93   Metadata *isCFProtectionSupported =
94       MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
95   if (!isCFProtectionSupported && !IndirectBranchTracking)
96     return false;
97 
98   // True if the current MF was changed and false otherwise.
99   bool Changed = false;
100 
101   TII = SubTarget.getInstrInfo();
102   EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
103 
104   // Non-internal function or function whose address was taken, can be
105   // accessed through indirect calls. Mark the first BB with ENDBR instruction
106   // unless nocf_check attribute is used.
107   if ((MF.getFunction().hasAddressTaken() ||
108        !MF.getFunction().hasLocalLinkage()) &&
109       !MF.getFunction().doesNoCfCheck()) {
110     auto MBB = MF.begin();
111     Changed |= addENDBR(*MBB);
112   }
113 
114   for (auto &MBB : MF)
115     // Find all basic blocks that their address was taken (for example
116     // in the case of indirect jump) and add ENDBR instruction.
117     if (MBB.hasAddressTaken())
118       Changed |= addENDBR(MBB);
119 
120   return Changed;
121 }
122