1//====- X86SpeculativeLoadHardening.cpp - A Spectre v1 mitigation ---------===//
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/// \file
9///
10/// Provide a pass which mitigates speculative execution attacks which operate
11/// by speculating incorrectly past some predicate (a type check, bounds check,
12/// or other condition) to reach a load with invalid inputs and leak the data
13/// accessed by that load using a side channel out of the speculative domain.
14///
15/// For details on the attacks, see the first variant in both the Project Zero
16/// writeup and the Spectre paper:
17/// https://googleprojectzero.blogspot.com/2018/01/reading-privileged-memory-with-side.html
18/// https://spectreattack.com/spectre.pdf
19///
20//===----------------------------------------------------------------------===//
21
22#include "X86.h"
23#include "X86InstrInfo.h"
24#include "X86Subtarget.h"
25#include "llvm/ADT/ArrayRef.h"
26#include "llvm/ADT/DenseMap.h"
27#include "llvm/ADT/STLExtras.h"
28#include "llvm/ADT/SmallPtrSet.h"
29#include "llvm/ADT/SmallSet.h"
30#include "llvm/ADT/SmallVector.h"
31#include "llvm/ADT/SparseBitVector.h"
32#include "llvm/ADT/Statistic.h"
33#include "llvm/CodeGen/MachineBasicBlock.h"
34#include "llvm/CodeGen/MachineConstantPool.h"
35#include "llvm/CodeGen/MachineFunction.h"
36#include "llvm/CodeGen/MachineFunctionPass.h"
37#include "llvm/CodeGen/MachineInstr.h"
38#include "llvm/CodeGen/MachineInstrBuilder.h"
39#include "llvm/CodeGen/MachineModuleInfo.h"
40#include "llvm/CodeGen/MachineOperand.h"
41#include "llvm/CodeGen/MachineRegisterInfo.h"
42#include "llvm/CodeGen/MachineSSAUpdater.h"
43#include "llvm/CodeGen/RegisterClassInfo.h"
44#include "llvm/CodeGen/TargetInstrInfo.h"
45#include "llvm/CodeGen/TargetRegisterInfo.h"
46#include "llvm/CodeGen/TargetSchedule.h"
47#include "llvm/CodeGen/TargetSubtargetInfo.h"
48#include "llvm/IR/DebugLoc.h"
49#include "llvm/MC/MCSchedule.h"
50#include "llvm/Pass.h"
51#include "llvm/Support/Debug.h"
52#include "llvm/Support/raw_ostream.h"
53#include "llvm/Target/TargetMachine.h"
54#include <cassert>
55#include <iterator>
56#include <optional>
57
58using namespace llvm;
59
60#define PASS_KEY "x86-slh"
61#define DEBUG_TYPE PASS_KEY
62
63STATISTIC(NumCondBranchesTraced, "Number of conditional branches traced");
64STATISTIC(NumBranchesUntraced, "Number of branches unable to trace");
65STATISTIC(NumAddrRegsHardened,
66 "Number of address mode used registers hardened");
67STATISTIC(NumPostLoadRegsHardened,
68 "Number of post-load register values hardened");
69STATISTIC(NumCallsOrJumpsHardened,
70 "Number of calls or jumps requiring extra hardening");
71STATISTIC(NumInstsInserted, "Number of instructions inserted");
72STATISTIC(NumLFENCEsInserted, "Number of lfence instructions inserted");
73
74namespace {
75
76constexpr StringRef X86SLHPassName = "X86 speculative load hardening";
77
78class X86SpeculativeLoadHardeningLegacy : public MachineFunctionPass {
79public:
80 X86SpeculativeLoadHardeningLegacy() : MachineFunctionPass(ID) {}
81
82 StringRef getPassName() const override { return X86SLHPassName; }
83 bool runOnMachineFunction(MachineFunction &MF) override;
84 void getAnalysisUsage(AnalysisUsage &AU) const override;
85
86 /// Pass identification, replacement for typeid.
87 static char ID;
88};
89
90class X86SpeculativeLoadHardeningImpl {
91public:
92 X86SpeculativeLoadHardeningImpl() = default;
93
94 bool run(MachineFunction &MF);
95
96private:
97 /// The information about a block's conditional terminators needed to trace
98 /// our predicate state through the exiting edges.
99 struct BlockCondInfo {
100 MachineBasicBlock *MBB;
101
102 // We mostly have one conditional branch, and in extremely rare cases have
103 // two. Three and more are so rare as to be unimportant for compile time.
104 SmallVector<MachineInstr *, 2> CondBrs;
105
106 MachineInstr *UncondBr;
107 };
108
109 /// Manages the predicate state traced through the program.
110 struct PredState {
111 Register InitialReg;
112 Register PoisonReg;
113
114 const TargetRegisterClass *RC;
115 MachineSSAUpdater SSA;
116
117 PredState(MachineFunction &MF, const TargetRegisterClass *RC)
118 : RC(RC), SSA(MF) {}
119 };
120
121 const X86Subtarget *Subtarget = nullptr;
122 MachineRegisterInfo *MRI = nullptr;
123 const X86InstrInfo *TII = nullptr;
124 const TargetRegisterInfo *TRI = nullptr;
125
126 std::optional<PredState> PS;
127
128 void hardenEdgesWithLFENCE(MachineFunction &MF);
129
130 SmallVector<BlockCondInfo, 16> collectBlockCondInfo(MachineFunction &MF);
131
132 SmallVector<MachineInstr *, 16>
133 tracePredStateThroughCFG(MachineFunction &MF, ArrayRef<BlockCondInfo> Infos);
134
135 void unfoldCallAndJumpLoads(MachineFunction &MF);
136
137 SmallVector<MachineInstr *, 16>
138 tracePredStateThroughIndirectBranches(MachineFunction &MF);
139
140 void tracePredStateThroughBlocksAndHarden(MachineFunction &MF);
141
142 Register saveEFLAGS(MachineBasicBlock &MBB,
143 MachineBasicBlock::iterator InsertPt,
144 const DebugLoc &Loc);
145 void restoreEFLAGS(MachineBasicBlock &MBB,
146 MachineBasicBlock::iterator InsertPt, const DebugLoc &Loc,
147 Register Reg);
148
149 void mergePredStateIntoSP(MachineBasicBlock &MBB,
150 MachineBasicBlock::iterator InsertPt,
151 const DebugLoc &Loc, Register PredStateReg);
152 Register extractPredStateFromSP(MachineBasicBlock &MBB,
153 MachineBasicBlock::iterator InsertPt,
154 const DebugLoc &Loc);
155
156 void
157 hardenLoadAddr(MachineInstr &MI, MachineOperand &BaseMO,
158 MachineOperand &IndexMO,
159 SmallDenseMap<Register, Register, 32> &AddrRegToHardenedReg);
160 MachineInstr *
161 sinkPostLoadHardenedInst(MachineInstr &MI,
162 SmallPtrSetImpl<MachineInstr *> &HardenedInstrs);
163 bool canHardenRegister(Register Reg);
164 Register hardenValueInRegister(Register Reg, MachineBasicBlock &MBB,
165 MachineBasicBlock::iterator InsertPt,
166 const DebugLoc &Loc);
167 Register hardenPostLoad(MachineInstr &MI);
168 void hardenReturnInstr(MachineInstr &MI);
169 void tracePredStateThroughCall(MachineInstr &MI);
170 void hardenIndirectCallOrJumpInstr(
171 MachineInstr &MI,
172 SmallDenseMap<Register, Register, 32> &AddrRegToHardenedReg);
173};
174
175} // end anonymous namespace
176
177bool X86SpeculativeLoadHardeningLegacy::runOnMachineFunction(
178 MachineFunction &MF) {
179 X86SpeculativeLoadHardeningImpl Impl;
180 bool Changed = Impl.run(MF);
181 LLVM_DEBUG(dbgs() << "Final speculative load hardened function:\n"; MF.dump();
182 dbgs() << "\n"; MF.verify(this));
183 return Changed;
184}
185
186char X86SpeculativeLoadHardeningLegacy::ID = 0;
187
188void X86SpeculativeLoadHardeningLegacy::getAnalysisUsage(
189 AnalysisUsage &AU) const {
190 AU.addPreserved<MachineRegisterClassInfoWrapperPass>();
191 MachineFunctionPass::getAnalysisUsage(AU);
192}
193
194static MachineBasicBlock &splitEdge(MachineBasicBlock &MBB,
195 MachineBasicBlock &Succ, int SuccCount,
196 MachineInstr *Br, MachineInstr *&UncondBr,
197 const X86InstrInfo &TII) {
198 assert(!Succ.isEHPad() && "Shouldn't get edges to EH pads!");
199
200 MachineFunction &MF = *MBB.getParent();
201
202 MachineBasicBlock &NewMBB = *MF.CreateMachineBasicBlock();
203
204 // We have to insert the new block immediately after the current one as we
205 // don't know what layout-successor relationships the successor has and we
206 // may not be able to (and generally don't want to) try to fix those up.
207 MF.insert(MBBI: std::next(x: MachineFunction::iterator(&MBB)), MBB: &NewMBB);
208
209 // Update the branch instruction if necessary.
210 if (Br) {
211 assert(Br->getOperand(0).getMBB() == &Succ &&
212 "Didn't start with the right target!");
213 Br->getOperand(i: 0).setMBB(&NewMBB);
214
215 // If this successor was reached through a branch rather than fallthrough,
216 // we might have *broken* fallthrough and so need to inject a new
217 // unconditional branch.
218 if (!UncondBr) {
219 MachineBasicBlock &OldLayoutSucc =
220 *std::next(x: MachineFunction::iterator(&NewMBB));
221 assert(MBB.isSuccessor(&OldLayoutSucc) &&
222 "Without an unconditional branch, the old layout successor should "
223 "be an actual successor!");
224 auto BrBuilder =
225 BuildMI(BB: &MBB, MIMD: DebugLoc(), MCID: TII.get(Opcode: X86::JMP_1)).addMBB(MBB: &OldLayoutSucc);
226 // Update the unconditional branch now that we've added one.
227 UncondBr = &*BrBuilder;
228 }
229
230 // Insert unconditional "jump Succ" instruction in the new block if
231 // necessary.
232 if (!NewMBB.isLayoutSuccessor(MBB: &Succ)) {
233 SmallVector<MachineOperand, 4> Cond;
234 TII.insertBranch(MBB&: NewMBB, TBB: &Succ, FBB: nullptr, Cond, DL: Br->getDebugLoc());
235 }
236 } else {
237 assert(!UncondBr &&
238 "Cannot have a branchless successor and an unconditional branch!");
239 assert(NewMBB.isLayoutSuccessor(&Succ) &&
240 "A non-branch successor must have been a layout successor before "
241 "and now is a layout successor of the new block.");
242 }
243
244 // If this is the only edge to the successor, we can just replace it in the
245 // CFG. Otherwise we need to add a new entry in the CFG for the new
246 // successor.
247 if (SuccCount == 1) {
248 MBB.replaceSuccessor(Old: &Succ, New: &NewMBB);
249 } else {
250 MBB.splitSuccessor(Old: &Succ, New: &NewMBB);
251 }
252
253 // Hook up the edge from the new basic block to the old successor in the CFG.
254 NewMBB.addSuccessor(Succ: &Succ);
255
256 // Fix PHI nodes in Succ so they refer to NewMBB instead of MBB.
257 for (MachineInstr &MI : Succ) {
258 if (!MI.isPHI())
259 break;
260 for (int OpIdx = 1, NumOps = MI.getNumOperands(); OpIdx < NumOps;
261 OpIdx += 2) {
262 MachineOperand &OpV = MI.getOperand(i: OpIdx);
263 MachineOperand &OpMBB = MI.getOperand(i: OpIdx + 1);
264 assert(OpMBB.isMBB() && "Block operand to a PHI is not a block!");
265 if (OpMBB.getMBB() != &MBB)
266 continue;
267
268 // If this is the last edge to the succesor, just replace MBB in the PHI
269 if (SuccCount == 1) {
270 OpMBB.setMBB(&NewMBB);
271 break;
272 }
273
274 // Otherwise, append a new pair of operands for the new incoming edge.
275 MI.addOperand(MF, Op: OpV);
276 MI.addOperand(MF, Op: MachineOperand::CreateMBB(MBB: &NewMBB));
277 break;
278 }
279 }
280
281 // Inherit live-ins from the successor
282 for (auto &LI : Succ.liveins())
283 NewMBB.addLiveIn(RegMaskPair: LI);
284
285 LLVM_DEBUG(dbgs() << " Split edge from '" << MBB.getName() << "' to '"
286 << Succ.getName() << "'.\n");
287 return NewMBB;
288}
289
290/// Removing duplicate PHI operands to leave the PHI in a canonical and
291/// predictable form.
292///
293/// FIXME: It's really frustrating that we have to do this, but SSA-form in MIR
294/// isn't what you might expect. We may have multiple entries in PHI nodes for
295/// a single predecessor. This makes CFG-updating extremely complex, so here we
296/// simplify all PHI nodes to a model even simpler than the IR's model: exactly
297/// one entry per predecessor, regardless of how many edges there are.
298static void canonicalizePHIOperands(MachineFunction &MF) {
299 SmallPtrSet<MachineBasicBlock *, 4> Preds;
300 SmallVector<int, 4> DupIndices;
301 for (auto &MBB : MF)
302 for (auto &MI : MBB) {
303 if (!MI.isPHI())
304 break;
305
306 // First we scan the operands of the PHI looking for duplicate entries
307 // a particular predecessor. We retain the operand index of each duplicate
308 // entry found.
309 for (int OpIdx = 1, NumOps = MI.getNumOperands(); OpIdx < NumOps;
310 OpIdx += 2)
311 if (!Preds.insert(Ptr: MI.getOperand(i: OpIdx + 1).getMBB()).second)
312 DupIndices.push_back(Elt: OpIdx);
313
314 // Now walk the duplicate indices, removing both the block and value. Note
315 // that these are stored as a vector making this element-wise removal
316 // potentially quadratic.
317 //
318 // FIXME: It is really frustrating that we have to use a quadratic
319 // removal algorithm here. There should be a better way, but the use-def
320 // updates required make that impossible using the public API.
321 //
322 // Note that we have to process these backwards so that we don't
323 // invalidate other indices with each removal.
324 while (!DupIndices.empty()) {
325 int OpIdx = DupIndices.pop_back_val();
326 // Remove both the block and value operand, again in reverse order to
327 // preserve indices.
328 MI.removeOperand(OpNo: OpIdx + 1);
329 MI.removeOperand(OpNo: OpIdx);
330 }
331
332 Preds.clear();
333 }
334}
335
336/// Helper to scan a function for loads vulnerable to misspeculation that we
337/// want to harden.
338///
339/// We use this to avoid making changes to functions where there is nothing we
340/// need to do to harden against misspeculation.
341static bool hasVulnerableLoad(MachineFunction &MF) {
342 for (MachineBasicBlock &MBB : MF) {
343 for (MachineInstr &MI : MBB) {
344 // Loads within this basic block after an LFENCE are not at risk of
345 // speculatively executing with invalid predicates from prior control
346 // flow. So break out of this block but continue scanning the function.
347 if (MI.getOpcode() == X86::LFENCE)
348 break;
349
350 // Looking for loads only.
351 if (!MI.mayLoad())
352 continue;
353
354 // An MFENCE is modeled as a load but isn't vulnerable to misspeculation.
355 if (MI.getOpcode() == X86::MFENCE)
356 continue;
357
358 // We found a load.
359 return true;
360 }
361 }
362
363 // No loads found.
364 return false;
365}
366
367bool X86SpeculativeLoadHardeningImpl::run(MachineFunction &MF) {
368 LLVM_DEBUG(dbgs() << "********** " << X86SLHPassName << " : " << MF.getName()
369 << " **********\n");
370
371 // Only run if this pass is forced enabled or we detect the relevant function
372 // attribute requesting SLH.
373 Subtarget = &MF.getSubtarget<X86Subtarget>();
374 const X86Options &CLOpts = Subtarget->getCLOpts();
375 if (!CLOpts.speculative_load_hardening &&
376 !MF.getFunction().hasFnAttribute(Kind: Attribute::SpeculativeLoadHardening))
377 return false;
378
379 MRI = &MF.getRegInfo();
380 TII = Subtarget->getInstrInfo();
381 TRI = Subtarget->getRegisterInfo();
382
383 // FIXME: Support for 32-bit.
384 PS.emplace(args&: MF, args: &X86::GR64_NOSPRegClass);
385
386 if (MF.begin() == MF.end())
387 // Nothing to do for a degenerate empty function...
388 return false;
389
390 // We support an alternative hardening technique based on a debug flag.
391 if (CLOpts.slh_lfence) {
392 hardenEdgesWithLFENCE(MF);
393 return true;
394 }
395
396 // Create a dummy debug loc to use for all the generated code here.
397 DebugLoc Loc;
398
399 MachineBasicBlock &Entry = *MF.begin();
400 auto EntryInsertPt = Entry.SkipPHIsLabelsAndDebug(I: Entry.begin());
401
402 // Do a quick scan to see if we have any checkable loads.
403 bool HasVulnerableLoad = hasVulnerableLoad(MF);
404
405 // See if we have any conditional branching blocks that we will need to trace
406 // predicate state through.
407 SmallVector<BlockCondInfo, 16> Infos = collectBlockCondInfo(MF);
408
409 // If we have no interesting conditions or loads, nothing to do here.
410 if (!HasVulnerableLoad && Infos.empty())
411 return true;
412
413 // The poison value is required to be an all-ones value for many aspects of
414 // this mitigation.
415 const int PoisonVal = -1;
416 PS->PoisonReg = MRI->createVirtualRegister(RegClass: PS->RC);
417 BuildMI(BB&: Entry, I: EntryInsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::MOV64ri32), DestReg: PS->PoisonReg)
418 .addImm(Val: PoisonVal);
419 ++NumInstsInserted;
420
421 // If we have loads being hardened and we've asked for call and ret edges to
422 // get a full fence-based mitigation, inject that fence.
423 if (HasVulnerableLoad && CLOpts.slh_fence_call_and_ret) {
424 // We need to insert an LFENCE at the start of the function to suspend any
425 // incoming misspeculation from the caller. This helps two-fold: the caller
426 // may not have been protected as this code has been, and this code gets to
427 // not take any specific action to protect across calls.
428 // FIXME: We could skip this for functions which unconditionally return
429 // a constant.
430 BuildMI(BB&: Entry, I: EntryInsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::LFENCE));
431 ++NumInstsInserted;
432 ++NumLFENCEsInserted;
433 }
434
435 // If we guarded the entry with an LFENCE and have no conditionals to protect
436 // in blocks, then we're done.
437 if (CLOpts.slh_fence_call_and_ret && Infos.empty())
438 // We may have changed the function's code at this point to insert fences.
439 return true;
440
441 // For every basic block in the function which can b
442 if (CLOpts.slh_ip && !CLOpts.slh_fence_call_and_ret) {
443 // Set up the predicate state by extracting it from the incoming stack
444 // pointer so we pick up any misspeculation in our caller.
445 PS->InitialReg = extractPredStateFromSP(MBB&: Entry, InsertPt: EntryInsertPt, Loc);
446 } else {
447 // Otherwise, just build the predicate state itself by zeroing a register
448 // as we don't need any initial state.
449 PS->InitialReg = MRI->createVirtualRegister(RegClass: PS->RC);
450 Register PredStateSubReg = MRI->createVirtualRegister(RegClass: &X86::GR32RegClass);
451 auto ZeroI = BuildMI(BB&: Entry, I: EntryInsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::MOV32r0),
452 DestReg: PredStateSubReg);
453 ++NumInstsInserted;
454 MachineOperand *ZeroEFLAGSDefOp =
455 ZeroI->findRegisterDefOperand(Reg: X86::EFLAGS, /*TRI=*/nullptr);
456 assert(ZeroEFLAGSDefOp && ZeroEFLAGSDefOp->isImplicit() &&
457 "Must have an implicit def of EFLAGS!");
458 ZeroEFLAGSDefOp->setIsDead(true);
459 BuildMI(BB&: Entry, I: EntryInsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::SUBREG_TO_REG),
460 DestReg: PS->InitialReg)
461 .addReg(RegNo: PredStateSubReg)
462 .addImm(Val: X86::sub_32bit);
463 }
464
465 // We're going to need to trace predicate state throughout the function's
466 // CFG. Prepare for this by setting up our initial state of PHIs with unique
467 // predecessor entries and all the initial predicate state.
468 canonicalizePHIOperands(MF);
469
470 // Track the updated values in an SSA updater to rewrite into SSA form at the
471 // end.
472 PS->SSA.Initialize(V: PS->InitialReg);
473 PS->SSA.AddAvailableValue(BB: &Entry, V: PS->InitialReg);
474
475 // Trace through the CFG.
476 auto CMovs = tracePredStateThroughCFG(MF, Infos);
477
478 // We may also enter basic blocks in this function via exception handling
479 // control flow. Here, if we are hardening interprocedurally, we need to
480 // re-capture the predicate state from the throwing code. In the Itanium ABI,
481 // the throw will always look like a call to __cxa_throw and will have the
482 // predicate state in the stack pointer, so extract fresh predicate state from
483 // the stack pointer and make it available in SSA.
484 // FIXME: Handle non-itanium ABI EH models.
485 if (CLOpts.slh_ip) {
486 for (MachineBasicBlock &MBB : MF) {
487 assert(!MBB.isEHScopeEntry() && "Only Itanium ABI EH supported!");
488 assert(!MBB.isEHFuncletEntry() && "Only Itanium ABI EH supported!");
489 assert(!MBB.isCleanupFuncletEntry() && "Only Itanium ABI EH supported!");
490 if (!MBB.isEHPad())
491 continue;
492 PS->SSA.AddAvailableValue(
493 BB: &MBB,
494 V: extractPredStateFromSP(MBB, InsertPt: MBB.SkipPHIsAndLabels(I: MBB.begin()), Loc));
495 }
496 }
497
498 if (CLOpts.slh_indirect) {
499 // If we are going to harden calls and jumps we need to unfold their memory
500 // operands.
501 unfoldCallAndJumpLoads(MF);
502
503 // Then we trace predicate state through the indirect branches.
504 auto IndirectBrCMovs = tracePredStateThroughIndirectBranches(MF);
505 CMovs.append(in_start: IndirectBrCMovs.begin(), in_end: IndirectBrCMovs.end());
506 }
507
508 // Now that we have the predicate state available at the start of each block
509 // in the CFG, trace it through each block, hardening vulnerable instructions
510 // as we go.
511 tracePredStateThroughBlocksAndHarden(MF);
512
513 // Now rewrite all the uses of the pred state using the SSA updater to insert
514 // PHIs connecting the state between blocks along the CFG edges.
515 for (MachineInstr *CMovI : CMovs)
516 for (MachineOperand &Op : CMovI->operands()) {
517 if (!Op.isReg() || Op.getReg() != PS->InitialReg)
518 continue;
519
520 PS->SSA.RewriteUse(U&: Op);
521 }
522
523 return true;
524}
525
526/// Implements the naive hardening approach of putting an LFENCE after every
527/// potentially mis-predicted control flow construct.
528///
529/// We include this as an alternative mostly for the purpose of comparison. The
530/// performance impact of this is expected to be extremely severe and not
531/// practical for any real-world users.
532void X86SpeculativeLoadHardeningImpl::hardenEdgesWithLFENCE(
533 MachineFunction &MF) {
534 // First, we scan the function looking for blocks that are reached along edges
535 // that we might want to harden.
536 SmallSetVector<MachineBasicBlock *, 8> Blocks;
537 for (MachineBasicBlock &MBB : MF) {
538 // If there are no or only one successor, nothing to do here.
539 if (MBB.succ_size() <= 1)
540 continue;
541
542 // Skip blocks unless their terminators start with a branch. Other
543 // terminators don't seem interesting for guarding against misspeculation.
544 auto TermIt = MBB.getFirstTerminator();
545 if (TermIt == MBB.end() || !TermIt->isBranch())
546 continue;
547
548 // Add all the non-EH-pad succossors to the blocks we want to harden. We
549 // skip EH pads because there isn't really a condition of interest on
550 // entering.
551 for (MachineBasicBlock *SuccMBB : MBB.successors())
552 if (!SuccMBB->isEHPad())
553 Blocks.insert(X: SuccMBB);
554 }
555
556 for (MachineBasicBlock *MBB : Blocks) {
557 auto InsertPt = MBB->SkipPHIsAndLabels(I: MBB->begin());
558 BuildMI(BB&: *MBB, I: InsertPt, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::LFENCE));
559 ++NumInstsInserted;
560 ++NumLFENCEsInserted;
561 }
562}
563
564SmallVector<X86SpeculativeLoadHardeningImpl::BlockCondInfo, 16>
565X86SpeculativeLoadHardeningImpl::collectBlockCondInfo(MachineFunction &MF) {
566 SmallVector<BlockCondInfo, 16> Infos;
567
568 // Walk the function and build up a summary for each block's conditions that
569 // we need to trace through.
570 for (MachineBasicBlock &MBB : MF) {
571 // If there are no or only one successor, nothing to do here.
572 if (MBB.succ_size() <= 1)
573 continue;
574
575 // We want to reliably handle any conditional branch terminators in the
576 // MBB, so we manually analyze the branch. We can handle all of the
577 // permutations here, including ones that analyze branch cannot.
578 //
579 // The approach is to walk backwards across the terminators, resetting at
580 // any unconditional non-indirect branch, and track all conditional edges
581 // to basic blocks as well as the fallthrough or unconditional successor
582 // edge. For each conditional edge, we track the target and the opposite
583 // condition code in order to inject a "no-op" cmov into that successor
584 // that will harden the predicate. For the fallthrough/unconditional
585 // edge, we inject a separate cmov for each conditional branch with
586 // matching condition codes. This effectively implements an "and" of the
587 // condition flags, even if there isn't a single condition flag that would
588 // directly implement that. We don't bother trying to optimize either of
589 // these cases because if such an optimization is possible, LLVM should
590 // have optimized the conditional *branches* in that way already to reduce
591 // instruction count. This late, we simply assume the minimal number of
592 // branch instructions is being emitted and use that to guide our cmov
593 // insertion.
594
595 BlockCondInfo Info = {.MBB: &MBB, .CondBrs: {}, .UncondBr: nullptr};
596
597 // Now walk backwards through the terminators and build up successors they
598 // reach and the conditions.
599 for (MachineInstr &MI : llvm::reverse(C&: MBB)) {
600 // Once we've handled all the terminators, we're done.
601 if (!MI.isTerminator())
602 break;
603
604 // If we see a non-branch terminator, we can't handle anything so bail.
605 if (!MI.isBranch()) {
606 Info.CondBrs.clear();
607 break;
608 }
609
610 // If we see an unconditional branch, reset our state, clear any
611 // fallthrough, and set this is the "else" successor.
612 if (MI.getOpcode() == X86::JMP_1) {
613 Info.CondBrs.clear();
614 Info.UncondBr = &MI;
615 continue;
616 }
617
618 // If we get an invalid condition, we have an indirect branch or some
619 // other unanalyzable "fallthrough" case. We model this as a nullptr for
620 // the destination so we can still guard any conditional successors.
621 // Consider code sequences like:
622 // ```
623 // jCC L1
624 // jmpq *%rax
625 // ```
626 // We still want to harden the edge to `L1`.
627 if (X86::getCondFromBranch(MI) == X86::COND_INVALID) {
628 Info.CondBrs.clear();
629 Info.UncondBr = &MI;
630 continue;
631 }
632
633 // We have a vanilla conditional branch, add it to our list.
634 Info.CondBrs.push_back(Elt: &MI);
635 }
636 if (Info.CondBrs.empty()) {
637 ++NumBranchesUntraced;
638 LLVM_DEBUG(dbgs() << "WARNING: unable to secure successors of block:\n";
639 MBB.dump());
640 continue;
641 }
642
643 Infos.push_back(Elt: Info);
644 }
645
646 return Infos;
647}
648
649/// Trace the predicate state through the CFG, instrumenting each conditional
650/// branch such that misspeculation through an edge will poison the predicate
651/// state.
652///
653/// Returns the list of inserted CMov instructions so that they can have their
654/// uses of the predicate state rewritten into proper SSA form once it is
655/// complete.
656SmallVector<MachineInstr *, 16>
657X86SpeculativeLoadHardeningImpl::tracePredStateThroughCFG(
658 MachineFunction &MF, ArrayRef<BlockCondInfo> Infos) {
659 // Collect the inserted cmov instructions so we can rewrite their uses of the
660 // predicate state into SSA form.
661 SmallVector<MachineInstr *, 16> CMovs;
662
663 // Now walk all of the basic blocks looking for ones that end in conditional
664 // jumps where we need to update this register along each edge.
665 for (const BlockCondInfo &Info : Infos) {
666 MachineBasicBlock &MBB = *Info.MBB;
667 const SmallVectorImpl<MachineInstr *> &CondBrs = Info.CondBrs;
668 MachineInstr *UncondBr = Info.UncondBr;
669
670 LLVM_DEBUG(dbgs() << "Tracing predicate through block: " << MBB.getName()
671 << "\n");
672 ++NumCondBranchesTraced;
673
674 // Compute the non-conditional successor as either the target of any
675 // unconditional branch or the layout successor.
676 MachineBasicBlock *UncondSucc =
677 UncondBr ? (UncondBr->getOpcode() == X86::JMP_1
678 ? UncondBr->getOperand(i: 0).getMBB()
679 : nullptr)
680 : &*std::next(x: MachineFunction::iterator(&MBB));
681
682 // Count how many edges there are to any given successor.
683 SmallDenseMap<MachineBasicBlock *, int> SuccCounts;
684 if (UncondSucc)
685 ++SuccCounts[UncondSucc];
686 for (auto *CondBr : CondBrs)
687 ++SuccCounts[CondBr->getOperand(i: 0).getMBB()];
688
689 // A lambda to insert cmov instructions into a block checking all of the
690 // condition codes in a sequence.
691 auto BuildCheckingBlockForSuccAndConds =
692 [&](MachineBasicBlock &MBB, MachineBasicBlock &Succ, int SuccCount,
693 MachineInstr *Br, MachineInstr *&UncondBr,
694 ArrayRef<X86::CondCode> Conds) {
695 // First, we split the edge to insert the checking block into a safe
696 // location.
697 auto &CheckingMBB =
698 (SuccCount == 1 && Succ.pred_size() == 1)
699 ? Succ
700 : splitEdge(MBB, Succ, SuccCount, Br, UncondBr, TII: *TII);
701
702 bool LiveEFLAGS = Succ.isLiveIn(Reg: X86::EFLAGS);
703 if (!LiveEFLAGS)
704 CheckingMBB.addLiveIn(PhysReg: X86::EFLAGS);
705
706 // Now insert the cmovs to implement the checks.
707 auto InsertPt = CheckingMBB.begin();
708 assert((InsertPt == CheckingMBB.end() || !InsertPt->isPHI()) &&
709 "Should never have a PHI in the initial checking block as it "
710 "always has a single predecessor!");
711
712 // We will wire each cmov to each other, but need to start with the
713 // incoming pred state.
714 Register CurStateReg = PS->InitialReg;
715
716 for (X86::CondCode Cond : Conds) {
717 int PredStateSizeInBytes = TRI->getRegSizeInBits(RC: *PS->RC) / 8;
718 auto CMovOp = X86::getCMovOpcode(RegBytes: PredStateSizeInBytes);
719
720 Register UpdatedStateReg = MRI->createVirtualRegister(RegClass: PS->RC);
721 // Note that we intentionally use an empty debug location so that
722 // this picks up the preceding location.
723 auto CMovI = BuildMI(BB&: CheckingMBB, I: InsertPt, MIMD: DebugLoc(),
724 MCID: TII->get(Opcode: CMovOp), DestReg: UpdatedStateReg)
725 .addReg(RegNo: CurStateReg)
726 .addReg(RegNo: PS->PoisonReg)
727 .addImm(Val: Cond);
728 // If this is the last cmov and the EFLAGS weren't originally
729 // live-in, mark them as killed.
730 if (!LiveEFLAGS && Cond == Conds.back())
731 CMovI->findRegisterUseOperand(Reg: X86::EFLAGS, /*TRI=*/nullptr)
732 ->setIsKill(true);
733
734 ++NumInstsInserted;
735 LLVM_DEBUG(dbgs() << " Inserting cmov: "; CMovI->dump();
736 dbgs() << "\n");
737
738 // The first one of the cmovs will be using the top level
739 // `PredStateReg` and need to get rewritten into SSA form.
740 if (CurStateReg == PS->InitialReg)
741 CMovs.push_back(Elt: &*CMovI);
742
743 // The next cmov should start from this one's def.
744 CurStateReg = UpdatedStateReg;
745 }
746
747 // And put the last one into the available values for SSA form of our
748 // predicate state.
749 PS->SSA.AddAvailableValue(BB: &CheckingMBB, V: CurStateReg);
750 };
751
752 std::vector<X86::CondCode> UncondCodeSeq;
753 for (auto *CondBr : CondBrs) {
754 MachineBasicBlock &Succ = *CondBr->getOperand(i: 0).getMBB();
755 int &SuccCount = SuccCounts[&Succ];
756
757 X86::CondCode Cond = X86::getCondFromBranch(MI: *CondBr);
758 X86::CondCode InvCond = X86::GetOppositeBranchCondition(CC: Cond);
759 UncondCodeSeq.push_back(x: Cond);
760
761 BuildCheckingBlockForSuccAndConds(MBB, Succ, SuccCount, CondBr, UncondBr,
762 {InvCond});
763
764 // Decrement the successor count now that we've split one of the edges.
765 // We need to keep the count of edges to the successor accurate in order
766 // to know above when to *replace* the successor in the CFG vs. just
767 // adding the new successor.
768 --SuccCount;
769 }
770
771 // Since we may have split edges and changed the number of successors,
772 // normalize the probabilities. This avoids doing it each time we split an
773 // edge.
774 MBB.normalizeSuccProbs();
775
776 // Finally, we need to insert cmovs into the "fallthrough" edge. Here, we
777 // need to intersect the other condition codes. We can do this by just
778 // doing a cmov for each one.
779 if (!UncondSucc)
780 // If we have no fallthrough to protect (perhaps it is an indirect jump?)
781 // just skip this and continue.
782 continue;
783
784 assert(SuccCounts[UncondSucc] == 1 &&
785 "We should never have more than one edge to the unconditional "
786 "successor at this point because every other edge must have been "
787 "split above!");
788
789 // Sort and unique the codes to minimize them.
790 llvm::sort(C&: UncondCodeSeq);
791 UncondCodeSeq.erase(first: llvm::unique(R&: UncondCodeSeq), last: UncondCodeSeq.end());
792
793 // Build a checking version of the successor.
794 BuildCheckingBlockForSuccAndConds(MBB, *UncondSucc, /*SuccCount*/ 1,
795 UncondBr, UncondBr, UncondCodeSeq);
796 }
797
798 return CMovs;
799}
800
801/// Compute the register class for the unfolded load.
802///
803/// FIXME: This should probably live in X86InstrInfo, potentially by adding
804/// a way to unfold into a newly created vreg rather than requiring a register
805/// input.
806static const TargetRegisterClass *
807getRegClassForUnfoldedLoad(const X86InstrInfo &TII, unsigned Opcode) {
808 unsigned Index;
809 unsigned UnfoldedOpc = TII.getOpcodeAfterMemoryUnfold(
810 Opc: Opcode, /*UnfoldLoad*/ true, /*UnfoldStore*/ false, LoadRegIndex: &Index);
811 const MCInstrDesc &MCID = TII.get(Opcode: UnfoldedOpc);
812 return TII.getRegClass(MCID, OpNum: Index);
813}
814
815void X86SpeculativeLoadHardeningImpl::unfoldCallAndJumpLoads(
816 MachineFunction &MF) {
817 for (MachineBasicBlock &MBB : MF)
818 // We use make_early_inc_range here so we can remove instructions if needed
819 // without disturbing the iteration.
820 for (MachineInstr &MI : llvm::make_early_inc_range(Range: MBB.instrs())) {
821 // Must either be a call or a branch.
822 if (!MI.isCall() && !MI.isBranch())
823 continue;
824 // We only care about loading variants of these instructions.
825 if (!MI.mayLoad())
826 continue;
827
828 switch (MI.getOpcode()) {
829 default: {
830 LLVM_DEBUG(
831 dbgs() << "ERROR: Found an unexpected loading branch or call "
832 "instruction:\n";
833 MI.dump(); dbgs() << "\n");
834 report_fatal_error(reason: "Unexpected loading branch or call!");
835 }
836
837 case X86::FARCALL16m:
838 case X86::FARCALL32m:
839 case X86::FARCALL64m:
840 case X86::FARJMP16m:
841 case X86::FARJMP32m:
842 case X86::FARJMP64m:
843 // We cannot mitigate far jumps or calls, but we also don't expect them
844 // to be vulnerable to Spectre v1.2 style attacks.
845 continue;
846
847 case X86::CALL16m:
848 case X86::CALL16m_NT:
849 case X86::CALL32m:
850 case X86::CALL32m_NT:
851 case X86::CALL64m:
852 case X86::CALL64m_NT:
853 case X86::JMP16m:
854 case X86::JMP16m_NT:
855 case X86::JMP32m:
856 case X86::JMP32m_NT:
857 case X86::JMP64m:
858 case X86::JMP64m_NT:
859 case X86::TAILJMPm64:
860 case X86::TAILJMPm64_REX:
861 case X86::TAILJMPm:
862 case X86::TCRETURNmi64:
863 case X86::TCRETURN_WINmi64:
864 case X86::TCRETURNmi: {
865 // Use the generic unfold logic now that we know we're dealing with
866 // expected instructions.
867 // FIXME: We don't have test coverage for all of these!
868 auto *UnfoldedRC = getRegClassForUnfoldedLoad(TII: *TII, Opcode: MI.getOpcode());
869 if (!UnfoldedRC) {
870 LLVM_DEBUG(dbgs()
871 << "ERROR: Unable to unfold load from instruction:\n";
872 MI.dump(); dbgs() << "\n");
873 report_fatal_error(reason: "Unable to unfold load!");
874 }
875 Register Reg = MRI->createVirtualRegister(RegClass: UnfoldedRC);
876 SmallVector<MachineInstr *, 2> NewMIs;
877 // If we were able to compute an unfolded reg class, any failure here
878 // is just a programming error so just assert.
879 bool Unfolded =
880 TII->unfoldMemoryOperand(MF, MI, Reg, /*UnfoldLoad*/ true,
881 /*UnfoldStore*/ false, NewMIs);
882 (void)Unfolded;
883 assert(Unfolded &&
884 "Computed unfolded register class but failed to unfold");
885 // Now stitch the new instructions into place and erase the old one.
886 for (auto *NewMI : NewMIs)
887 MBB.insert(I: MI.getIterator(), M: NewMI);
888
889 // Update the call info.
890 if (MI.isCandidateForAdditionalCallInfo())
891 MF.eraseAdditionalCallInfo(MI: &MI);
892
893 MI.eraseFromParent();
894 LLVM_DEBUG({
895 dbgs() << "Unfolded load successfully into:\n";
896 for (auto *NewMI : NewMIs) {
897 NewMI->dump();
898 dbgs() << "\n";
899 }
900 });
901 continue;
902 }
903 }
904 llvm_unreachable("Escaped switch with default!");
905 }
906}
907
908/// Trace the predicate state through indirect branches, instrumenting them to
909/// poison the state if a target is reached that does not match the expected
910/// target.
911///
912/// This is designed to mitigate Spectre variant 1 attacks where an indirect
913/// branch is trained to predict a particular target and then mispredicts that
914/// target in a way that can leak data. Despite using an indirect branch, this
915/// is really a variant 1 style attack: it does not steer execution to an
916/// arbitrary or attacker controlled address, and it does not require any
917/// special code executing next to the victim. This attack can also be mitigated
918/// through retpolines, but those require either replacing indirect branches
919/// with conditional direct branches or lowering them through a device that
920/// blocks speculation. This mitigation can replace these retpoline-style
921/// mitigations for jump tables and other indirect branches within a function
922/// when variant 2 isn't a risk while allowing limited speculation. Indirect
923/// calls, however, cannot be mitigated through this technique without changing
924/// the ABI in a fundamental way.
925SmallVector<MachineInstr *, 16>
926X86SpeculativeLoadHardeningImpl::tracePredStateThroughIndirectBranches(
927 MachineFunction &MF) {
928 // We use the SSAUpdater to insert PHI nodes for the target addresses of
929 // indirect branches. We don't actually need the full power of the SSA updater
930 // in this particular case as we always have immediately available values, but
931 // this avoids us having to re-implement the PHI construction logic.
932 MachineSSAUpdater TargetAddrSSA(MF);
933 TargetAddrSSA.Initialize(V: MRI->createVirtualRegister(RegClass: &X86::GR64RegClass));
934
935 // Track which blocks were terminated with an indirect branch.
936 SmallPtrSet<MachineBasicBlock *, 4> IndirectTerminatedMBBs;
937
938 // We need to know what blocks end up reached via indirect branches. We
939 // expect this to be a subset of those whose address is taken and so track it
940 // directly via the CFG.
941 SmallPtrSet<MachineBasicBlock *, 4> IndirectTargetMBBs;
942
943 // Walk all the blocks which end in an indirect branch and make the
944 // target address available.
945 for (MachineBasicBlock &MBB : MF) {
946 // Find the last terminator.
947 auto MII = MBB.instr_rbegin();
948 while (MII != MBB.instr_rend() && MII->isDebugInstr())
949 ++MII;
950 if (MII == MBB.instr_rend())
951 continue;
952 MachineInstr &TI = *MII;
953 if (!TI.isTerminator() || !TI.isBranch())
954 // No terminator or non-branch terminator.
955 continue;
956
957 Register TargetReg;
958
959 switch (TI.getOpcode()) {
960 default:
961 // Direct branch or conditional branch (leading to fallthrough).
962 continue;
963
964 case X86::FARJMP16m:
965 case X86::FARJMP32m:
966 case X86::FARJMP64m:
967 // We cannot mitigate far jumps or calls, but we also don't expect them
968 // to be vulnerable to Spectre v1.2 or v2 (self trained) style attacks.
969 continue;
970
971 case X86::JMP16m:
972 case X86::JMP16m_NT:
973 case X86::JMP32m:
974 case X86::JMP32m_NT:
975 case X86::JMP64m:
976 case X86::JMP64m_NT:
977 // Mostly as documentation.
978 report_fatal_error(reason: "Memory operand jumps should have been unfolded!");
979
980 case X86::JMP16r:
981 report_fatal_error(
982 reason: "Support for 16-bit indirect branches is not implemented.");
983 case X86::JMP32r:
984 report_fatal_error(
985 reason: "Support for 32-bit indirect branches is not implemented.");
986
987 case X86::JMP64r:
988 TargetReg = TI.getOperand(i: 0).getReg();
989 }
990
991 // We have definitely found an indirect branch. Verify that there are no
992 // preceding conditional branches as we don't yet support that.
993 if (llvm::any_of(Range: MBB.terminators(), P: [&](MachineInstr &OtherTI) {
994 return !OtherTI.isDebugInstr() && &OtherTI != &TI;
995 })) {
996 LLVM_DEBUG({
997 dbgs() << "ERROR: Found other terminators in a block with an indirect "
998 "branch! This is not yet supported! Terminator sequence:\n";
999 for (MachineInstr &MI : MBB.terminators()) {
1000 MI.dump();
1001 dbgs() << '\n';
1002 }
1003 });
1004 report_fatal_error(reason: "Unimplemented terminator sequence!");
1005 }
1006
1007 // Make the target register an available value for this block.
1008 TargetAddrSSA.AddAvailableValue(BB: &MBB, V: TargetReg);
1009 IndirectTerminatedMBBs.insert(Ptr: &MBB);
1010
1011 // Add all the successors to our target candidates.
1012 IndirectTargetMBBs.insert_range(R: MBB.successors());
1013 }
1014
1015 // Keep track of the cmov instructions we insert so we can return them.
1016 SmallVector<MachineInstr *, 16> CMovs;
1017
1018 // If we didn't find any indirect branches with targets, nothing to do here.
1019 if (IndirectTargetMBBs.empty())
1020 return CMovs;
1021
1022 // We found indirect branches and targets that need to be instrumented to
1023 // harden loads within them. Walk the blocks of the function (to get a stable
1024 // ordering) and instrument each target of an indirect branch.
1025 for (MachineBasicBlock &MBB : MF) {
1026 // Skip the blocks that aren't candidate targets.
1027 if (!IndirectTargetMBBs.count(Ptr: &MBB))
1028 continue;
1029
1030 // We don't expect EH pads to ever be reached via an indirect branch. If
1031 // this is desired for some reason, we could simply skip them here rather
1032 // than asserting.
1033 assert(!MBB.isEHPad() &&
1034 "Unexpected EH pad as target of an indirect branch!");
1035
1036 // We should never end up threading EFLAGS into a block to harden
1037 // conditional jumps as there would be an additional successor via the
1038 // indirect branch. As a consequence, all such edges would be split before
1039 // reaching here, and the inserted block will handle the EFLAGS-based
1040 // hardening.
1041 assert(!MBB.isLiveIn(X86::EFLAGS) &&
1042 "Cannot check within a block that already has live-in EFLAGS!");
1043
1044 // We can't handle having non-indirect edges into this block unless this is
1045 // the only successor and we can synthesize the necessary target address.
1046 for (MachineBasicBlock *Pred : MBB.predecessors()) {
1047 // If we've already handled this by extracting the target directly,
1048 // nothing to do.
1049 if (IndirectTerminatedMBBs.count(Ptr: Pred))
1050 continue;
1051
1052 // Otherwise, we have to be the only successor. We generally expect this
1053 // to be true as conditional branches should have had a critical edge
1054 // split already. We don't however need to worry about EH pad successors
1055 // as they'll happily ignore the target and their hardening strategy is
1056 // resilient to all ways in which they could be reached speculatively.
1057 if (!llvm::all_of(Range: Pred->successors(), P: [&](MachineBasicBlock *Succ) {
1058 return Succ->isEHPad() || Succ == &MBB;
1059 })) {
1060 LLVM_DEBUG({
1061 dbgs() << "ERROR: Found conditional entry to target of indirect "
1062 "branch!\n";
1063 Pred->dump();
1064 MBB.dump();
1065 });
1066 report_fatal_error(reason: "Cannot harden a conditional entry to a target of "
1067 "an indirect branch!");
1068 }
1069
1070 // Now we need to compute the address of this block and install it as a
1071 // synthetic target in the predecessor. We do this at the bottom of the
1072 // predecessor.
1073 auto InsertPt = Pred->getFirstTerminator();
1074 Register TargetReg = MRI->createVirtualRegister(RegClass: &X86::GR64RegClass);
1075 if (MF.getTarget().getCodeModel() == CodeModel::Small &&
1076 !Subtarget->isPositionIndependent()) {
1077 // Directly materialize it into an immediate.
1078 auto AddrI = BuildMI(BB&: *Pred, I: InsertPt, MIMD: DebugLoc(),
1079 MCID: TII->get(Opcode: X86::MOV64ri32), DestReg: TargetReg)
1080 .addMBB(MBB: &MBB);
1081 ++NumInstsInserted;
1082 (void)AddrI;
1083 LLVM_DEBUG(dbgs() << " Inserting mov: "; AddrI->dump();
1084 dbgs() << "\n");
1085 } else {
1086 auto AddrI = BuildMI(BB&: *Pred, I: InsertPt, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::LEA64r),
1087 DestReg: TargetReg)
1088 .addReg(/*Base*/ RegNo: X86::RIP)
1089 .addImm(/*Scale*/ Val: 1)
1090 .addReg(/*Index*/ RegNo: 0)
1091 .addMBB(MBB: &MBB)
1092 .addReg(/*Segment*/ RegNo: 0);
1093 ++NumInstsInserted;
1094 (void)AddrI;
1095 LLVM_DEBUG(dbgs() << " Inserting lea: "; AddrI->dump();
1096 dbgs() << "\n");
1097 }
1098 // And make this available.
1099 TargetAddrSSA.AddAvailableValue(BB: Pred, V: TargetReg);
1100 }
1101
1102 // Materialize the needed SSA value of the target. Note that we need the
1103 // middle of the block as this block might at the bottom have an indirect
1104 // branch back to itself. We can do this here because at this point, every
1105 // predecessor of this block has an available value. This is basically just
1106 // automating the construction of a PHI node for this target.
1107 Register TargetReg = TargetAddrSSA.GetValueInMiddleOfBlock(BB: &MBB);
1108
1109 // Insert a comparison of the incoming target register with this block's
1110 // address. This also requires us to mark the block as having its address
1111 // taken explicitly.
1112 MBB.setMachineBlockAddressTaken();
1113 auto InsertPt = MBB.SkipPHIsLabelsAndDebug(I: MBB.begin());
1114 if (MF.getTarget().getCodeModel() == CodeModel::Small &&
1115 !Subtarget->isPositionIndependent()) {
1116 // Check directly against a relocated immediate when we can.
1117 auto CheckI = BuildMI(BB&: MBB, I: InsertPt, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::CMP64ri32))
1118 .addReg(RegNo: TargetReg, Flags: RegState::Kill)
1119 .addMBB(MBB: &MBB);
1120 ++NumInstsInserted;
1121 (void)CheckI;
1122 LLVM_DEBUG(dbgs() << " Inserting cmp: "; CheckI->dump(); dbgs() << "\n");
1123 } else {
1124 // Otherwise compute the address into a register first.
1125 Register AddrReg = MRI->createVirtualRegister(RegClass: &X86::GR64RegClass);
1126 auto AddrI =
1127 BuildMI(BB&: MBB, I: InsertPt, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::LEA64r), DestReg: AddrReg)
1128 .addReg(/*Base*/ RegNo: X86::RIP)
1129 .addImm(/*Scale*/ Val: 1)
1130 .addReg(/*Index*/ RegNo: 0)
1131 .addMBB(MBB: &MBB)
1132 .addReg(/*Segment*/ RegNo: 0);
1133 ++NumInstsInserted;
1134 (void)AddrI;
1135 LLVM_DEBUG(dbgs() << " Inserting lea: "; AddrI->dump(); dbgs() << "\n");
1136 auto CheckI = BuildMI(BB&: MBB, I: InsertPt, MIMD: DebugLoc(), MCID: TII->get(Opcode: X86::CMP64rr))
1137 .addReg(RegNo: TargetReg, Flags: RegState::Kill)
1138 .addReg(RegNo: AddrReg, Flags: RegState::Kill);
1139 ++NumInstsInserted;
1140 (void)CheckI;
1141 LLVM_DEBUG(dbgs() << " Inserting cmp: "; CheckI->dump(); dbgs() << "\n");
1142 }
1143
1144 // Now cmov over the predicate if the comparison wasn't equal.
1145 int PredStateSizeInBytes = TRI->getRegSizeInBits(RC: *PS->RC) / 8;
1146 auto CMovOp = X86::getCMovOpcode(RegBytes: PredStateSizeInBytes);
1147 Register UpdatedStateReg = MRI->createVirtualRegister(RegClass: PS->RC);
1148 auto CMovI =
1149 BuildMI(BB&: MBB, I: InsertPt, MIMD: DebugLoc(), MCID: TII->get(Opcode: CMovOp), DestReg: UpdatedStateReg)
1150 .addReg(RegNo: PS->InitialReg)
1151 .addReg(RegNo: PS->PoisonReg)
1152 .addImm(Val: X86::COND_NE);
1153 CMovI->findRegisterUseOperand(Reg: X86::EFLAGS, /*TRI=*/nullptr)
1154 ->setIsKill(true);
1155 ++NumInstsInserted;
1156 LLVM_DEBUG(dbgs() << " Inserting cmov: "; CMovI->dump(); dbgs() << "\n");
1157 CMovs.push_back(Elt: &*CMovI);
1158
1159 // And put the new value into the available values for SSA form of our
1160 // predicate state.
1161 PS->SSA.AddAvailableValue(BB: &MBB, V: UpdatedStateReg);
1162 }
1163
1164 // Return all the newly inserted cmov instructions of the predicate state.
1165 return CMovs;
1166}
1167
1168// Returns true if the MI has EFLAGS as a register def operand and it's live,
1169// otherwise it returns false
1170static bool isEFLAGSDefLive(const MachineInstr &MI) {
1171 if (const MachineOperand *DefOp =
1172 MI.findRegisterDefOperand(Reg: X86::EFLAGS, /*TRI=*/nullptr)) {
1173 return !DefOp->isDead();
1174 }
1175 return false;
1176}
1177
1178static bool isEFLAGSLive(MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
1179 const TargetRegisterInfo &TRI) {
1180 // Check if EFLAGS are alive by seeing if there is a def of them or they
1181 // live-in, and then seeing if that def is in turn used.
1182 for (MachineInstr &MI : llvm::reverse(C: llvm::make_range(x: MBB.begin(), y: I))) {
1183 if (MachineOperand *DefOp =
1184 MI.findRegisterDefOperand(Reg: X86::EFLAGS, /*TRI=*/nullptr)) {
1185 // If the def is dead, then EFLAGS is not live.
1186 if (DefOp->isDead())
1187 return false;
1188
1189 // Otherwise we've def'ed it, and it is live.
1190 return true;
1191 }
1192 // While at this instruction, also check if we use and kill EFLAGS
1193 // which means it isn't live.
1194 if (MI.killsRegister(Reg: X86::EFLAGS, TRI: &TRI))
1195 return false;
1196 }
1197
1198 // If we didn't find anything conclusive (neither definitely alive or
1199 // definitely dead) return whether it lives into the block.
1200 return MBB.isLiveIn(Reg: X86::EFLAGS);
1201}
1202
1203/// Trace the predicate state through each of the blocks in the function,
1204/// hardening everything necessary along the way.
1205///
1206/// We call this routine once the initial predicate state has been established
1207/// for each basic block in the function in the SSA updater. This routine traces
1208/// it through the instructions within each basic block, and for non-returning
1209/// blocks informs the SSA updater about the final state that lives out of the
1210/// block. Along the way, it hardens any vulnerable instruction using the
1211/// currently valid predicate state. We have to do these two things together
1212/// because the SSA updater only works across blocks. Within a block, we track
1213/// the current predicate state directly and update it as it changes.
1214///
1215/// This operates in two passes over each block. First, we analyze the loads in
1216/// the block to determine which strategy will be used to harden them: hardening
1217/// the address or hardening the loaded value when loaded into a register
1218/// amenable to hardening. We have to process these first because the two
1219/// strategies may interact -- later hardening may change what strategy we wish
1220/// to use. We also will analyze data dependencies between loads and avoid
1221/// hardening those loads that are data dependent on a load with a hardened
1222/// address. We also skip hardening loads already behind an LFENCE as that is
1223/// sufficient to harden them against misspeculation.
1224///
1225/// Second, we actively trace the predicate state through the block, applying
1226/// the hardening steps we determined necessary in the first pass as we go.
1227///
1228/// These two passes are applied to each basic block. We operate one block at a
1229/// time to simplify reasoning about reachability and sequencing.
1230void X86SpeculativeLoadHardeningImpl::tracePredStateThroughBlocksAndHarden(
1231 MachineFunction &MF) {
1232 const X86Options &CLOpts = Subtarget->getCLOpts();
1233 SmallPtrSet<MachineInstr *, 16> HardenPostLoad;
1234 SmallPtrSet<MachineInstr *, 16> HardenLoadAddr;
1235
1236 SmallSet<Register, 16> HardenedAddrRegs;
1237
1238 SmallDenseMap<Register, Register, 32> AddrRegToHardenedReg;
1239
1240 // Track the set of load-dependent registers through the basic block. Because
1241 // the values of these registers have an existing data dependency on a loaded
1242 // value which we would have checked, we can omit any checks on them.
1243 SparseBitVector<> LoadDepRegs;
1244
1245 for (MachineBasicBlock &MBB : MF) {
1246 // The first pass over the block: collect all the loads which can have their
1247 // loaded value hardened and all the loads that instead need their address
1248 // hardened. During this walk we propagate load dependence for address
1249 // hardened loads and also look for LFENCE to stop hardening wherever
1250 // possible. When deciding whether or not to harden the loaded value or not,
1251 // we check to see if any registers used in the address will have been
1252 // hardened at this point and if so, harden any remaining address registers
1253 // as that often successfully re-uses hardened addresses and minimizes
1254 // instructions.
1255 //
1256 // FIXME: We should consider an aggressive mode where we continue to keep as
1257 // many loads value hardened even when some address register hardening would
1258 // be free (due to reuse).
1259 //
1260 // Note that we only need this pass if we are actually hardening loads.
1261 if (CLOpts.slh_loads)
1262 for (MachineInstr &MI : MBB) {
1263 // We naively assume that all def'ed registers of an instruction have
1264 // a data dependency on all of their operands.
1265 // FIXME: Do a more careful analysis of x86 to build a conservative
1266 // model here.
1267 if (llvm::any_of(Range: MI.uses(), P: [&](MachineOperand &Op) {
1268 return Op.isReg() && LoadDepRegs.test(Idx: Op.getReg().id());
1269 }))
1270 for (MachineOperand &Def : MI.defs())
1271 if (Def.isReg())
1272 LoadDepRegs.set(Def.getReg().id());
1273
1274 // Both Intel and AMD are guiding that they will change the semantics of
1275 // LFENCE to be a speculation barrier, so if we see an LFENCE, there is
1276 // no more need to guard things in this block.
1277 if (MI.getOpcode() == X86::LFENCE)
1278 break;
1279
1280 // If this instruction cannot load, nothing to do.
1281 if (!MI.mayLoad())
1282 continue;
1283
1284 // Some instructions which "load" are trivially safe or unimportant.
1285 if (MI.getOpcode() == X86::MFENCE)
1286 continue;
1287
1288 // Extract the memory operand information about this instruction.
1289 const int MemRefBeginIdx = X86::getFirstAddrOperandIdx(MI);
1290 if (MemRefBeginIdx < 0) {
1291 LLVM_DEBUG(dbgs()
1292 << "WARNING: unable to harden loading instruction: ";
1293 MI.dump());
1294 continue;
1295 }
1296
1297 MachineOperand &BaseMO =
1298 MI.getOperand(i: MemRefBeginIdx + X86::AddrBaseReg);
1299 MachineOperand &IndexMO =
1300 MI.getOperand(i: MemRefBeginIdx + X86::AddrIndexReg);
1301
1302 // If we have at least one (non-frame-index, non-RIP) register operand,
1303 // and neither operand is load-dependent, we need to check the load.
1304 Register BaseReg, IndexReg;
1305 if (!BaseMO.isFI() && BaseMO.getReg() != X86::RIP &&
1306 BaseMO.getReg().isValid())
1307 BaseReg = BaseMO.getReg();
1308 if (IndexMO.getReg().isValid())
1309 IndexReg = IndexMO.getReg();
1310
1311 if (!BaseReg && !IndexReg)
1312 // No register operands!
1313 continue;
1314
1315 // If any register operand is dependent, this load is dependent and we
1316 // needn't check it.
1317 // FIXME: Is this true in the case where we are hardening loads after
1318 // they complete? Unclear, need to investigate.
1319 if ((BaseReg && LoadDepRegs.test(Idx: BaseReg.id())) ||
1320 (IndexReg && LoadDepRegs.test(Idx: IndexReg.id())))
1321 continue;
1322
1323 // If post-load hardening is enabled, this load is compatible with
1324 // post-load hardening, and we aren't already going to harden one of the
1325 // address registers, queue it up to be hardened post-load. Notably,
1326 // even once hardened this won't introduce a useful dependency that
1327 // could prune out subsequent loads.
1328 if (CLOpts.slh_post_load && X86InstrInfo::isDataInvariantLoad(MI) &&
1329 !isEFLAGSDefLive(MI) && MI.getDesc().getNumDefs() == 1 &&
1330 MI.getOperand(i: 0).isReg() &&
1331 canHardenRegister(Reg: MI.getOperand(i: 0).getReg()) &&
1332 !HardenedAddrRegs.count(V: BaseReg) &&
1333 !HardenedAddrRegs.count(V: IndexReg)) {
1334 HardenPostLoad.insert(Ptr: &MI);
1335 HardenedAddrRegs.insert(V: MI.getOperand(i: 0).getReg());
1336 continue;
1337 }
1338
1339 // Record this instruction for address hardening and record its register
1340 // operands as being address-hardened.
1341 HardenLoadAddr.insert(Ptr: &MI);
1342 if (BaseReg)
1343 HardenedAddrRegs.insert(V: BaseReg);
1344 if (IndexReg)
1345 HardenedAddrRegs.insert(V: IndexReg);
1346
1347 for (MachineOperand &Def : MI.defs())
1348 if (Def.isReg())
1349 LoadDepRegs.set(Def.getReg().id());
1350 }
1351
1352 // Now re-walk the instructions in the basic block, and apply whichever
1353 // hardening strategy we have elected. Note that we do this in a second
1354 // pass specifically so that we have the complete set of instructions for
1355 // which we will do post-load hardening and can defer it in certain
1356 // circumstances.
1357 for (MachineInstr &MI : MBB) {
1358 if (CLOpts.slh_loads) {
1359 // We cannot both require hardening the def of a load and its address.
1360 assert(!(HardenLoadAddr.count(&MI) && HardenPostLoad.count(&MI)) &&
1361 "Requested to harden both the address and def of a load!");
1362
1363 // Check if this is a load whose address needs to be hardened.
1364 if (HardenLoadAddr.erase(Ptr: &MI)) {
1365 const int MemRefBeginIdx = X86::getFirstAddrOperandIdx(MI);
1366 assert(MemRefBeginIdx >= 0 && "Cannot have an invalid index here!");
1367
1368 MachineOperand &BaseMO =
1369 MI.getOperand(i: MemRefBeginIdx + X86::AddrBaseReg);
1370 MachineOperand &IndexMO =
1371 MI.getOperand(i: MemRefBeginIdx + X86::AddrIndexReg);
1372 hardenLoadAddr(MI, BaseMO, IndexMO, AddrRegToHardenedReg);
1373 continue;
1374 }
1375
1376 // Test if this instruction is one of our post load instructions (and
1377 // remove it from the set if so).
1378 if (HardenPostLoad.erase(Ptr: &MI)) {
1379 assert(!MI.isCall() && "Must not try to post-load harden a call!");
1380
1381 // If this is a data-invariant load and there is no EFLAGS
1382 // interference, we want to try and sink any hardening as far as
1383 // possible.
1384 if (X86InstrInfo::isDataInvariantLoad(MI) && !isEFLAGSDefLive(MI)) {
1385 // Sink the instruction we'll need to harden as far as we can down
1386 // the graph.
1387 MachineInstr *SunkMI = sinkPostLoadHardenedInst(MI, HardenedInstrs&: HardenPostLoad);
1388
1389 // If we managed to sink this instruction, update everything so we
1390 // harden that instruction when we reach it in the instruction
1391 // sequence.
1392 if (SunkMI != &MI) {
1393 // If in sinking there was no instruction needing to be hardened,
1394 // we're done.
1395 if (!SunkMI)
1396 continue;
1397
1398 // Otherwise, add this to the set of defs we harden.
1399 HardenPostLoad.insert(Ptr: SunkMI);
1400 continue;
1401 }
1402 }
1403
1404 Register HardenedReg = hardenPostLoad(MI);
1405
1406 // Mark the resulting hardened register as such so we don't re-harden.
1407 AddrRegToHardenedReg[HardenedReg] = HardenedReg;
1408
1409 continue;
1410 }
1411
1412 // Check for an indirect call or branch that may need its input hardened
1413 // even if we couldn't find the specific load used, or were able to
1414 // avoid hardening it for some reason. Note that here we cannot break
1415 // out afterward as we may still need to handle any call aspect of this
1416 // instruction.
1417 if ((MI.isCall() || MI.isBranch()) && CLOpts.slh_indirect)
1418 hardenIndirectCallOrJumpInstr(MI, AddrRegToHardenedReg);
1419 }
1420
1421 // After we finish hardening loads we handle interprocedural hardening if
1422 // enabled and relevant for this instruction.
1423 if (!CLOpts.slh_ip)
1424 continue;
1425 if (!MI.isCall() && !MI.isReturn())
1426 continue;
1427
1428 // If this is a direct return (IE, not a tail call) just directly harden
1429 // it.
1430 if (MI.isReturn() && !MI.isCall()) {
1431 hardenReturnInstr(MI);
1432 continue;
1433 }
1434
1435 // Otherwise we have a call. We need to handle transferring the predicate
1436 // state into a call and recovering it after the call returns (unless this
1437 // is a tail call).
1438 assert(MI.isCall() && "Should only reach here for calls!");
1439 tracePredStateThroughCall(MI);
1440 }
1441
1442 HardenPostLoad.clear();
1443 HardenLoadAddr.clear();
1444 HardenedAddrRegs.clear();
1445 AddrRegToHardenedReg.clear();
1446
1447 // Currently, we only track data-dependent loads within a basic block.
1448 // FIXME: We should see if this is necessary or if we could be more
1449 // aggressive here without opening up attack avenues.
1450 LoadDepRegs.clear();
1451 }
1452}
1453
1454/// Save EFLAGS into the returned GPR. This can in turn be restored with
1455/// `restoreEFLAGS`.
1456///
1457/// Note that LLVM can only lower very simple patterns of saved and restored
1458/// EFLAGS registers. The restore should always be within the same basic block
1459/// as the save so that no PHI nodes are inserted.
1460Register X86SpeculativeLoadHardeningImpl::saveEFLAGS(
1461 MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertPt,
1462 const DebugLoc &Loc) {
1463 // FIXME: Hard coding this to a 32-bit register class seems weird, but matches
1464 // what instruction selection does.
1465 Register Reg = MRI->createVirtualRegister(RegClass: &X86::GR32RegClass);
1466 // We directly copy the FLAGS register and rely on later lowering to clean
1467 // this up into the appropriate setCC instructions.
1468 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::COPY), DestReg: Reg).addReg(RegNo: X86::EFLAGS);
1469 ++NumInstsInserted;
1470 return Reg;
1471}
1472
1473/// Restore EFLAGS from the provided GPR. This should be produced by
1474/// `saveEFLAGS`.
1475///
1476/// This must be done within the same basic block as the save in order to
1477/// reliably lower.
1478void X86SpeculativeLoadHardeningImpl::restoreEFLAGS(
1479 MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertPt,
1480 const DebugLoc &Loc, Register Reg) {
1481 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::COPY), DestReg: X86::EFLAGS).addReg(RegNo: Reg);
1482 ++NumInstsInserted;
1483}
1484
1485/// Takes the current predicate state (in a register) and merges it into the
1486/// stack pointer. The state is essentially a single bit, but we merge this in
1487/// a way that won't form non-canonical pointers and also will be preserved
1488/// across normal stack adjustments.
1489void X86SpeculativeLoadHardeningImpl::mergePredStateIntoSP(
1490 MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertPt,
1491 const DebugLoc &Loc, Register PredStateReg) {
1492 Register TmpReg = MRI->createVirtualRegister(RegClass: PS->RC);
1493 // FIXME: This hard codes a shift distance based on the number of bits needed
1494 // to stay canonical on 64-bit. We should compute this somehow and support
1495 // 32-bit as part of that.
1496 auto ShiftI = BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::SHL64ri), DestReg: TmpReg)
1497 .addReg(RegNo: PredStateReg, Flags: RegState::Kill)
1498 .addImm(Val: 47);
1499 ShiftI->addRegisterDead(Reg: X86::EFLAGS, RegInfo: TRI);
1500 ++NumInstsInserted;
1501 auto OrI = BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::OR64rr), DestReg: X86::RSP)
1502 .addReg(RegNo: X86::RSP)
1503 .addReg(RegNo: TmpReg, Flags: RegState::Kill);
1504 OrI->addRegisterDead(Reg: X86::EFLAGS, RegInfo: TRI);
1505 ++NumInstsInserted;
1506}
1507
1508/// Extracts the predicate state stored in the high bits of the stack pointer.
1509Register X86SpeculativeLoadHardeningImpl::extractPredStateFromSP(
1510 MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertPt,
1511 const DebugLoc &Loc) {
1512 Register PredStateReg = MRI->createVirtualRegister(RegClass: PS->RC);
1513 Register TmpReg = MRI->createVirtualRegister(RegClass: PS->RC);
1514
1515 // We know that the stack pointer will have any preserved predicate state in
1516 // its high bit. We just want to smear this across the other bits. Turns out,
1517 // this is exactly what an arithmetic right shift does.
1518 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: TmpReg)
1519 .addReg(RegNo: X86::RSP);
1520 auto ShiftI =
1521 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::SAR64ri), DestReg: PredStateReg)
1522 .addReg(RegNo: TmpReg, Flags: RegState::Kill)
1523 .addImm(Val: TRI->getRegSizeInBits(RC: *PS->RC) - 1);
1524 ShiftI->addRegisterDead(Reg: X86::EFLAGS, RegInfo: TRI);
1525 ++NumInstsInserted;
1526
1527 return PredStateReg;
1528}
1529
1530void X86SpeculativeLoadHardeningImpl::hardenLoadAddr(
1531 MachineInstr &MI, MachineOperand &BaseMO, MachineOperand &IndexMO,
1532 SmallDenseMap<Register, Register, 32> &AddrRegToHardenedReg) {
1533 MachineBasicBlock &MBB = *MI.getParent();
1534 const DebugLoc &Loc = MI.getDebugLoc();
1535
1536 // Check if EFLAGS are alive by seeing if there is a def of them or they
1537 // live-in, and then seeing if that def is in turn used.
1538 bool EFLAGSLive = isEFLAGSLive(MBB, I: MI.getIterator(), TRI: *TRI);
1539
1540 SmallVector<MachineOperand *, 2> HardenOpRegs;
1541
1542 if (BaseMO.isFI()) {
1543 // A frame index is never a dynamically controllable load, so only
1544 // harden it if we're covering fixed address loads as well.
1545 LLVM_DEBUG(
1546 dbgs() << " Skipping hardening base of explicit stack frame load: ";
1547 MI.dump(); dbgs() << "\n");
1548 } else if (BaseMO.getReg() == X86::RSP) {
1549 // Some idempotent atomic operations are lowered directly to a locked
1550 // OR with 0 to the top of stack(or slightly offset from top) which uses an
1551 // explicit RSP register as the base.
1552 assert(IndexMO.getReg() == X86::NoRegister &&
1553 "Explicit RSP access with dynamic index!");
1554 LLVM_DEBUG(
1555 dbgs() << " Cannot harden base of explicit RSP offset in a load!");
1556 } else if (BaseMO.getReg() == X86::RIP ||
1557 BaseMO.getReg() == X86::NoRegister) {
1558 // For both RIP-relative addressed loads or absolute loads, we cannot
1559 // meaningfully harden them because the address being loaded has no
1560 // dynamic component.
1561 //
1562 // FIXME: When using a segment base (like TLS does) we end up with the
1563 // dynamic address being the base plus -1 because we can't mutate the
1564 // segment register here. This allows the signed 32-bit offset to point at
1565 // valid segment-relative addresses and load them successfully.
1566 LLVM_DEBUG(
1567 dbgs() << " Cannot harden base of "
1568 << (BaseMO.getReg() == X86::RIP ? "RIP-relative" : "no-base")
1569 << " address in a load!");
1570 } else {
1571 assert(BaseMO.isReg() &&
1572 "Only allowed to have a frame index or register base.");
1573 HardenOpRegs.push_back(Elt: &BaseMO);
1574 }
1575
1576 if (IndexMO.getReg() != X86::NoRegister &&
1577 (HardenOpRegs.empty() ||
1578 HardenOpRegs.front()->getReg() != IndexMO.getReg()))
1579 HardenOpRegs.push_back(Elt: &IndexMO);
1580
1581 assert((HardenOpRegs.size() == 1 || HardenOpRegs.size() == 2) &&
1582 "Should have exactly one or two registers to harden!");
1583 assert((HardenOpRegs.size() == 1 ||
1584 HardenOpRegs[0]->getReg() != HardenOpRegs[1]->getReg()) &&
1585 "Should not have two of the same registers!");
1586
1587 // Remove any registers that have alreaded been checked.
1588 llvm::erase_if(C&: HardenOpRegs, P: [&](MachineOperand *Op) {
1589 // See if this operand's register has already been checked.
1590 auto It = AddrRegToHardenedReg.find(Val: Op->getReg());
1591 if (It == AddrRegToHardenedReg.end())
1592 // Not checked, so retain this one.
1593 return false;
1594
1595 // Otherwise, we can directly update this operand and remove it.
1596 Op->setReg(It->second);
1597 return true;
1598 });
1599 // If there are none left, we're done.
1600 if (HardenOpRegs.empty())
1601 return;
1602
1603 // Compute the current predicate state.
1604 Register StateReg = PS->SSA.GetValueAtEndOfBlock(BB: &MBB);
1605
1606 auto InsertPt = MI.getIterator();
1607
1608 // If EFLAGS are live and we don't have access to instructions that avoid
1609 // clobbering EFLAGS we need to save and restore them. This in turn makes
1610 // the EFLAGS no longer live.
1611 Register FlagsReg;
1612 if (EFLAGSLive && !Subtarget->hasBMI2()) {
1613 EFLAGSLive = false;
1614 FlagsReg = saveEFLAGS(MBB, InsertPt, Loc);
1615 }
1616
1617 for (MachineOperand *Op : HardenOpRegs) {
1618 Register OpReg = Op->getReg();
1619 auto *OpRC = MRI->getRegClass(Reg: OpReg);
1620 Register TmpReg = MRI->createVirtualRegister(RegClass: OpRC);
1621
1622 // If this is a vector register, we'll need somewhat custom logic to handle
1623 // hardening it.
1624 if (!Subtarget->hasVLX() && (OpRC->hasSuperClassEq(RC: &X86::VR128RegClass) ||
1625 OpRC->hasSuperClassEq(RC: &X86::VR256RegClass))) {
1626 assert(Subtarget->hasAVX2() && "AVX2-specific register classes!");
1627 bool Is128Bit = OpRC->hasSuperClassEq(RC: &X86::VR128RegClass);
1628
1629 // Move our state into a vector register.
1630 // FIXME: We could skip this at the cost of longer encodings with AVX-512
1631 // but that doesn't seem likely worth it.
1632 Register VStateReg = MRI->createVirtualRegister(RegClass: &X86::VR128RegClass);
1633 auto MovI =
1634 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::VMOV64toPQIrr), DestReg: VStateReg)
1635 .addReg(RegNo: StateReg);
1636 (void)MovI;
1637 ++NumInstsInserted;
1638 LLVM_DEBUG(dbgs() << " Inserting mov: "; MovI->dump(); dbgs() << "\n");
1639
1640 // Broadcast it across the vector register.
1641 Register VBStateReg = MRI->createVirtualRegister(RegClass: OpRC);
1642 auto BroadcastI = BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc,
1643 MCID: TII->get(Opcode: Is128Bit ? X86::VPBROADCASTQrr
1644 : X86::VPBROADCASTQYrr),
1645 DestReg: VBStateReg)
1646 .addReg(RegNo: VStateReg);
1647 (void)BroadcastI;
1648 ++NumInstsInserted;
1649 LLVM_DEBUG(dbgs() << " Inserting broadcast: "; BroadcastI->dump();
1650 dbgs() << "\n");
1651
1652 // Merge our potential poison state into the value with a vector or.
1653 auto OrI =
1654 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc,
1655 MCID: TII->get(Opcode: Is128Bit ? X86::VPORrr : X86::VPORYrr), DestReg: TmpReg)
1656 .addReg(RegNo: VBStateReg)
1657 .addReg(RegNo: OpReg);
1658 (void)OrI;
1659 ++NumInstsInserted;
1660 LLVM_DEBUG(dbgs() << " Inserting or: "; OrI->dump(); dbgs() << "\n");
1661 } else if (OpRC->hasSuperClassEq(RC: &X86::VR128XRegClass) ||
1662 OpRC->hasSuperClassEq(RC: &X86::VR256XRegClass) ||
1663 OpRC->hasSuperClassEq(RC: &X86::VR512RegClass)) {
1664 assert(Subtarget->hasAVX512() && "AVX512-specific register classes!");
1665 bool Is128Bit = OpRC->hasSuperClassEq(RC: &X86::VR128XRegClass);
1666 bool Is256Bit = OpRC->hasSuperClassEq(RC: &X86::VR256XRegClass);
1667 if (Is128Bit || Is256Bit)
1668 assert(Subtarget->hasVLX() && "AVX512VL-specific register classes!");
1669
1670 // Broadcast our state into a vector register.
1671 Register VStateReg = MRI->createVirtualRegister(RegClass: OpRC);
1672 unsigned BroadcastOp = Is128Bit ? X86::VPBROADCASTQrZ128rr
1673 : Is256Bit ? X86::VPBROADCASTQrZ256rr
1674 : X86::VPBROADCASTQrZrr;
1675 auto BroadcastI =
1676 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: BroadcastOp), DestReg: VStateReg)
1677 .addReg(RegNo: StateReg);
1678 (void)BroadcastI;
1679 ++NumInstsInserted;
1680 LLVM_DEBUG(dbgs() << " Inserting broadcast: "; BroadcastI->dump();
1681 dbgs() << "\n");
1682
1683 // Merge our potential poison state into the value with a vector or.
1684 unsigned OrOp = Is128Bit ? X86::VPORQZ128rr
1685 : Is256Bit ? X86::VPORQZ256rr : X86::VPORQZrr;
1686 auto OrI = BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: OrOp), DestReg: TmpReg)
1687 .addReg(RegNo: VStateReg)
1688 .addReg(RegNo: OpReg);
1689 (void)OrI;
1690 ++NumInstsInserted;
1691 LLVM_DEBUG(dbgs() << " Inserting or: "; OrI->dump(); dbgs() << "\n");
1692 } else {
1693 // FIXME: Need to support GR32 here for 32-bit code.
1694 assert(OpRC->hasSuperClassEq(&X86::GR64RegClass) &&
1695 "Not a supported register class for address hardening!");
1696
1697 if (!EFLAGSLive) {
1698 // Merge our potential poison state into the value with an or.
1699 auto OrI = BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::OR64rr), DestReg: TmpReg)
1700 .addReg(RegNo: StateReg)
1701 .addReg(RegNo: OpReg);
1702 OrI->addRegisterDead(Reg: X86::EFLAGS, RegInfo: TRI);
1703 ++NumInstsInserted;
1704 LLVM_DEBUG(dbgs() << " Inserting or: "; OrI->dump(); dbgs() << "\n");
1705 } else {
1706 // We need to avoid touching EFLAGS so shift out all but the least
1707 // significant bit using the instruction that doesn't update flags.
1708 auto ShiftI =
1709 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::SHRX64rr), DestReg: TmpReg)
1710 .addReg(RegNo: OpReg)
1711 .addReg(RegNo: StateReg);
1712 (void)ShiftI;
1713 ++NumInstsInserted;
1714 LLVM_DEBUG(dbgs() << " Inserting shrx: "; ShiftI->dump();
1715 dbgs() << "\n");
1716 }
1717 }
1718
1719 // Record this register as checked and update the operand.
1720 assert(!AddrRegToHardenedReg.count(Op->getReg()) &&
1721 "Should not have checked this register yet!");
1722 AddrRegToHardenedReg[Op->getReg()] = TmpReg;
1723 Op->setReg(TmpReg);
1724 ++NumAddrRegsHardened;
1725 }
1726
1727 // And restore the flags if needed.
1728 if (FlagsReg)
1729 restoreEFLAGS(MBB, InsertPt, Loc, Reg: FlagsReg);
1730}
1731
1732MachineInstr *X86SpeculativeLoadHardeningImpl::sinkPostLoadHardenedInst(
1733 MachineInstr &InitialMI, SmallPtrSetImpl<MachineInstr *> &HardenedInstrs) {
1734 assert(X86InstrInfo::isDataInvariantLoad(InitialMI) &&
1735 "Cannot get here with a non-invariant load!");
1736 assert(!isEFLAGSDefLive(InitialMI) &&
1737 "Cannot get here with a data invariant load "
1738 "that interferes with EFLAGS!");
1739
1740 // See if we can sink hardening the loaded value.
1741 auto SinkCheckToSingleUse =
1742 [&](MachineInstr &MI) -> std::optional<MachineInstr *> {
1743 Register DefReg = MI.getOperand(i: 0).getReg();
1744
1745 // We need to find a single use which we can sink the check. We can
1746 // primarily do this because many uses may already end up checked on their
1747 // own.
1748 MachineInstr *SingleUseMI = nullptr;
1749 for (MachineInstr &UseMI : MRI->use_instructions(Reg: DefReg)) {
1750 // If we're already going to harden this use, it is data invariant, it
1751 // does not interfere with EFLAGS, and within our block.
1752 if (HardenedInstrs.count(Ptr: &UseMI)) {
1753 if (!X86InstrInfo::isDataInvariantLoad(MI&: UseMI) || isEFLAGSDefLive(MI: UseMI)) {
1754 // If we've already decided to harden a non-load, we must have sunk
1755 // some other post-load hardened instruction to it and it must itself
1756 // be data-invariant.
1757 assert(X86InstrInfo::isDataInvariant(UseMI) &&
1758 "Data variant instruction being hardened!");
1759 continue;
1760 }
1761
1762 // Otherwise, this is a load and the load component can't be data
1763 // invariant so check how this register is being used.
1764 const int MemRefBeginIdx = X86::getFirstAddrOperandIdx(MI: UseMI);
1765 assert(MemRefBeginIdx >= 0 &&
1766 "Should always have mem references here!");
1767
1768 MachineOperand &BaseMO =
1769 UseMI.getOperand(i: MemRefBeginIdx + X86::AddrBaseReg);
1770 MachineOperand &IndexMO =
1771 UseMI.getOperand(i: MemRefBeginIdx + X86::AddrIndexReg);
1772 if ((BaseMO.isReg() && BaseMO.getReg() == DefReg) ||
1773 (IndexMO.isReg() && IndexMO.getReg() == DefReg))
1774 // The load uses the register as part of its address making it not
1775 // invariant.
1776 return {};
1777
1778 continue;
1779 }
1780
1781 if (SingleUseMI)
1782 // We already have a single use, this would make two. Bail.
1783 return {};
1784
1785 // If this single use isn't data invariant, isn't in this block, or has
1786 // interfering EFLAGS, we can't sink the hardening to it.
1787 if (!X86InstrInfo::isDataInvariant(MI&: UseMI) || UseMI.getParent() != MI.getParent() ||
1788 isEFLAGSDefLive(MI: UseMI))
1789 return {};
1790
1791 // If this instruction defines multiple registers bail as we won't harden
1792 // all of them.
1793 if (UseMI.getDesc().getNumDefs() > 1)
1794 return {};
1795
1796 // If this register isn't a virtual register we can't walk uses of sanely,
1797 // just bail. Also check that its register class is one of the ones we
1798 // can harden.
1799 Register UseDefReg = UseMI.getOperand(i: 0).getReg();
1800 if (!canHardenRegister(Reg: UseDefReg))
1801 return {};
1802
1803 SingleUseMI = &UseMI;
1804 }
1805
1806 // If SingleUseMI is still null, there is no use that needs its own
1807 // checking. Otherwise, it is the single use that needs checking.
1808 return {SingleUseMI};
1809 };
1810
1811 MachineInstr *MI = &InitialMI;
1812 while (std::optional<MachineInstr *> SingleUse = SinkCheckToSingleUse(*MI)) {
1813 // Update which MI we're checking now.
1814 MI = *SingleUse;
1815 if (!MI)
1816 break;
1817 }
1818
1819 return MI;
1820}
1821
1822bool X86SpeculativeLoadHardeningImpl::canHardenRegister(Register Reg) {
1823 // We only support hardening virtual registers.
1824 if (!Reg.isVirtual())
1825 return false;
1826
1827 auto *RC = MRI->getRegClass(Reg);
1828 int RegBytes = TRI->getRegSizeInBits(RC: *RC) / 8;
1829 if (RegBytes > 8)
1830 // We don't support post-load hardening of vectors.
1831 return false;
1832
1833 unsigned RegIdx = Log2_32(Value: RegBytes);
1834 assert(RegIdx < 4 && "Unsupported register size");
1835
1836 // If this register class is explicitly constrained to a class that doesn't
1837 // require REX prefix, we may not be able to satisfy that constraint when
1838 // emitting the hardening instructions, so bail out here.
1839 // FIXME: This seems like a pretty lame hack. The way this comes up is when we
1840 // end up both with a NOREX and REX-only register as operands to the hardening
1841 // instructions. It would be better to fix that code to handle this situation
1842 // rather than hack around it in this way.
1843 const TargetRegisterClass *NOREXRegClasses[] = {
1844 &X86::GR8_NOREXRegClass, &X86::GR16_NOREXRegClass,
1845 &X86::GR32_NOREXRegClass, &X86::GR64_NOREXRegClass};
1846 if (RC == NOREXRegClasses[RegIdx])
1847 return false;
1848
1849 const TargetRegisterClass *GPRRegClasses[] = {
1850 &X86::GR8RegClass, &X86::GR16RegClass, &X86::GR32RegClass,
1851 &X86::GR64RegClass};
1852 return RC->hasSuperClassEq(RC: GPRRegClasses[RegIdx]);
1853}
1854
1855/// Harden a value in a register.
1856///
1857/// This is the low-level logic to fully harden a value sitting in a register
1858/// against leaking during speculative execution.
1859///
1860/// Unlike hardening an address that is used by a load, this routine is required
1861/// to hide *all* incoming bits in the register.
1862///
1863/// `Reg` must be a virtual register. Currently, it is required to be a GPR no
1864/// larger than the predicate state register. FIXME: We should support vector
1865/// registers here by broadcasting the predicate state.
1866///
1867/// The new, hardened virtual register is returned. It will have the same
1868/// register class as `Reg`.
1869Register X86SpeculativeLoadHardeningImpl::hardenValueInRegister(
1870 Register Reg, MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertPt,
1871 const DebugLoc &Loc) {
1872 assert(canHardenRegister(Reg) && "Cannot harden this register!");
1873
1874 auto *RC = MRI->getRegClass(Reg);
1875 int Bytes = TRI->getRegSizeInBits(RC: *RC) / 8;
1876 Register StateReg = PS->SSA.GetValueAtEndOfBlock(BB: &MBB);
1877 assert((Bytes == 1 || Bytes == 2 || Bytes == 4 || Bytes == 8) &&
1878 "Unknown register size");
1879
1880 // FIXME: Need to teach this about 32-bit mode.
1881 if (Bytes != 8) {
1882 unsigned SubRegImms[] = {X86::sub_8bit, X86::sub_16bit, X86::sub_32bit};
1883 unsigned SubRegImm = SubRegImms[Log2_32(Value: Bytes)];
1884 Register NarrowStateReg = MRI->createVirtualRegister(RegClass: RC);
1885 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: NarrowStateReg)
1886 .addReg(RegNo: StateReg, Flags: {}, SubReg: SubRegImm);
1887 StateReg = NarrowStateReg;
1888 }
1889
1890 Register FlagsReg;
1891 if (isEFLAGSLive(MBB, I: InsertPt, TRI: *TRI))
1892 FlagsReg = saveEFLAGS(MBB, InsertPt, Loc);
1893
1894 Register NewReg = MRI->createVirtualRegister(RegClass: RC);
1895 unsigned OrOpCodes[] = {X86::OR8rr, X86::OR16rr, X86::OR32rr, X86::OR64rr};
1896 unsigned OrOpCode = OrOpCodes[Log2_32(Value: Bytes)];
1897 auto OrI = BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: OrOpCode), DestReg: NewReg)
1898 .addReg(RegNo: StateReg)
1899 .addReg(RegNo: Reg);
1900 OrI->addRegisterDead(Reg: X86::EFLAGS, RegInfo: TRI);
1901 ++NumInstsInserted;
1902 LLVM_DEBUG(dbgs() << " Inserting or: "; OrI->dump(); dbgs() << "\n");
1903
1904 if (FlagsReg)
1905 restoreEFLAGS(MBB, InsertPt, Loc, Reg: FlagsReg);
1906
1907 return NewReg;
1908}
1909
1910/// Harden a load by hardening the loaded value in the defined register.
1911///
1912/// We can harden a non-leaking load into a register without touching the
1913/// address by just hiding all of the loaded bits during misspeculation. We use
1914/// an `or` instruction to do this because we set up our poison value as all
1915/// ones. And the goal is just for the loaded bits to not be exposed to
1916/// execution and coercing them to one is sufficient.
1917///
1918/// Returns the newly hardened register.
1919Register X86SpeculativeLoadHardeningImpl::hardenPostLoad(MachineInstr &MI) {
1920 MachineBasicBlock &MBB = *MI.getParent();
1921 const DebugLoc &Loc = MI.getDebugLoc();
1922
1923 auto &DefOp = MI.getOperand(i: 0);
1924 Register OldDefReg = DefOp.getReg();
1925 auto *DefRC = MRI->getRegClass(Reg: OldDefReg);
1926
1927 // Because we want to completely replace the uses of this def'ed value with
1928 // the hardened value, create a dedicated new register that will only be used
1929 // to communicate the unhardened value to the hardening.
1930 Register UnhardenedReg = MRI->createVirtualRegister(RegClass: DefRC);
1931 DefOp.setReg(UnhardenedReg);
1932
1933 // Now harden this register's value, getting a hardened reg that is safe to
1934 // use. Note that we insert the instructions to compute this *after* the
1935 // defining instruction, not before it.
1936 Register HardenedReg = hardenValueInRegister(
1937 Reg: UnhardenedReg, MBB, InsertPt: std::next(x: MI.getIterator()), Loc);
1938
1939 // Finally, replace the old register (which now only has the uses of the
1940 // original def) with the hardened register.
1941 MRI->replaceRegWith(/*FromReg*/ OldDefReg, /*ToReg*/ HardenedReg);
1942
1943 ++NumPostLoadRegsHardened;
1944 return HardenedReg;
1945}
1946
1947/// Harden a return instruction.
1948///
1949/// Returns implicitly perform a load which we need to harden. Without hardening
1950/// this load, an attacker my speculatively write over the return address to
1951/// steer speculation of the return to an attacker controlled address. This is
1952/// called Spectre v1.1 or Bounds Check Bypass Store (BCBS) and is described in
1953/// this paper:
1954/// https://people.csail.mit.edu/vlk/spectre11.pdf
1955///
1956/// We can harden this by introducing an LFENCE that will delay any load of the
1957/// return address until prior instructions have retired (and thus are not being
1958/// speculated), or we can harden the address used by the implicit load: the
1959/// stack pointer.
1960///
1961/// If we are not using an LFENCE, hardening the stack pointer has an additional
1962/// benefit: it allows us to pass the predicate state accumulated in this
1963/// function back to the caller. In the absence of a BCBS attack on the return,
1964/// the caller will typically be resumed and speculatively executed due to the
1965/// Return Stack Buffer (RSB) prediction which is very accurate and has a high
1966/// priority. It is possible that some code from the caller will be executed
1967/// speculatively even during a BCBS-attacked return until the steering takes
1968/// effect. Whenever this happens, the caller can recover the (poisoned)
1969/// predicate state from the stack pointer and continue to harden loads.
1970void X86SpeculativeLoadHardeningImpl::hardenReturnInstr(MachineInstr &MI) {
1971 MachineBasicBlock &MBB = *MI.getParent();
1972 const DebugLoc &Loc = MI.getDebugLoc();
1973 auto InsertPt = MI.getIterator();
1974
1975 if (Subtarget->getCLOpts().slh_fence_call_and_ret)
1976 // No need to fence here as we'll fence at the return site itself. That
1977 // handles more cases than we can handle here.
1978 return;
1979
1980 // Take our predicate state, shift it to the high 17 bits (so that we keep
1981 // pointers canonical) and merge it into RSP. This will allow the caller to
1982 // extract it when we return (speculatively).
1983 mergePredStateIntoSP(MBB, InsertPt, Loc, PredStateReg: PS->SSA.GetValueAtEndOfBlock(BB: &MBB));
1984}
1985
1986/// Trace the predicate state through a call.
1987///
1988/// There are several layers of this needed to handle the full complexity of
1989/// calls.
1990///
1991/// First, we need to send the predicate state into the called function. We do
1992/// this by merging it into the high bits of the stack pointer.
1993///
1994/// For tail calls, this is all we need to do.
1995///
1996/// For calls where we might return and resume the control flow, we need to
1997/// extract the predicate state from the high bits of the stack pointer after
1998/// control returns from the called function.
1999///
2000/// We also need to verify that we intended to return to this location in the
2001/// code. An attacker might arrange for the processor to mispredict the return
2002/// to this valid but incorrect return address in the program rather than the
2003/// correct one. See the paper on this attack, called "ret2spec" by the
2004/// researchers, here:
2005/// https://christian-rossow.de/publications/ret2spec-ccs2018.pdf
2006///
2007/// The way we verify that we returned to the correct location is by preserving
2008/// the expected return address across the call. One technique involves taking
2009/// advantage of the red-zone to load the return address from `8(%rsp)` where it
2010/// was left by the RET instruction when it popped `%rsp`. Alternatively, we can
2011/// directly save the address into a register that will be preserved across the
2012/// call. We compare this intended return address against the address
2013/// immediately following the call (the observed return address). If these
2014/// mismatch, we have detected misspeculation and can poison our predicate
2015/// state.
2016void X86SpeculativeLoadHardeningImpl::tracePredStateThroughCall(
2017 MachineInstr &MI) {
2018 MachineBasicBlock &MBB = *MI.getParent();
2019 MachineFunction &MF = *MBB.getParent();
2020 auto InsertPt = MI.getIterator();
2021 const DebugLoc &Loc = MI.getDebugLoc();
2022
2023 if (Subtarget->getCLOpts().slh_fence_call_and_ret) {
2024 if (MI.isReturn())
2025 // Tail call, we don't return to this function.
2026 // FIXME: We should also handle noreturn calls.
2027 return;
2028
2029 // We don't need to fence before the call because the function should fence
2030 // in its entry. However, we do need to fence after the call returns.
2031 // Fencing before the return doesn't correctly handle cases where the return
2032 // itself is mispredicted.
2033 BuildMI(BB&: MBB, I: std::next(x: InsertPt), MIMD: Loc, MCID: TII->get(Opcode: X86::LFENCE));
2034 ++NumInstsInserted;
2035 ++NumLFENCEsInserted;
2036 return;
2037 }
2038
2039 // First, we transfer the predicate state into the called function by merging
2040 // it into the stack pointer. This will kill the current def of the state.
2041 Register StateReg = PS->SSA.GetValueAtEndOfBlock(BB: &MBB);
2042 mergePredStateIntoSP(MBB, InsertPt, Loc, PredStateReg: StateReg);
2043
2044 // If this call is also a return, it is a tail call and we don't need anything
2045 // else to handle it so just return. Also, if there are no further
2046 // instructions and no successors, this call does not return so we can also
2047 // bail.
2048 if (MI.isReturn() || (std::next(x: InsertPt) == MBB.end() && MBB.succ_empty()))
2049 return;
2050
2051 // Create a symbol to track the return address and attach it to the call
2052 // machine instruction. We will lower extra symbols attached to call
2053 // instructions as label immediately following the call.
2054 MCSymbol *RetSymbol =
2055 MF.getContext().createTempSymbol(Name: "slh_ret_addr",
2056 /*AlwaysAddSuffix*/ true);
2057 MI.setPostInstrSymbol(MF, Symbol: RetSymbol);
2058
2059 const TargetRegisterClass *AddrRC = &X86::GR64RegClass;
2060 Register ExpectedRetAddrReg;
2061
2062 // If we have no red zones or if the function returns twice (possibly without
2063 // using the `ret` instruction) like setjmp, we need to save the expected
2064 // return address prior to the call.
2065 if (!Subtarget->getFrameLowering()->has128ByteRedZone(MF) ||
2066 MF.exposesReturnsTwice()) {
2067 // If we don't have red zones, we need to compute the expected return
2068 // address prior to the call and store it in a register that lives across
2069 // the call.
2070 //
2071 // In some ways, this is doubly satisfying as a mitigation because it will
2072 // also successfully detect stack smashing bugs in some cases (typically,
2073 // when a callee-saved register is used and the callee doesn't push it onto
2074 // the stack). But that isn't our primary goal, so we only use it as
2075 // a fallback.
2076 //
2077 // FIXME: It isn't clear that this is reliable in the face of
2078 // rematerialization in the register allocator. We somehow need to force
2079 // that to not occur for this particular instruction, and instead to spill
2080 // or otherwise preserve the value computed *prior* to the call.
2081 //
2082 // FIXME: It is even less clear why MachineCSE can't just fold this when we
2083 // end up having to use identical instructions both before and after the
2084 // call to feed the comparison.
2085 ExpectedRetAddrReg = MRI->createVirtualRegister(RegClass: AddrRC);
2086 if (MF.getTarget().getCodeModel() == CodeModel::Small &&
2087 !Subtarget->isPositionIndependent()) {
2088 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::MOV64ri32), DestReg: ExpectedRetAddrReg)
2089 .addSym(Sym: RetSymbol);
2090 } else {
2091 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::LEA64r), DestReg: ExpectedRetAddrReg)
2092 .addReg(/*Base*/ RegNo: X86::RIP)
2093 .addImm(/*Scale*/ Val: 1)
2094 .addReg(/*Index*/ RegNo: 0)
2095 .addSym(Sym: RetSymbol)
2096 .addReg(/*Segment*/ RegNo: 0);
2097 }
2098 }
2099
2100 // Step past the call to handle when it returns.
2101 ++InsertPt;
2102
2103 // If we didn't pre-compute the expected return address into a register, then
2104 // red zones are enabled and the return address is still available on the
2105 // stack immediately after the call. As the very first instruction, we load it
2106 // into a register.
2107 if (!ExpectedRetAddrReg) {
2108 ExpectedRetAddrReg = MRI->createVirtualRegister(RegClass: AddrRC);
2109 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::MOV64rm), DestReg: ExpectedRetAddrReg)
2110 .addReg(/*Base*/ RegNo: X86::RSP)
2111 .addImm(/*Scale*/ Val: 1)
2112 .addReg(/*Index*/ RegNo: 0)
2113 .addImm(/*Displacement*/ Val: -8) // The stack pointer has been popped, so
2114 // the return address is 8-bytes past it.
2115 .addReg(/*Segment*/ RegNo: 0);
2116 }
2117
2118 // Now we extract the callee's predicate state from the stack pointer.
2119 Register NewStateReg = extractPredStateFromSP(MBB, InsertPt, Loc);
2120
2121 // Test the expected return address against our actual address. If we can
2122 // form this basic block's address as an immediate, this is easy. Otherwise
2123 // we compute it.
2124 if (MF.getTarget().getCodeModel() == CodeModel::Small &&
2125 !Subtarget->isPositionIndependent()) {
2126 // FIXME: Could we fold this with the load? It would require careful EFLAGS
2127 // management.
2128 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::CMP64ri32))
2129 .addReg(RegNo: ExpectedRetAddrReg, Flags: RegState::Kill)
2130 .addSym(Sym: RetSymbol);
2131 } else {
2132 Register ActualRetAddrReg = MRI->createVirtualRegister(RegClass: AddrRC);
2133 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::LEA64r), DestReg: ActualRetAddrReg)
2134 .addReg(/*Base*/ RegNo: X86::RIP)
2135 .addImm(/*Scale*/ Val: 1)
2136 .addReg(/*Index*/ RegNo: 0)
2137 .addSym(Sym: RetSymbol)
2138 .addReg(/*Segment*/ RegNo: 0);
2139 BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: X86::CMP64rr))
2140 .addReg(RegNo: ExpectedRetAddrReg, Flags: RegState::Kill)
2141 .addReg(RegNo: ActualRetAddrReg, Flags: RegState::Kill);
2142 }
2143
2144 // Now conditionally update the predicate state we just extracted if we ended
2145 // up at a different return address than expected.
2146 int PredStateSizeInBytes = TRI->getRegSizeInBits(RC: *PS->RC) / 8;
2147 auto CMovOp = X86::getCMovOpcode(RegBytes: PredStateSizeInBytes);
2148
2149 Register UpdatedStateReg = MRI->createVirtualRegister(RegClass: PS->RC);
2150 auto CMovI = BuildMI(BB&: MBB, I: InsertPt, MIMD: Loc, MCID: TII->get(Opcode: CMovOp), DestReg: UpdatedStateReg)
2151 .addReg(RegNo: NewStateReg, Flags: RegState::Kill)
2152 .addReg(RegNo: PS->PoisonReg)
2153 .addImm(Val: X86::COND_NE);
2154 CMovI->findRegisterUseOperand(Reg: X86::EFLAGS, /*TRI=*/nullptr)->setIsKill(true);
2155 ++NumInstsInserted;
2156 LLVM_DEBUG(dbgs() << " Inserting cmov: "; CMovI->dump(); dbgs() << "\n");
2157
2158 PS->SSA.AddAvailableValue(BB: &MBB, V: UpdatedStateReg);
2159}
2160
2161/// An attacker may speculatively store over a value that is then speculatively
2162/// loaded and used as the target of an indirect call or jump instruction. This
2163/// is called Spectre v1.2 or Bounds Check Bypass Store (BCBS) and is described
2164/// in this paper:
2165/// https://people.csail.mit.edu/vlk/spectre11.pdf
2166///
2167/// When this happens, the speculative execution of the call or jump will end up
2168/// being steered to this attacker controlled address. While most such loads
2169/// will be adequately hardened already, we want to ensure that they are
2170/// definitively treated as needing post-load hardening. While address hardening
2171/// is sufficient to prevent secret data from leaking to the attacker, it may
2172/// not be sufficient to prevent an attacker from steering speculative
2173/// execution. We forcibly unfolded all relevant loads above and so will always
2174/// have an opportunity to post-load harden here, we just need to scan for cases
2175/// not already flagged and add them.
2176void X86SpeculativeLoadHardeningImpl::hardenIndirectCallOrJumpInstr(
2177 MachineInstr &MI,
2178 SmallDenseMap<Register, Register, 32> &AddrRegToHardenedReg) {
2179 switch (MI.getOpcode()) {
2180 case X86::FARCALL16m:
2181 case X86::FARCALL32m:
2182 case X86::FARCALL64m:
2183 case X86::FARJMP16m:
2184 case X86::FARJMP32m:
2185 case X86::FARJMP64m:
2186 // We don't need to harden either far calls or far jumps as they are
2187 // safe from Spectre.
2188 return;
2189
2190 default:
2191 break;
2192 }
2193
2194 // We should never see a loading instruction at this point, as those should
2195 // have been unfolded.
2196 assert(!MI.mayLoad() && "Found a lingering loading instruction!");
2197
2198 // If the first operand isn't a register, this is a branch or call
2199 // instruction with an immediate operand which doesn't need to be hardened.
2200 if (!MI.getOperand(i: 0).isReg())
2201 return;
2202
2203 // For all of these, the target register is the first operand of the
2204 // instruction.
2205 auto &TargetOp = MI.getOperand(i: 0);
2206 Register OldTargetReg = TargetOp.getReg();
2207
2208 // Try to lookup a hardened version of this register. We retain a reference
2209 // here as we want to update the map to track any newly computed hardened
2210 // register.
2211 Register &HardenedTargetReg = AddrRegToHardenedReg[OldTargetReg];
2212
2213 // If we don't have a hardened register yet, compute one. Otherwise, just use
2214 // the already hardened register.
2215 //
2216 // FIXME: It is a little suspect that we use partially hardened registers that
2217 // only feed addresses. The complexity of partial hardening with SHRX
2218 // continues to pile up. Should definitively measure its value and consider
2219 // eliminating it.
2220 if (!HardenedTargetReg)
2221 HardenedTargetReg = hardenValueInRegister(
2222 Reg: OldTargetReg, MBB&: *MI.getParent(), InsertPt: MI.getIterator(), Loc: MI.getDebugLoc());
2223
2224 // Set the target operand to the hardened register.
2225 TargetOp.setReg(HardenedTargetReg);
2226
2227 ++NumCallsOrJumpsHardened;
2228}
2229
2230PreservedAnalyses
2231X86SpeculativeLoadHardeningPass::run(MachineFunction &MF,
2232 MachineFunctionAnalysisManager &MFAM) {
2233 X86SpeculativeLoadHardeningImpl Impl;
2234 const bool Changed = Impl.run(MF);
2235 LLVM_DEBUG(dbgs() << "Final speculative load hardened function:\n"; MF.dump();
2236 dbgs() << "\n"; MF.verify(MFAM));
2237 return Changed ? getMachineFunctionPassPreservedAnalyses()
2238 .preserveSet<CFGAnalyses>()
2239 : PreservedAnalyses::all();
2240}
2241
2242INITIALIZE_PASS_BEGIN(X86SpeculativeLoadHardeningLegacy, PASS_KEY,
2243 "X86 speculative load hardener", false, false)
2244INITIALIZE_PASS_END(X86SpeculativeLoadHardeningLegacy, PASS_KEY,
2245 "X86 speculative load hardener", false, false)
2246
2247FunctionPass *llvm::createX86SpeculativeLoadHardeningLegacyPass() {
2248 return new X86SpeculativeLoadHardeningLegacy();
2249}
2250