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