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/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 }