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/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   // Make sure that the target supports IBT instruction.
  93   if (!SubTarget.hasIBT())
  94     return false;
  95
  96   // Check that the cf-protection-branch is enabled.
  97   Metadata *isCFProtectionSupported =
  98       MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
  99   if (!isCFProtectionSupported && !IndirectBranchTracking)
 100     return false;
 101
 102   // True if the current MF was changed and false otherwise.
 103   bool Changed = false;
 104
 105   TII = SubTarget.getInstrInfo();
 106   EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
 107
 108   // Non-internal function or function whose address was taken, can be
 109   // accessed through indirect calls. Mark the first BB with ENDBR instruction
 110   // unless nocf_check attribute is used.
 111   if ((MF.getFunction().hasAddressTaken() ||
 112        !MF.getFunction().hasLocalLinkage()) &&
 113       !MF.getFunction().doesNoCfCheck()) {
 114     auto MBB = MF.begin();
 115     Changed |= addENDBR(*MBB);
 116   }
 117
 118   for (auto &MBB : MF)
 119     // Find all basic blocks that their address was taken (for example
 120     // in the case of indirect jump) and add ENDBR instruction.
 121     if (MBB.hasAddressTaken())
 122       Changed |= addENDBR(MBB);
 123
 124   return Changed;
 125 }