contrib/llvm-project/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 "llvm/ADT/Statistic.h"
  22 #include "llvm/CodeGen/MachineFunctionPass.h"
  23 #include "llvm/CodeGen/MachineInstrBuilder.h"
  24 #include "llvm/CodeGen/MachineModuleInfo.h"
  25
  26 using namespace llvm;
  27
  28 #define DEBUG_TYPE "x86-indirect-branch-tracking"
  29
  30 static cl::opt<bool> IndirectBranchTracking(
  31     "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
  32     cl::desc("Enable X86 indirect branch tracking pass."));
  33
  34 STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
  35
  36 namespace {
  37 class X86IndirectBranchTrackingPass : public MachineFunctionPass {
  38 public:
  39   X86IndirectBranchTrackingPass() : MachineFunctionPass(ID) {}
  40
  41   StringRef getPassName() const override {
  42     return "X86 Indirect Branch Tracking";
  43   }
  44
  45   bool runOnMachineFunction(MachineFunction &MF) override;
  46
  47 private:
  48   static char ID;
  49
  50   /// Machine instruction info used throughout the class.
  51   const X86InstrInfo *TII;
  52
  53   /// Endbr opcode for the current machine function.
  54   unsigned int EndbrOpcode;
  55
  56   /// Adds a new ENDBR instruction to the begining of the MBB.
  57   /// The function will not add it if already exists.
  58   /// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
  59   /// \returns true if the ENDBR was added and false otherwise.
  60   bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const;
  61 };
  62
  63 } // end anonymous namespace
  64
  65 char X86IndirectBranchTrackingPass::ID = 0;
  66
  67 FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
  68   return new X86IndirectBranchTrackingPass();
  69 }
  70
  71 bool X86IndirectBranchTrackingPass::addENDBR(
  72     MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const {
  73   assert(TII && "Target instruction info was not initialized");
  74   assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
  75          "Unexpected Endbr opcode");
  76
  77   // If the MBB/I is empty or the current instruction is not ENDBR,
  78   // insert ENDBR instruction to the location of I.
  79   if (I == MBB.end() || I->getOpcode() != EndbrOpcode) {
  80     BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode));
  81     ++NumEndBranchAdded;
  82     return true;
  83   }
  84   return false;
  85 }
  86
  87 bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
  88   if (!MOp.isGlobal())
  89     return false;
  90   auto *CalleeFn = dyn_cast<Function>(MOp.getGlobal());
  91   if (!CalleeFn)
  92     return false;
  93   AttributeList Attrs = CalleeFn->getAttributes();
  94   if (Attrs.hasAttribute(AttributeList::FunctionIndex, Attribute::ReturnsTwice))
  95     return true;
  96   return false;
  97 }
  98
  99 bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
 100   const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
 101
 102   // Check that the cf-protection-branch is enabled.
 103   Metadata *isCFProtectionSupported =
 104       MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
 105   if (!isCFProtectionSupported && !IndirectBranchTracking)
 106     return false;
 107
 108   // True if the current MF was changed and false otherwise.
 109   bool Changed = false;
 110
 111   TII = SubTarget.getInstrInfo();
 112   EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
 113
 114   // Non-internal function or function whose address was taken, can be
 115   // accessed through indirect calls. Mark the first BB with ENDBR instruction
 116   // unless nocf_check attribute is used.
 117   if ((MF.getFunction().hasAddressTaken() ||
 118        !MF.getFunction().hasLocalLinkage()) &&
 119       !MF.getFunction().doesNoCfCheck()) {
 120     auto MBB = MF.begin();
 121     Changed |= addENDBR(*MBB, MBB->begin());
 122   }
 123
 124   for (auto &MBB : MF) {
 125     // Find all basic blocks that their address was taken (for example
 126     // in the case of indirect jump) and add ENDBR instruction.
 127     if (MBB.hasAddressTaken())
 128       Changed |= addENDBR(MBB, MBB.begin());
 129
 130     for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
 131       if (!I->isCall())
 132         continue;
 133       if (IsCallReturnTwice(I->getOperand(0)))
 134         Changed |= addENDBR(MBB, std::next(I));
 135     }
 136   }
 137   return Changed;
 138 }