brintos

brintos / llvm-project-archived public Read only

0
0
Text · 6.5 KiB · 52be142 Raw
192 lines · 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-exception6//7//===----------------------------------------------------------------------===//8//9// This file defines a pass that enables Indirect Branch Tracking (IBT) as part10// of Control-Flow Enforcement Technology (CET).11// The pass adds ENDBR (End Branch) machine instructions at the beginning of12// each basic block or function that is referenced by an indrect jump/call13// instruction.14// The ENDBR instructions have a NOP encoding and as such are ignored in15// 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 28using namespace llvm;29 30#define DEBUG_TYPE "x86-indirect-branch-tracking"31 32cl::opt<bool> IndirectBranchTracking(33    "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,34    cl::desc("Enable X86 indirect branch tracking pass."));35 36STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");37 38namespace {39class X86IndirectBranchTrackingPass : public MachineFunctionPass {40public: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 49private: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 namespace66 67char X86IndirectBranchTrackingPass::ID = 0;68 69FunctionPass *llvm::createX86IndirectBranchTrackingPass() {70  return new X86IndirectBranchTrackingPass();71}72 73bool 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 89static 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 prologue100static 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 116bool 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 CET124  //  enabled.125  const X86TargetMachine *TM =126      static_cast<const X86TargetMachine *>(&MF.getTarget());127#ifdef __CET__128  bool isJITwithCET = TM->isJIT();129#else130  bool isJITwithCET = false;131#endif132  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 example149    // in the case of indirect jump) and add ENDBR instruction.150    if (MBB.isMachineBlockAddressTaken() || MBB.isIRBlockAddressTaken())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 add161    // ENDBR before catch pad instructions. For SjLj exception model, it will162    // 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) with173          // 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}192