llvm/lib/Target/X86/X86IndirectBranchTracking.cpp

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