lib/Target/X86/X86IndirectBranchTracking.cpp

   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   /// Adds a new ENDBR instruction to the begining of the MBB.
  59   /// The function will not add it if already exists.
  60   /// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
  61   void 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 void 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   }
  84 }
  85
  86 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
  87   const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
  88
  89   // Make sure that the target supports ENDBR instruction.
  90   if (!SubTarget.hasIBT())
  91     return false;
  92
  93   // Check that the cf-protection-branch is enabled.
  94   Metadata *isCFProtectionSupported =
  95       MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
  96   if (!isCFProtectionSupported && !IndirectBranchTracking)
  97     return false;
  98
  99   // True if the current MF was changed and false otherwise.
 100   bool Changed = false;
 101
 102   TII = SubTarget.getInstrInfo();
 103   EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
 104
 105   // Non-internal function or function whose address was taken, can be
 106   // invoked through indirect calls. Mark the first BB with ENDBR instruction.
 107   // TODO: Do not add ENDBR instruction in case notrack attribute is used.
 108   if (MF.getFunction().hasAddressTaken() ||
 109       !MF.getFunction().hasLocalLinkage()) {
 110     auto MBB = MF.begin();
 111     addENDBR(*MBB);
 112     Changed = true;
 113   }
 114
 115   for (auto &MBB : MF) {
 116     // Find all basic blocks that thier address was taken (for example
 117     // in the case of indirect jump) and add ENDBR instruction.
 118     if (MBB.hasAddressTaken()) {
 119       addENDBR(MBB);
 120       Changed = true;
 121     }
 122   }
 123
 124   // Adds ENDBR instructions to MBB destinations of the jump table.
 125   // TODO: In case of more than 50 destinations, do not add ENDBR and
 126   // instead add DS_PREFIX.
 127   if (MachineJumpTableInfo *JTI = MF.getJumpTableInfo()) {
 128     for (const auto &JT : JTI->getJumpTables()) {
 129       for (auto *MBB : JT.MBBs) {
 130         // This assert verifies the assumption that this MBB has an indirect
 131         // jump terminator in one of its predecessor.
 132         // Jump tables are generated when lowering switch-case statements or
 133         // setjmp/longjump functions. As a result only indirect jumps use jump
 134         // tables.
 135         #ifndef NDEBUG
 136         bool hasIndirectJumpTerm = false;
 137         for (auto &PredMBB : MBB->predecessors())
 138           for (auto &TermI : PredMBB->terminators())
 139             if (TermI.isIndirectBranch())
 140               hasIndirectJumpTerm = true;
 141         assert(hasIndirectJumpTerm &&
 142                "The MBB is not the destination of an indirect jump");
 143         (void)hasIndirectJumpTerm;
 144         #endif
 145         addENDBR(*MBB);
 146         Changed = true;
 147       }
 148     }
 149   }
 150
 151   return Changed;
 152 }