X86IndirectBranchTracking.cpp revision 360784 
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
26using namespace llvm;
27
28#define DEBUG_TYPE "x86-indirect-branch-tracking"
29
30static cl::opt<bool> IndirectBranchTracking(
31    "x86-indirect-branch-tracking", cl::init(false), cl::Hidden,
32    cl::desc("Enable X86 indirect branch tracking pass."));
33
34STATISTIC(NumEndBranchAdded, "Number of ENDBR instructions added");
35
36namespace {
37class X86IndirectBranchTrackingPass : public MachineFunctionPass {
38public:
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
47private:
48  static char ID;
49
50  /// Machine instruction info used throughout the class.
51  const X86InstrInfo *TII = nullptr;
52
53  /// Endbr opcode for the current machine function.
54  unsigned int EndbrOpcode = 0;
55
56  /// Adds a new ENDBR instruction to the beginning 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
65char X86IndirectBranchTrackingPass::ID = 0;
66
67FunctionPass *llvm::createX86IndirectBranchTrackingPass() {
68  return new X86IndirectBranchTrackingPass();
69}
70
71bool 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
87static 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
99bool 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}
139