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
31cl::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  return Attrs.hasFnAttr(Attribute::ReturnsTwice);
96}
97
98// Checks if function should have an ENDBR in its prologue
99static bool needsPrologueENDBR(MachineFunction &MF, const Module *M) {
100  Function &F = MF.getFunction();
101
102  if (F.doesNoCfCheck())
103    return false;
104
105  switch (MF.getTarget().getCodeModel()) {
106  // Large code model functions always reachable through indirect calls.
107  case CodeModel::Large:
108    return true;
109  // Address taken or externally linked functions may be reachable.
110  default:
111    return (F.hasAddressTaken() || !F.hasLocalLinkage());
112  }
113}
114
115bool X86IndirectBranchTrackingPass::runOnMachineFunction(MachineFunction &MF) {
116  const X86Subtarget &SubTarget = MF.getSubtarget<X86Subtarget>();
117
118  const Module *M = MF.getMMI().getModule();
119  // Check that the cf-protection-branch is enabled.
120  Metadata *isCFProtectionSupported = M->getModuleFlag("cf-protection-branch");
121
122  //  NB: We need to enable IBT in jitted code if JIT compiler is CET
123  //  enabled.
124  const X86TargetMachine *TM =
125      static_cast<const X86TargetMachine *>(&MF.getTarget());
126#ifdef __CET__
127  bool isJITwithCET = TM->isJIT();
128#else
129  bool isJITwithCET = false;
130#endif
131  if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET)
132    return false;
133
134  // True if the current MF was changed and false otherwise.
135  bool Changed = false;
136
137  TII = SubTarget.getInstrInfo();
138  EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32;
139
140  // If function is reachable indirectly, mark the first BB with ENDBR.
141  if (needsPrologueENDBR(MF, M)) {
142    auto MBB = MF.begin();
143    Changed |= addENDBR(*MBB, MBB->begin());
144  }
145
146  for (auto &MBB : MF) {
147    // Find all basic blocks that their address was taken (for example
148    // in the case of indirect jump) and add ENDBR instruction.
149    if (MBB.hasAddressTaken())
150      Changed |= addENDBR(MBB, MBB.begin());
151
152    for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
153      if (I->isCall() && I->getNumOperands() > 0 &&
154          IsCallReturnTwice(I->getOperand(0))) {
155        Changed |= addENDBR(MBB, std::next(I));
156      }
157    }
158
159    // Exception handle may indirectly jump to catch pad, So we should add
160    // ENDBR before catch pad instructions. For SjLj exception model, it will
161    // create a new BB(new landingpad) indirectly jump to the old landingpad.
162    if (TM->Options.ExceptionModel == ExceptionHandling::SjLj) {
163      for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
164        // New Landingpad BB without EHLabel.
165        if (MBB.isEHPad()) {
166          if (I->isDebugInstr())
167            continue;
168          Changed |= addENDBR(MBB, I);
169          break;
170        } else if (I->isEHLabel()) {
171          // Old Landingpad BB (is not Landingpad now) with
172          // the the old "callee" EHLabel.
173          MCSymbol *Sym = I->getOperand(0).getMCSymbol();
174          if (!MF.hasCallSiteLandingPad(Sym))
175            continue;
176          Changed |= addENDBR(MBB, std::next(I));
177          break;
178        }
179      }
180    } else if (MBB.isEHPad()){
181      for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
182        if (!I->isEHLabel())
183          continue;
184        Changed |= addENDBR(MBB, std::next(I));
185        break;
186      }
187    }
188  }
189  return Changed;
190}
191