1//===-- CFGuardLongjmp.cpp - Longjmp symbols for CFGuard --------*- C++ -*-===//
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/// \file
10/// This file contains a machine function pass to insert a symbol after each
11/// call to _setjmp and store this in the MachineFunction's LongjmpTargets
12/// vector. This will be used to emit the table of valid longjmp targets used
13/// by Control Flow Guard.
14///
15//===----------------------------------------------------------------------===//
16
17#include "llvm/ADT/Statistic.h"
18#include "llvm/CodeGen/MachineBasicBlock.h"
19#include "llvm/CodeGen/MachineFunctionPass.h"
20#include "llvm/CodeGen/MachineInstr.h"
21#include "llvm/CodeGen/MachineModuleInfo.h"
22#include "llvm/CodeGen/MachineOperand.h"
23#include "llvm/CodeGen/Passes.h"
24#include "llvm/InitializePasses.h"
25
26using namespace llvm;
27
28#define DEBUG_TYPE "cfguard-longjmp"
29
30STATISTIC(CFGuardLongjmpTargets,
31          "Number of Control Flow Guard longjmp targets");
32
33namespace {
34
35/// MachineFunction pass to insert a symbol after each call to _setjmp and store
36/// this in the MachineFunction's LongjmpTargets vector.
37class CFGuardLongjmp : public MachineFunctionPass {
38public:
39  static char ID;
40
41  CFGuardLongjmp() : MachineFunctionPass(ID) {
42    initializeCFGuardLongjmpPass(*PassRegistry::getPassRegistry());
43  }
44
45  StringRef getPassName() const override {
46    return "Control Flow Guard longjmp targets";
47  }
48
49  bool runOnMachineFunction(MachineFunction &MF) override;
50};
51
52} // end anonymous namespace
53
54char CFGuardLongjmp::ID = 0;
55
56INITIALIZE_PASS(CFGuardLongjmp, "CFGuardLongjmp",
57                "Insert symbols at valid longjmp targets for /guard:cf", false,
58                false)
59FunctionPass *llvm::createCFGuardLongjmpPass() { return new CFGuardLongjmp(); }
60
61bool CFGuardLongjmp::runOnMachineFunction(MachineFunction &MF) {
62
63  // Skip modules for which the cfguard flag is not set.
64  if (!MF.getMMI().getModule()->getModuleFlag("cfguard"))
65    return false;
66
67  // Skip functions that do not have calls to _setjmp.
68  if (!MF.getFunction().callsFunctionThatReturnsTwice())
69    return false;
70
71  SmallVector<MachineInstr *, 8> SetjmpCalls;
72
73  // Iterate over all instructions in the function and add calls to functions
74  // that return twice to the list of targets.
75  for (MachineBasicBlock &MBB : MF) {
76    for (MachineInstr &MI : MBB) {
77
78      // Skip instructions that are not calls.
79      if (!MI.isCall() || MI.getNumOperands() < 1)
80        continue;
81
82      // Iterate over operands to find calls to global functions.
83      for (MachineOperand &MO : MI.operands()) {
84        if (!MO.isGlobal())
85          continue;
86
87        auto *F = dyn_cast<Function>(MO.getGlobal());
88        if (!F)
89          continue;
90
91        // If the instruction calls a function that returns twice, add
92        // it to the list of targets.
93        if (F->hasFnAttribute(Attribute::ReturnsTwice)) {
94          SetjmpCalls.push_back(&MI);
95          break;
96        }
97      }
98    }
99  }
100
101  if (SetjmpCalls.empty())
102    return false;
103
104  unsigned SetjmpNum = 0;
105
106  // For each possible target, create a new symbol and insert it immediately
107  // after the call to setjmp. Add this symbol to the MachineFunction's list
108  // of longjmp targets.
109  for (MachineInstr *Setjmp : SetjmpCalls) {
110    SmallString<128> SymbolName;
111    raw_svector_ostream(SymbolName) << "$cfgsj_" << MF.getName() << SetjmpNum++;
112    MCSymbol *SjSymbol = MF.getContext().getOrCreateSymbol(SymbolName);
113
114    Setjmp->setPostInstrSymbol(MF, SjSymbol);
115    MF.addLongjmpTarget(SjSymbol);
116    CFGuardLongjmpTargets++;
117  }
118
119  return true;
120}
121