text: //===---- X86IndirectBranchTracking.cpp - Enables CET IBT mechanism -------===//
  //
  // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  // See https://llvm.org/LICENSE.txt for license information.
  // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  //
  //===----------------------------------------------------------------------===//
  //
  // This file defines a pass that enables Indirect Branch Tracking (IBT) as part
  // of Control-Flow Enforcement Technology (CET).
  // The pass adds ENDBR (End Branch) machine instructions at the beginning of
  // each basic block or function that is referenced by an indrect jump/call
  // instruction.
  // The ENDBR instructions have a NOP encoding and as such are ignored in
  // targets that do not support CET IBT mechanism.
  //===----------------------------------------------------------------------===//
  
  #include "X86.h"
  #include "X86InstrInfo.h"
  #include "X86Subtarget.h"
  #include "llvm/ADT/Statistic.h"
  #include "llvm/CodeGen/MachineFunctionPass.h"
  #include "llvm/CodeGen/MachineInstrBuilder.h"
  #include "llvm/CodeGen/MachineModuleInfo.h"
  
  using namespace llvm;
  
  #define DEBUG_TYPE "x86-indirect-branch-tracking"
  
  static cl::opt<bool> IndirectBranchTracking(
      "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
      cl::desc("Enable X86 indirect branch tracking pass."));
  
  STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
  
  namespace {
  class X86IndirectBranchTrackingPass : public MachineFunctionPass {
  public:
    X86IndirectBranchTrackingPass() : MachineFunctionPass(ID) {}
  
    StringRef getPassName() const override {
      return "X86 Indirect Branch Tracking";
    }
  
    bool runOnMachineFunction(MachineFunction &MF) override;
  
  private:
    static char ID;
  
    /// Machine instruction info used throughout the class.
    const X86InstrInfo *TII;
  
    /// Endbr opcode for the current machine function.
    unsigned int EndbrOpcode;
  
    /// Adds a new ENDBR instruction to the beginning of the MBB.
    /// The function will not add it if already exists.
    /// It will add ENDBR32 or ENDBR64 opcode, depending on the target.
    /// \returns true if the ENDBR was added and false otherwise.
    bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const;
  };
  
  } // end anonymous namespace
  
  char X86IndirectBranchTrackingPass::ID = 0;
  
  FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
    return new X86IndirectBranchTrackingPass();
  }
  
  bool X86IndirectBranchTrackingPass::addENDBR(
      MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const {
    assert(TII && "Target instruction info was not initialized");
    assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) &&
           "Unexpected Endbr opcode");
  
    // If the MBB/I is empty or the current instruction is not ENDBR,
    // insert ENDBR instruction to the location of I.
    if (I == MBB.end() || I->getOpcode() != EndbrOpcode) {
      BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode));
      ++NumEndBranchAdded;
      return true;
    }
    return false;
  }
  
  static bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
    if (!MOp.isGlobal())
      return false;
    auto *CalleeFn = dyn_cast<Function>(MOp.getGlobal());
    if (!CalleeFn)
      return false;
    AttributeList Attrs = CalleeFn->getAttributes();
    if (Attrs.hasAttribute(AttributeList::FunctionIndex, Attribute::ReturnsTwice))
      return true;
    return false;
  }
  
  bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
    const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
  
    // Check that the cf-protection-branch is enabled.
    Metadata *isCFProtectionSupported =
        MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
    if (!isCFProtectionSupported && !IndirectBranchTracking)
      return false;
  
    // True if the current MF was changed and false otherwise.
    bool Changed = false;
  
    TII = SubTarget.getInstrInfo();
    EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
  
    // Non-internal function or function whose address was taken, can be
    // accessed through indirect calls. Mark the first BB with ENDBR instruction
    // unless nocf_check attribute is used.
    if ((MF.getFunction().hasAddressTaken() ||
         !MF.getFunction().hasLocalLinkage()) &&
        !MF.getFunction().doesNoCfCheck()) {
      auto MBB = MF.begin();
      Changed |= addENDBR(*MBB, MBB->begin());
    }
  
    for (auto &MBB : MF) {
      // Find all basic blocks that their address was taken (for example
      // in the case of indirect jump) and add ENDBR instruction.
      if (MBB.hasAddressTaken())
        Changed |= addENDBR(MBB, MBB.begin());
  
      for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
        if (!I->isCall())
          continue;
        if (IsCallReturnTwice(I->getOperand(0)))
          Changed |= addENDBR(MBB, std::next(I));
      }
    }
    return Changed;
  }
  