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 
26 using namespace llvm;
27 
28 #define DEBUG_TYPE "cfguard-longjmp"
29 
30 STATISTIC(CFGuardLongjmpTargets,
31           "Number of Control Flow Guard longjmp targets");
32 
33 namespace {
34 
35 /// MachineFunction pass to insert a symbol after each call to _setjmp and store
36 /// this in the MachineFunction's LongjmpTargets vector.
37 class CFGuardLongjmp : public MachineFunctionPass {
38 public:
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 
54 char CFGuardLongjmp::ID = 0;
55 
56 INITIALIZE_PASS(CFGuardLongjmp, "CFGuardLongjmp",
57                 "Insert symbols at valid longjmp targets for /guard:cf", false,
58                 false)
59 FunctionPass *llvm::createCFGuardLongjmpPass() { return new CFGuardLongjmp(); }
60 
61 bool 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