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
31using namespace llvm;
32
33#define PASS_KEY "x86-lvi-ret"
34#define DEBUG_TYPE PASS_KEY
35
36STATISTIC(NumFences, "Number of LFENCEs inserted for LVI mitigation");
37STATISTIC(NumFunctionsConsidered, "Number of functions analyzed");
38STATISTIC(NumFunctionsMitigated, "Number of functions for which mitigations "
39 "were deployed");
40
41namespace {
42
43class X86LoadValueInjectionRetHardeningPass : public MachineFunctionPass {
44public:
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
56char X86LoadValueInjectionRetHardeningPass::ID = 0;
57
58bool X86LoadValueInjectionRetHardeningPass::runOnMachineFunction(
59 MachineFunction &MF) {
60 const X86Subtarget *Subtarget = &MF.getSubtarget<X86Subtarget>();
61 if (!Subtarget->useLVIControlFlowIntegrity() || !Subtarget->is64Bit())
62 return false; // FIXME: support 32-bit
63
64 // Don't skip functions with the "optnone" attr but participate in opt-bisect.
65 const Function &F = MF.getFunction();
66 if (!F.hasOptNone() && skipFunction(F))
67 return false;
68
69 LLVM_DEBUG(dbgs() << "***** " << getPassName() << " : " << MF.getName()
70 << " *****\n");
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(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::POP64r))
84 .addReg(RegNo: ClobberReg, flags: RegState::Define)
85 .setMIFlag(MachineInstr::FrameDestroy);
86 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::LFENCE));
87 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::JMP64r))
88 .addReg(RegNo: ClobberReg);
89 MBB.erase(I: 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(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::LFENCE));
97 addRegOffset(MIB: BuildMI(BB&: MBB, I: Fence, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::SHL64mi)),
98 Reg: X86::RSP, isKill: false, Offset: 0)
99 .addImm(Val: 0)
100 ->addRegisterDead(Reg: X86::EFLAGS, RegInfo: TRI);
101 }
102
103 ++NumFences;
104 Modified = true;
105 break;
106 }
107 }
108
109 if (Modified)
110 ++NumFunctionsMitigated;
111 return Modified;
112}
113
114INITIALIZE_PASS(X86LoadValueInjectionRetHardeningPass, PASS_KEY,
115 "X86 LVI ret hardener", false, false)
116
117FunctionPass *llvm::createX86LoadValueInjectionRetHardeningPass() {
118 return new X86LoadValueInjectionRetHardeningPass();
119}
120