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
43constexpr StringRef X86LVIRetPassName =
44 "X86 Load Value Injection (LVI) Ret-Hardening";
45
46class X86LoadValueInjectionRetHardeningLegacy : public MachineFunctionPass {
47public:
48 X86LoadValueInjectionRetHardeningLegacy() : MachineFunctionPass(ID) {}
49 StringRef getPassName() const override { return X86LVIRetPassName; }
50 bool runOnMachineFunction(MachineFunction &MF) override;
51
52 static char ID;
53};
54
55} // end anonymous namespace
56
57char X86LoadValueInjectionRetHardeningLegacy::ID = 0;
58
59static bool runX86LoadValueInjectionRetHardening(MachineFunction &MF) {
60 const X86Subtarget *Subtarget = &MF.getSubtarget<X86Subtarget>();
61 if (!Subtarget->useLVIControlFlowIntegrity() || !Subtarget->is64Bit())
62 return false; // FIXME: support 32-bit
63
64 LLVM_DEBUG(dbgs() << "***** " << X86LVIRetPassName << " : " << MF.getName()
65 << " *****\n");
66 ++NumFunctionsConsidered;
67 const X86RegisterInfo *TRI = Subtarget->getRegisterInfo();
68 const X86InstrInfo *TII = Subtarget->getInstrInfo();
69
70 bool Modified = false;
71 for (auto &MBB : MF) {
72 for (auto MBBI = MBB.begin(); MBBI != MBB.end(); ++MBBI) {
73 if (MBBI->getOpcode() != X86::RET64)
74 continue;
75
76 unsigned ClobberReg = TRI->findDeadCallerSavedReg(MBB, MBBI);
77 if (ClobberReg != X86::NoRegister) {
78 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::POP64r))
79 .addReg(RegNo: ClobberReg, Flags: RegState::Define)
80 .setMIFlag(MachineInstr::FrameDestroy);
81 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::LFENCE));
82 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::JMP64r))
83 .addReg(RegNo: ClobberReg);
84 MBB.erase(I: MBBI);
85 } else {
86 // In case there is no available scratch register, we can still read
87 // from RSP to assert that RSP points to a valid page. The write to RSP
88 // is also helpful because it verifies that the stack's write
89 // permissions are intact.
90 MachineInstr *Fence =
91 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::LFENCE));
92 addRegOffset(MIB: BuildMI(BB&: MBB, I: Fence, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::SHL64mi)),
93 Reg: X86::RSP, isKill: false, Offset: 0)
94 .addImm(Val: 0)
95 ->addRegisterDead(Reg: X86::EFLAGS, RegInfo: TRI);
96 }
97
98 ++NumFences;
99 Modified = true;
100 break;
101 }
102 }
103
104 if (Modified)
105 ++NumFunctionsMitigated;
106 return Modified;
107}
108
109bool X86LoadValueInjectionRetHardeningLegacy::runOnMachineFunction(
110 MachineFunction &MF) {
111 // Don't skip functions with the "optnone" attr but participate in opt-bisect.
112 // Note: NewPM implements this behavior by default.
113 const Function &F = MF.getFunction();
114 if (!F.hasOptNone() && skipFunction(F))
115 return false;
116
117 return runX86LoadValueInjectionRetHardening(MF);
118}
119
120PreservedAnalyses X86LoadValueInjectionRetHardeningPass::run(
121 MachineFunction &MF, MachineFunctionAnalysisManager &MFAM) {
122 return runX86LoadValueInjectionRetHardening(MF)
123 ? getMachineFunctionPassPreservedAnalyses()
124 .preserveSet<CFGAnalyses>()
125 : PreservedAnalyses::all();
126}
127
128INITIALIZE_PASS(X86LoadValueInjectionRetHardeningLegacy, PASS_KEY,
129 "X86 LVI ret hardener", false, false)
130
131FunctionPass *llvm::createX86LoadValueInjectionRetHardeningLegacyPass() {
132 return new X86LoadValueInjectionRetHardeningLegacy();
133}
134