llvm/lib/CodeGen/CFGuardLongjmp.cpp

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