1362593Sdim//===-- X86LoadValueInjectionRetHardening.cpp - LVI RET hardening for x86 --==// 2362593Sdim// 3362593Sdim// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4362593Sdim// See https://llvm.org/LICENSE.txt for license information. 5362593Sdim// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6362593Sdim// 7362593Sdim//===----------------------------------------------------------------------===// 8362593Sdim/// 9362593Sdim/// Description: Replaces every `ret` instruction with the sequence: 10362593Sdim/// ``` 11362593Sdim/// pop <scratch-reg> 12362593Sdim/// lfence 13362593Sdim/// jmp *<scratch-reg> 14362593Sdim/// ``` 15362593Sdim/// where `<scratch-reg>` is some available scratch register, according to the 16362593Sdim/// calling convention of the function being mitigated. 17362593Sdim/// 18362593Sdim//===----------------------------------------------------------------------===// 19362593Sdim 20362593Sdim#include "X86.h" 21362593Sdim#include "X86InstrBuilder.h" 22362593Sdim#include "X86Subtarget.h" 23362593Sdim#include "llvm/ADT/Statistic.h" 24362593Sdim#include "llvm/CodeGen/MachineBasicBlock.h" 25362593Sdim#include "llvm/CodeGen/MachineFunction.h" 26362593Sdim#include "llvm/CodeGen/MachineFunctionPass.h" 27362593Sdim#include "llvm/CodeGen/MachineInstrBuilder.h" 28362593Sdim#include "llvm/IR/Function.h" 29362593Sdim#include "llvm/Support/Debug.h" 30362593Sdim#include <bitset> 31362593Sdim 32362593Sdimusing namespace llvm; 33362593Sdim 34362593Sdim#define PASS_KEY "x86-lvi-ret" 35362593Sdim#define DEBUG_TYPE PASS_KEY 36362593Sdim 37362593SdimSTATISTIC(NumFences, "Number of LFENCEs inserted for LVI mitigation"); 38362593SdimSTATISTIC(NumFunctionsConsidered, "Number of functions analyzed"); 39362593SdimSTATISTIC(NumFunctionsMitigated, "Number of functions for which mitigations " 40362593Sdim "were deployed"); 41362593Sdim 42362593Sdimnamespace { 43362593Sdim 44362593Sdimclass X86LoadValueInjectionRetHardeningPass : public MachineFunctionPass { 45362593Sdimpublic: 46362593Sdim X86LoadValueInjectionRetHardeningPass() : MachineFunctionPass(ID) {} 47362593Sdim StringRef getPassName() const override { 48362593Sdim return "X86 Load Value Injection (LVI) Ret-Hardening"; 49362593Sdim } 50362593Sdim bool runOnMachineFunction(MachineFunction &MF) override; 51362593Sdim 52362593Sdim static char ID; 53362593Sdim}; 54362593Sdim 55362593Sdim} // end anonymous namespace 56362593Sdim 57362593Sdimchar X86LoadValueInjectionRetHardeningPass::ID = 0; 58362593Sdim 59362593Sdimbool X86LoadValueInjectionRetHardeningPass::runOnMachineFunction( 60362593Sdim MachineFunction &MF) { 61362593Sdim LLVM_DEBUG(dbgs() << "***** " << getPassName() << " : " << MF.getName() 62362593Sdim << " *****\n"); 63362593Sdim const X86Subtarget *Subtarget = &MF.getSubtarget<X86Subtarget>(); 64362593Sdim if (!Subtarget->useLVIControlFlowIntegrity() || !Subtarget->is64Bit()) 65362593Sdim return false; // FIXME: support 32-bit 66362593Sdim 67362593Sdim // Don't skip functions with the "optnone" attr but participate in opt-bisect. 68362593Sdim const Function &F = MF.getFunction(); 69362593Sdim if (!F.hasOptNone() && skipFunction(F)) 70362593Sdim return false; 71362593Sdim 72362593Sdim ++NumFunctionsConsidered; 73362593Sdim const X86RegisterInfo *TRI = Subtarget->getRegisterInfo(); 74362593Sdim const X86InstrInfo *TII = Subtarget->getInstrInfo(); 75362593Sdim unsigned ClobberReg = X86::NoRegister; 76362593Sdim std::bitset<X86::NUM_TARGET_REGS> UnclobberableGR64s; 77362593Sdim UnclobberableGR64s.set(X86::RSP); // can't clobber stack pointer 78362593Sdim UnclobberableGR64s.set(X86::RIP); // can't clobber instruction pointer 79362593Sdim UnclobberableGR64s.set(X86::RAX); // used for function return 80362593Sdim UnclobberableGR64s.set(X86::RDX); // used for function return 81362593Sdim 82362593Sdim // We can clobber any register allowed by the function's calling convention. 83362593Sdim for (const MCPhysReg *PR = TRI->getCalleeSavedRegs(&MF); auto Reg = *PR; ++PR) 84362593Sdim UnclobberableGR64s.set(Reg); 85362593Sdim for (auto &Reg : X86::GR64RegClass) { 86362593Sdim if (!UnclobberableGR64s.test(Reg)) { 87362593Sdim ClobberReg = Reg; 88362593Sdim break; 89362593Sdim } 90362593Sdim } 91362593Sdim 92362593Sdim if (ClobberReg != X86::NoRegister) { 93362593Sdim LLVM_DEBUG(dbgs() << "Selected register " 94362593Sdim << Subtarget->getRegisterInfo()->getRegAsmName(ClobberReg) 95362593Sdim << " to clobber\n"); 96362593Sdim } else { 97362593Sdim LLVM_DEBUG(dbgs() << "Could not find a register to clobber\n"); 98362593Sdim } 99362593Sdim 100362593Sdim bool Modified = false; 101362593Sdim for (auto &MBB : MF) { 102362593Sdim if (MBB.empty()) 103362593Sdim continue; 104362593Sdim 105362593Sdim MachineInstr &MI = MBB.back(); 106362593Sdim if (MI.getOpcode() != X86::RETQ) 107362593Sdim continue; 108362593Sdim 109362593Sdim if (ClobberReg != X86::NoRegister) { 110362593Sdim MBB.erase_instr(&MI); 111362593Sdim BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::POP64r)) 112362593Sdim .addReg(ClobberReg, RegState::Define) 113362593Sdim .setMIFlag(MachineInstr::FrameDestroy); 114362593Sdim BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::LFENCE)); 115362593Sdim BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::JMP64r)) 116362593Sdim .addReg(ClobberReg); 117362593Sdim } else { 118362593Sdim // In case there is no available scratch register, we can still read from 119362593Sdim // RSP to assert that RSP points to a valid page. The write to RSP is 120362593Sdim // also helpful because it verifies that the stack's write permissions 121362593Sdim // are intact. 122362593Sdim MachineInstr *Fence = BuildMI(MBB, MI, DebugLoc(), TII->get(X86::LFENCE)); 123362593Sdim addRegOffset(BuildMI(MBB, Fence, DebugLoc(), TII->get(X86::SHL64mi)), 124362593Sdim X86::RSP, false, 0) 125362593Sdim .addImm(0) 126362593Sdim ->addRegisterDead(X86::EFLAGS, TRI); 127362593Sdim } 128362593Sdim 129362593Sdim ++NumFences; 130362593Sdim Modified = true; 131362593Sdim } 132362593Sdim 133362593Sdim if (Modified) 134362593Sdim ++NumFunctionsMitigated; 135362593Sdim return Modified; 136362593Sdim} 137362593Sdim 138362593SdimINITIALIZE_PASS(X86LoadValueInjectionRetHardeningPass, PASS_KEY, 139362593Sdim "X86 LVI ret hardener", false, false) 140362593Sdim 141362593SdimFunctionPass *llvm::createX86LoadValueInjectionRetHardeningPass() { 142362593Sdim return new X86LoadValueInjectionRetHardeningPass(); 143362593Sdim} 144