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 #include "llvm/IR/Module.h"
  27
  28 using namespace llvm;
  29
  30 #define DEBUG_TYPE "x86-indirect-branch-tracking"
  31
  32 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 = nullptr;
  54
  55   /// Endbr opcode for the current machine function.
  56   unsigned int EndbrOpcode = 0;
  57
  58   /// Adds a new ENDBR instruction to the beginning 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   /// \returns true if the ENDBR was added and false otherwise.
  62   bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const;
  63 };
  64
  65 } // end anonymous namespace
  66
  67 char X86IndirectBranchTrackingPass::ID = 0;
  68
  69 FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
  70   return new X86IndirectBranchTrackingPass();
  71 }
  72
  73 bool X86IndirectBranchTrackingPass::addENDBR(
  74     MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const {
  75   assert(TII && "Target instruction info was not initialized");
  76   assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
  77          "Unexpected Endbr opcode");
  78
  79   // If the MBB/I is empty or the current instruction is not ENDBR,
  80   // insert ENDBR instruction to the location of I.
  81   if (I == MBB.end() || I->getOpcode() != EndbrOpcode) {
  82     BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode));
  83     ++NumEndBranchAdded;
  84     return true;
  85   }
  86   return false;
  87 }
  88
  89 static bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
  90   if (!MOp.isGlobal())
  91     return false;
  92   auto *CalleeFn = dyn_cast<Function>(MOp.getGlobal());
  93   if (!CalleeFn)
  94     return false;
  95   AttributeList Attrs = CalleeFn->getAttributes();
  96   return Attrs.hasFnAttr(Attribute::ReturnsTwice);
  97 }
  98
  99 // Checks if function should have an ENDBR in its prologue
 100 static bool needsPrologueENDBR(MachineFunction &MF, const Module *M) {
 101   Function &F = MF.getFunction();
 102
 103   if (F.doesNoCfCheck())
 104     return false;
 105
 106   switch (MF.getTarget().getCodeModel()) {
 107   // Large code model functions always reachable through indirect calls.
 108   case CodeModel::Large:
 109     return true;
 110   // Address taken or externally linked functions may be reachable.
 111   default:
 112     return (F.hasAddressTaken() || !F.hasLocalLinkage());
 113   }
 114 }
 115
 116 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
 117   const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
 118
 119   const Module *M = MF.getFunction().getParent();
 120   // Check that the cf-protection-branch is enabled.
 121   Metadata *isCFProtectionSupported = M->getModuleFlag("cf-protection-branch");
 122
 123   //  NB: We need to enable IBT in jitted code if JIT compiler is CET
 124   //  enabled.
 125   const X86TargetMachine *TM =
 126       static_cast<const X86TargetMachine *>(&MF.getTarget());
 127 #ifdef __CET__
 128   bool isJITwithCET = TM->isJIT();
 129 #else
 130   bool isJITwithCET = false;
 131 #endif
 132   if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET)
 133     return false;
 134
 135   // True if the current MF was changed and false otherwise.
 136   bool Changed = false;
 137
 138   TII = SubTarget.getInstrInfo();
 139   EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
 140
 141   // If function is reachable indirectly, mark the first BB with ENDBR.
 142   if (needsPrologueENDBR(MF, M)) {
 143     auto MBB = MF.begin();
 144     Changed |= addENDBR(*MBB, MBB->begin());
 145   }
 146
 147   for (auto &MBB : MF) {
 148     // Find all basic blocks that their address was taken (for example
 149     // in the case of indirect jump) and add ENDBR instruction.
 150     if (MBB.hasAddressTaken())
 151       Changed |= addENDBR(MBB, MBB.begin());
 152
 153     for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
 154       if (I->isCall() && I->getNumOperands() > 0 &&
 155           IsCallReturnTwice(I->getOperand(0))) {
 156         Changed |= addENDBR(MBB, std::next(I));
 157       }
 158     }
 159
 160     // Exception handle may indirectly jump to catch pad, So we should add
 161     // ENDBR before catch pad instructions. For SjLj exception model, it will
 162     // create a new BB(new landingpad) indirectly jump to the old landingpad.
 163     if (TM->Options.ExceptionModel == ExceptionHandling::SjLj) {
 164       for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
 165         // New Landingpad BB without EHLabel.
 166         if (MBB.isEHPad()) {
 167           if (I->isDebugInstr())
 168             continue;
 169           Changed |= addENDBR(MBB, I);
 170           break;
 171         } else if (I->isEHLabel()) {
 172           // Old Landingpad BB (is not Landingpad now) with
 173           // the old "callee" EHLabel.
 174           MCSymbol *Sym = I->getOperand(0).getMCSymbol();
 175           if (!MF.hasCallSiteLandingPad(Sym))
 176             continue;
 177           Changed |= addENDBR(MBB, std::next(I));
 178           break;
 179         }
 180       }
 181     } else if (MBB.isEHPad()){
 182       for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
 183         if (!I->isEHLabel())
 184           continue;
 185         Changed |= addENDBR(MBB, std::next(I));
 186         break;
 187       }
 188     }
 189   }
 190   return Changed;
 191 }