llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

   1 //===-- X86LoadValueInjectionRetHardening.cpp - LVI RET hardening for x86 --==//
   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 /// Description: Replaces every `ret` instruction with the sequence:
  10 /// ```
  11 /// pop <scratch-reg>
  12 /// lfence
  13 /// jmp *<scratch-reg>
  14 /// ```
  15 /// where `<scratch-reg>` is some available scratch register, according to the
  16 /// calling convention of the function being mitigated.
  17 ///
  18 //===----------------------------------------------------------------------===//
  19
  20 #include "X86.h"
  21 #include "X86InstrBuilder.h"
  22 #include "X86Subtarget.h"
  23 #include "llvm/ADT/Statistic.h"
  24 #include "llvm/CodeGen/MachineBasicBlock.h"
  25 #include "llvm/CodeGen/MachineFunction.h"
  26 #include "llvm/CodeGen/MachineFunctionPass.h"
  27 #include "llvm/CodeGen/MachineInstrBuilder.h"
  28 #include "llvm/IR/Function.h"
  29 #include "llvm/Support/Debug.h"
  30
  31 using namespace llvm;
  32
  33 #define PASS_KEY "x86-lvi-ret"
  34 #define DEBUG_TYPE PASS_KEY
  35
  36 STATISTIC(NumFences, "Number of LFENCEs inserted for LVI mitigation");
  37 STATISTIC(NumFunctionsConsidered, "Number of functions analyzed");
  38 STATISTIC(NumFunctionsMitigated, "Number of functions for which mitigations "
  39                                  "were deployed");
  40
  41 namespace {
  42
  43 class X86LoadValueInjectionRetHardeningPass : public MachineFunctionPass {
  44 public:
  45   X86LoadValueInjectionRetHardeningPass() : MachineFunctionPass(ID) {}
  46   StringRef getPassName() const override {
  47     return "X86 Load Value Injection (LVI) Ret-Hardening";
  48   }
  49   bool runOnMachineFunction(MachineFunction &MF) override;
  50
  51   static char ID;
  52 };
  53
  54 } // end anonymous namespace
  55
  56 char X86LoadValueInjectionRetHardeningPass::ID = 0;
  57
  58 bool X86LoadValueInjectionRetHardeningPass::runOnMachineFunction(
  59     MachineFunction &MF) {
  60   LLVM_DEBUG(dbgs() << "***** " << getPassName() << " : " << MF.getName()
  61                     << " *****\n");
  62   const X86Subtarget *Subtarget = &MF.getSubtarget<X86Subtarget>();
  63   if (!Subtarget->useLVIControlFlowIntegrity() || !Subtarget->is64Bit())
  64     return false; // FIXME: support 32-bit
  65
  66   // Don't skip functions with the "optnone" attr but participate in opt-bisect.
  67   const Function &F = MF.getFunction();
  68   if (!F.hasOptNone() && skipFunction(F))
  69     return false;
  70
  71   ++NumFunctionsConsidered;
  72   const X86RegisterInfo *TRI = Subtarget->getRegisterInfo();
  73   const X86InstrInfo *TII = Subtarget->getInstrInfo();
  74
  75   bool Modified = false;
  76   for (auto &MBB : MF) {
  77     for (auto MBBI = MBB.begin(); MBBI != MBB.end(); ++MBBI) {
  78       if (MBBI->getOpcode() != X86::RET64)
  79         continue;
  80
  81       unsigned ClobberReg = TRI->findDeadCallerSavedReg(MBB, MBBI);
  82       if (ClobberReg != X86::NoRegister) {
  83         BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::POP64r))
  84             .addReg(ClobberReg, RegState::Define)
  85             .setMIFlag(MachineInstr::FrameDestroy);
  86         BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::LFENCE));
  87         BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::JMP64r))
  88             .addReg(ClobberReg);
  89         MBB.erase(MBBI);
  90       } else {
  91         // In case there is no available scratch register, we can still read
  92         // from RSP to assert that RSP points to a valid page. The write to RSP
  93         // is also helpful because it verifies that the stack's write
  94         // permissions are intact.
  95         MachineInstr *Fence =
  96             BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::LFENCE));
  97         addRegOffset(BuildMI(MBB, Fence, DebugLoc(), TII->get(X86::SHL64mi)),
  98                      X86::RSP, false, 0)
  99             .addImm(0)
 100             ->addRegisterDead(X86::EFLAGS, TRI);
 101       }
 102
 103       ++NumFences;
 104       Modified = true;
 105       break;
 106     }
 107   }
 108
 109   if (Modified)
 110     ++NumFunctionsMitigated;
 111   return Modified;
 112 }
 113
 114 INITIALIZE_PASS(X86LoadValueInjectionRetHardeningPass, PASS_KEY,
 115                 "X86 LVI ret hardener", false, false)
 116
 117 FunctionPass *llvm::createX86LoadValueInjectionRetHardeningPass() {
 118   return new X86LoadValueInjectionRetHardeningPass();
 119 }