X86IndirectBranchTracking.cpp revision 363496 
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
27using namespace llvm;
28
29#define DEBUG_TYPE "x86-indirect-branch-tracking"
30
31static cl::opt<bool> IndirectBranchTracking(
32    "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
33    cl::desc("Enable X86 indirect branch tracking pass."));
34
35STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
36
37namespace {
38class X86IndirectBranchTrackingPass : public MachineFunctionPass {
39public:
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
48private:
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
66char X86IndirectBranchTrackingPass::ID = 0;
67
68FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
69  return new X86IndirectBranchTrackingPass();
70}
71
72bool 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
88static 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  if (Attrs.hasAttribute(AttributeList::FunctionIndex, Attribute::ReturnsTwice))
96    return true;
97  return false;
98}
99
100bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
101  const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
102
103  // Check that the cf-protection-branch is enabled.
104  Metadata *isCFProtectionSupported =
105      MF.getMMI().getModule()->getModuleFlag("cf-protection-branch");
106  // NB: We need to enable IBT in jitted code if JIT compiler is CET
107  // enabled.
108  const X86TargetMachine *TM =
109      static_cast<const X86TargetMachine *>(&MF.getTarget());
110#ifdef __CET__
111  bool isJITwithCET = TM->isJIT();
112#else
113  bool isJITwithCET = false;
114#endif
115  if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET)
116    return false;
117
118  // True if the current MF was changed and false otherwise.
119  bool Changed = false;
120
121  TII = SubTarget.getInstrInfo();
122  EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
123
124  // Large code model, non-internal function or function whose address
125  // was taken, can be accessed through indirect calls. Mark the first
126  // BB with ENDBR instruction unless nocf_check attribute is used.
127  if ((TM->getCodeModel() == CodeModel::Large ||
128       MF.getFunction().hasAddressTaken() ||
129       !MF.getFunction().hasLocalLinkage()) &&
130      !MF.getFunction().doesNoCfCheck()) {
131    auto MBB = MF.begin();
132    Changed |= addENDBR(*MBB, MBB->begin());
133  }
134
135  for (auto &MBB : MF) {
136    // Find all basic blocks that their address was taken (for example
137    // in the case of indirect jump) and add ENDBR instruction.
138    if (MBB.hasAddressTaken())
139      Changed |= addENDBR(MBB, MBB.begin());
140
141    // Exception handle may indirectly jump to catch pad, So we should add
142    // ENDBR before catch pad instructions.
143    bool EHPadIBTNeeded = MBB.isEHPad();
144
145    for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
146      if (I->isCall() && IsCallReturnTwice(I->getOperand(0)))
147        Changed |= addENDBR(MBB, std::next(I));
148
149      if (EHPadIBTNeeded && I->isEHLabel()) {
150        Changed |= addENDBR(MBB, std::next(I));
151        EHPadIBTNeeded = false;
152      }
153    }
154  }
155  return Changed;
156}
157