1//===- AArch64SpeculationHardening.cpp - Harden Against Missspeculation --===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file contains a pass to insert code to mitigate against side channel
10// vulnerabilities that may happen under control flow miss-speculation.
11//
12// The pass implements tracking of control flow miss-speculation into a "taint"
13// register. That taint register can then be used to mask off registers with
14// sensitive data when executing under miss-speculation, a.k.a. "transient
15// execution".
16// This pass is aimed at mitigating against SpectreV1-style vulnerabilities.
17//
18// It also implements speculative load hardening, i.e. using the taint register
19// to automatically mask off loaded data.
20//
21// As a possible follow-on improvement, also an intrinsics-based approach as
22// explained at https://lwn.net/Articles/759423/ could be implemented on top of
23// the current design.
24//
25// For AArch64, the following implementation choices are made to implement the
26// tracking of control flow miss-speculation into a taint register:
27// Some of these are different than the implementation choices made in
28// the similar pass implemented in X86SpeculativeLoadHardening.cpp, as
29// the instruction set characteristics result in different trade-offs.
30// - The speculation hardening is done after register allocation. With a
31// relative abundance of registers, one register is reserved (X16) to be
32// the taint register. X16 is expected to not clash with other register
33// reservation mechanisms with very high probability because:
34// . The AArch64 ABI doesn't guarantee X16 to be retained across any call.
35// . The only way to request X16 to be used as a programmer is through
36// inline assembly. In the rare case a function explicitly demands to
37// use X16/W16, this pass falls back to hardening against speculation
38// by inserting a DSB SYS/ISB barrier pair which will prevent control
39// flow speculation.
40// - It is easy to insert mask operations at this late stage as we have
41// mask operations available that don't set flags.
42// - The taint variable contains all-ones when no miss-speculation is detected,
43// and contains all-zeros when miss-speculation is detected. Therefore, when
44// masking, an AND instruction (which only changes the register to be masked,
45// no other side effects) can easily be inserted anywhere that's needed.
46// - The tracking of miss-speculation is done by using a data-flow conditional
47// select instruction (CSEL) to evaluate the flags that were also used to
48// make conditional branch direction decisions. Speculation of the CSEL
49// instruction can be limited with a CSDB instruction - so the combination of
50// CSEL + a later CSDB gives the guarantee that the flags as used in the CSEL
51// aren't speculated. When conditional branch direction gets miss-speculated,
52// the semantics of the inserted CSEL instruction is such that the taint
53// register will contain all zero bits.
54// One key requirement for this to work is that the conditional branch is
55// followed by an execution of the CSEL instruction, where the CSEL
56// instruction needs to use the same flags status as the conditional branch.
57// This means that the conditional branches must not be implemented as one
58// of the AArch64 conditional branches that do not use the flags as input
59// (CB(N)Z and TB(N)Z). This is implemented by ensuring in the instruction
60// selectors to not produce these instructions when speculation hardening
61// is enabled. This pass will assert if it does encounter such an instruction.
62// - On function call boundaries, the miss-speculation state is transferred from
63// the taint register X16 to be encoded in the SP register as value 0.
64//
65// For the aspect of automatically hardening loads, using the taint register,
66// (a.k.a. speculative load hardening, see
67// https://llvm.org/docs/SpeculativeLoadHardening.html), the following
68// implementation choices are made for AArch64:
69// - Many of the optimizations described at
70// https://llvm.org/docs/SpeculativeLoadHardening.html to harden fewer
71// loads haven't been implemented yet - but for some of them there are
72// FIXMEs in the code.
73// - loads that load into general purpose (X or W) registers get hardened by
74// masking the loaded data. For loads that load into other registers, the
75// address loaded from gets hardened. It is expected that hardening the
76// loaded data may be more efficient; but masking data in registers other
77// than X or W is not easy and may result in being slower than just
78// hardening the X address register loaded from.
79// - On AArch64, CSDB instructions are inserted between the masking of the
80// register and its first use, to ensure there's no non-control-flow
81// speculation that might undermine the hardening mechanism.
82//
83// Future extensions/improvements could be:
84// - Implement this functionality using full speculation barriers, akin to the
85// x86-slh-lfence option. This may be more useful for the intrinsics-based
86// approach than for the SLH approach to masking.
87// Note that this pass already inserts the full speculation barriers if the
88// function for some niche reason makes use of X16/W16.
89// - no indirect branch misprediction gets protected/instrumented; but this
90// could be done for some indirect branches, such as switch jump tables.
91//===----------------------------------------------------------------------===//
92
93#include "AArch64Subtarget.h"
94#include "Utils/AArch64BaseInfo.h"
95#include "llvm/ADT/BitVector.h"
96#include "llvm/ADT/SmallVector.h"
97#include "llvm/CodeGen/MachineBasicBlock.h"
98#include "llvm/CodeGen/MachineFunction.h"
99#include "llvm/CodeGen/MachineFunctionPass.h"
100#include "llvm/CodeGen/MachineInstr.h"
101#include "llvm/CodeGen/MachineInstrBuilder.h"
102#include "llvm/CodeGen/MachineOperand.h"
103#include "llvm/CodeGen/RegisterScavenging.h"
104#include "llvm/IR/DebugLoc.h"
105#include "llvm/Pass.h"
106#include "llvm/Support/Debug.h"
107#include "llvm/Target/TargetMachine.h"
108#include <cassert>
109
110using namespace llvm;
111
112#define DEBUG_TYPE "aarch64-speculation-hardening"
113
114#define AARCH64_SPECULATION_HARDENING_NAME "AArch64 speculation hardening pass"
115
116namespace {
117
118class AArch64SpeculationHardening : public MachineFunctionPass {
119public:
120 const TargetInstrInfo *TII;
121 const TargetRegisterInfo *TRI;
122
123 static char ID;
124
125 AArch64SpeculationHardening() : MachineFunctionPass(ID) {}
126
127 bool runOnMachineFunction(MachineFunction &Fn) override;
128
129 StringRef getPassName() const override {
130 return AARCH64_SPECULATION_HARDENING_NAME;
131 }
132
133private:
134 unsigned MisspeculatingTaintReg;
135 unsigned MisspeculatingTaintReg32Bit;
136 bool UseControlFlowSpeculationBarrier;
137 BitVector RegsNeedingCSDBBeforeUse;
138 BitVector RegsAlreadyMasked;
139
140 bool functionUsesHardeningRegister(MachineFunction &MF) const;
141 bool instrumentControlFlow(MachineBasicBlock &MBB,
142 bool &UsesFullSpeculationBarrier);
143 bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
144 MachineBasicBlock *&FBB,
145 AArch64CC::CondCode &CondCode) const;
146 void insertTrackingCode(MachineBasicBlock &SplitEdgeBB,
147 AArch64CC::CondCode &CondCode, DebugLoc DL) const;
148 void insertSPToRegTaintPropagation(MachineBasicBlock &MBB,
149 MachineBasicBlock::iterator MBBI) const;
150 void insertRegToSPTaintPropagation(MachineBasicBlock &MBB,
151 MachineBasicBlock::iterator MBBI,
152 unsigned TmpReg) const;
153 void insertFullSpeculationBarrier(MachineBasicBlock &MBB,
154 MachineBasicBlock::iterator MBBI,
155 DebugLoc DL) const;
156
157 bool slhLoads(MachineBasicBlock &MBB);
158 bool makeGPRSpeculationSafe(MachineBasicBlock &MBB,
159 MachineBasicBlock::iterator MBBI,
160 MachineInstr &MI, unsigned Reg);
161 bool lowerSpeculationSafeValuePseudos(MachineBasicBlock &MBB,
162 bool UsesFullSpeculationBarrier);
163 bool expandSpeculationSafeValue(MachineBasicBlock &MBB,
164 MachineBasicBlock::iterator MBBI,
165 bool UsesFullSpeculationBarrier);
166 bool insertCSDB(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
167 DebugLoc DL);
168};
169
170} // end anonymous namespace
171
172char AArch64SpeculationHardening::ID = 0;
173
174INITIALIZE_PASS(AArch64SpeculationHardening, "aarch64-speculation-hardening",
175 AARCH64_SPECULATION_HARDENING_NAME, false, false)
176
177bool AArch64SpeculationHardening::endsWithCondControlFlow(
178 MachineBasicBlock &MBB, MachineBasicBlock *&TBB, MachineBasicBlock *&FBB,
179 AArch64CC::CondCode &CondCode) const {
180 SmallVector<MachineOperand, 1> analyzeBranchCondCode;
181 if (TII->analyzeBranch(MBB, TBB, FBB, Cond&: analyzeBranchCondCode, AllowModify: false))
182 return false;
183
184 // Ignore if the BB ends in an unconditional branch/fall-through.
185 if (analyzeBranchCondCode.empty())
186 return false;
187
188 // If the BB ends with a single conditional branch, FBB will be set to
189 // nullptr (see API docs for TII->analyzeBranch). For the rest of the
190 // analysis we want the FBB block to be set always.
191 assert(TBB != nullptr);
192 if (FBB == nullptr)
193 FBB = MBB.getFallThrough();
194
195 // If both the true and the false condition jump to the same basic block,
196 // there isn't need for any protection - whether the branch is speculated
197 // correctly or not, we end up executing the architecturally correct code.
198 if (TBB == FBB)
199 return false;
200
201 assert(MBB.succ_size() == 2);
202 // translate analyzeBranchCondCode to CondCode.
203 assert(analyzeBranchCondCode.size() == 1 && "unknown Cond array format");
204 CondCode = AArch64CC::CondCode(analyzeBranchCondCode[0].getImm());
205 return true;
206}
207
208void AArch64SpeculationHardening::insertFullSpeculationBarrier(
209 MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
210 DebugLoc DL) const {
211 // A full control flow speculation barrier consists of (DSB SYS + ISB)
212 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::DSB)).addImm(Val: 0xf);
213 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::ISB)).addImm(Val: 0xf);
214}
215
216void AArch64SpeculationHardening::insertTrackingCode(
217 MachineBasicBlock &SplitEdgeBB, AArch64CC::CondCode &CondCode,
218 DebugLoc DL) const {
219 if (UseControlFlowSpeculationBarrier) {
220 insertFullSpeculationBarrier(MBB&: SplitEdgeBB, MBBI: SplitEdgeBB.begin(), DL);
221 } else {
222 BuildMI(BB&: SplitEdgeBB, I: SplitEdgeBB.begin(), MIMD: DL, MCID: TII->get(Opcode: AArch64::CSELXr))
223 .addDef(RegNo: MisspeculatingTaintReg)
224 .addUse(RegNo: MisspeculatingTaintReg)
225 .addUse(RegNo: AArch64::XZR)
226 .addImm(Val: CondCode);
227 SplitEdgeBB.addLiveIn(PhysReg: AArch64::NZCV);
228 }
229}
230
231bool AArch64SpeculationHardening::instrumentControlFlow(
232 MachineBasicBlock &MBB, bool &UsesFullSpeculationBarrier) {
233 LLVM_DEBUG(dbgs() << "Instrument control flow tracking on MBB: " << MBB);
234
235 bool Modified = false;
236 MachineBasicBlock *TBB = nullptr;
237 MachineBasicBlock *FBB = nullptr;
238 AArch64CC::CondCode CondCode;
239
240 if (!endsWithCondControlFlow(MBB, TBB, FBB, CondCode)) {
241 LLVM_DEBUG(dbgs() << "... doesn't end with CondControlFlow\n");
242 } else {
243 // Now insert:
244 // "CSEL MisSpeculatingR, MisSpeculatingR, XZR, cond" on the True edge and
245 // "CSEL MisSpeculatingR, MisSpeculatingR, XZR, Invertcond" on the False
246 // edge.
247 AArch64CC::CondCode InvCondCode = AArch64CC::getInvertedCondCode(Code: CondCode);
248
249 MachineBasicBlock *SplitEdgeTBB = MBB.SplitCriticalEdge(Succ: TBB, P&: *this);
250 MachineBasicBlock *SplitEdgeFBB = MBB.SplitCriticalEdge(Succ: FBB, P&: *this);
251
252 assert(SplitEdgeTBB != nullptr);
253 assert(SplitEdgeFBB != nullptr);
254
255 DebugLoc DL;
256 if (MBB.instr_end() != MBB.instr_begin())
257 DL = (--MBB.instr_end())->getDebugLoc();
258
259 insertTrackingCode(SplitEdgeBB&: *SplitEdgeTBB, CondCode, DL);
260 insertTrackingCode(SplitEdgeBB&: *SplitEdgeFBB, CondCode&: InvCondCode, DL);
261
262 LLVM_DEBUG(dbgs() << "SplitEdgeTBB: " << *SplitEdgeTBB << "\n");
263 LLVM_DEBUG(dbgs() << "SplitEdgeFBB: " << *SplitEdgeFBB << "\n");
264 Modified = true;
265 }
266
267 // Perform correct code generation around function calls and before returns.
268 // The below variables record the return/terminator instructions and the call
269 // instructions respectively; including which register is available as a
270 // temporary register just before the recorded instructions.
271 SmallVector<std::pair<MachineInstr *, unsigned>, 4> ReturnInstructions;
272 SmallVector<std::pair<MachineInstr *, unsigned>, 4> CallInstructions;
273 // if a temporary register is not available for at least one of the
274 // instructions for which we need to transfer taint to the stack pointer, we
275 // need to insert a full speculation barrier.
276 // TmpRegisterNotAvailableEverywhere tracks that condition.
277 bool TmpRegisterNotAvailableEverywhere = false;
278
279 RegScavenger RS;
280 RS.enterBasicBlockEnd(MBB);
281
282 for (MachineBasicBlock::iterator I = MBB.end(); I != MBB.begin(); ) {
283 MachineInstr &MI = *--I;
284 if (!MI.isReturn() && !MI.isCall())
285 continue;
286
287 // The RegScavenger represents registers available *after* the MI
288 // instruction pointed to by RS.getCurrentPosition().
289 // We need to have a register that is available *before* the MI is executed.
290 if (I == MBB.begin())
291 RS.enterBasicBlock(MBB);
292 else
293 RS.backward(I);
294 // FIXME: The below just finds *a* unused register. Maybe code could be
295 // optimized more if this looks for the register that isn't used for the
296 // longest time around this place, to enable more scheduling freedom. Not
297 // sure if that would actually result in a big performance difference
298 // though. Maybe RegisterScavenger::findSurvivorBackwards has some logic
299 // already to do this - but it's unclear if that could easily be used here.
300 Register TmpReg = RS.FindUnusedReg(RC: &AArch64::GPR64commonRegClass);
301 LLVM_DEBUG(dbgs() << "RS finds "
302 << ((TmpReg == 0) ? "no register " : "register ");
303 if (TmpReg != 0) dbgs() << printReg(TmpReg, TRI) << " ";
304 dbgs() << "to be available at MI " << MI);
305 if (TmpReg == 0)
306 TmpRegisterNotAvailableEverywhere = true;
307 if (MI.isReturn())
308 ReturnInstructions.push_back(Elt: {&MI, TmpReg});
309 else if (MI.isCall())
310 CallInstructions.push_back(Elt: {&MI, TmpReg});
311 }
312
313 if (TmpRegisterNotAvailableEverywhere) {
314 // When a temporary register is not available everywhere in this basic
315 // basic block where a propagate-taint-to-sp operation is needed, just
316 // emit a full speculation barrier at the start of this basic block, which
317 // renders the taint/speculation tracking in this basic block unnecessary.
318 insertFullSpeculationBarrier(MBB, MBBI: MBB.begin(),
319 DL: (MBB.begin())->getDebugLoc());
320 UsesFullSpeculationBarrier = true;
321 Modified = true;
322 } else {
323 for (auto MI_Reg : ReturnInstructions) {
324 assert(MI_Reg.second != 0);
325 LLVM_DEBUG(
326 dbgs()
327 << " About to insert Reg to SP taint propagation with temp register "
328 << printReg(MI_Reg.second, TRI)
329 << " on instruction: " << *MI_Reg.first);
330 insertRegToSPTaintPropagation(MBB, MBBI: MI_Reg.first, TmpReg: MI_Reg.second);
331 Modified = true;
332 }
333
334 for (auto MI_Reg : CallInstructions) {
335 assert(MI_Reg.second != 0);
336 LLVM_DEBUG(dbgs() << " About to insert Reg to SP and back taint "
337 "propagation with temp register "
338 << printReg(MI_Reg.second, TRI)
339 << " around instruction: " << *MI_Reg.first);
340 // Just after the call:
341 insertSPToRegTaintPropagation(
342 MBB, MBBI: std::next(x: (MachineBasicBlock::iterator)MI_Reg.first));
343 // Just before the call:
344 insertRegToSPTaintPropagation(MBB, MBBI: MI_Reg.first, TmpReg: MI_Reg.second);
345 Modified = true;
346 }
347 }
348 return Modified;
349}
350
351void AArch64SpeculationHardening::insertSPToRegTaintPropagation(
352 MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) const {
353 // If full control flow speculation barriers are used, emit a control flow
354 // barrier to block potential miss-speculation in flight coming in to this
355 // function.
356 if (UseControlFlowSpeculationBarrier) {
357 insertFullSpeculationBarrier(MBB, MBBI, DL: DebugLoc());
358 return;
359 }
360
361 // CMP SP, #0 === SUBS xzr, SP, #0
362 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: AArch64::SUBSXri))
363 .addDef(RegNo: AArch64::XZR)
364 .addUse(RegNo: AArch64::SP)
365 .addImm(Val: 0)
366 .addImm(Val: 0); // no shift
367 // CSETM x16, NE === CSINV x16, xzr, xzr, EQ
368 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: AArch64::CSINVXr))
369 .addDef(RegNo: MisspeculatingTaintReg)
370 .addUse(RegNo: AArch64::XZR)
371 .addUse(RegNo: AArch64::XZR)
372 .addImm(Val: AArch64CC::EQ);
373}
374
375void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
376 MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
377 unsigned TmpReg) const {
378 // If full control flow speculation barriers are used, there will not be
379 // miss-speculation when returning from this function, and therefore, also
380 // no need to encode potential miss-speculation into the stack pointer.
381 if (UseControlFlowSpeculationBarrier)
382 return;
383
384 // mov Xtmp, SP === ADD Xtmp, SP, #0
385 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: AArch64::ADDXri))
386 .addDef(RegNo: TmpReg)
387 .addUse(RegNo: AArch64::SP)
388 .addImm(Val: 0)
389 .addImm(Val: 0); // no shift
390 // and Xtmp, Xtmp, TaintReg === AND Xtmp, Xtmp, TaintReg, #0
391 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: AArch64::ANDXrs))
392 .addDef(RegNo: TmpReg, Flags: RegState::Renamable)
393 .addUse(RegNo: TmpReg, Flags: RegState::Kill | RegState::Renamable)
394 .addUse(RegNo: MisspeculatingTaintReg, Flags: RegState::Kill)
395 .addImm(Val: 0);
396 // mov SP, Xtmp === ADD SP, Xtmp, #0
397 BuildMI(BB&: MBB, I: MBBI, MIMD: DebugLoc(), MCID: TII->get(Opcode: AArch64::ADDXri))
398 .addDef(RegNo: AArch64::SP)
399 .addUse(RegNo: TmpReg, Flags: RegState::Kill)
400 .addImm(Val: 0)
401 .addImm(Val: 0); // no shift
402}
403
404bool AArch64SpeculationHardening::functionUsesHardeningRegister(
405 MachineFunction &MF) const {
406 for (MachineBasicBlock &MBB : MF) {
407 for (MachineInstr &MI : MBB) {
408 // treat function calls specially, as the hardening register does not
409 // need to remain live across function calls.
410 if (MI.isCall())
411 continue;
412 if (MI.readsRegister(Reg: MisspeculatingTaintReg, TRI) ||
413 MI.modifiesRegister(Reg: MisspeculatingTaintReg, TRI))
414 return true;
415 }
416 }
417 return false;
418}
419
420// Make GPR register Reg speculation-safe by putting it through the
421// SpeculationSafeValue pseudo instruction, if we can't prove that
422// the value in the register has already been hardened.
423bool AArch64SpeculationHardening::makeGPRSpeculationSafe(
424 MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, MachineInstr &MI,
425 unsigned Reg) {
426 assert(AArch64::GPR32allRegClass.contains(Reg) ||
427 AArch64::GPR64allRegClass.contains(Reg));
428
429 // Loads cannot directly load a value into the SP (nor WSP).
430 // Therefore, if Reg is SP or WSP, it is because the instruction loads from
431 // the stack through the stack pointer.
432 //
433 // Since the stack pointer is never dynamically controllable, don't harden it.
434 if (Reg == AArch64::SP || Reg == AArch64::WSP)
435 return false;
436
437 // Do not harden the register again if already hardened before.
438 if (RegsAlreadyMasked[Reg])
439 return false;
440
441 const bool Is64Bit = AArch64::GPR64allRegClass.contains(Reg);
442 LLVM_DEBUG(dbgs() << "About to harden register : " << Reg << "\n");
443 BuildMI(BB&: MBB, I: MBBI, MIMD: MI.getDebugLoc(),
444 MCID: TII->get(Opcode: Is64Bit ? AArch64::SpeculationSafeValueX
445 : AArch64::SpeculationSafeValueW))
446 .addDef(RegNo: Reg)
447 .addUse(RegNo: Reg);
448 RegsAlreadyMasked.set(Reg);
449 return true;
450}
451
452bool AArch64SpeculationHardening::slhLoads(MachineBasicBlock &MBB) {
453 bool Modified = false;
454
455 LLVM_DEBUG(dbgs() << "slhLoads running on MBB: " << MBB);
456
457 RegsAlreadyMasked.reset();
458
459 MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
460 MachineBasicBlock::iterator NextMBBI;
461 for (; MBBI != E; MBBI = NextMBBI) {
462 MachineInstr &MI = *MBBI;
463 NextMBBI = std::next(x: MBBI);
464 // Only harden loaded values or addresses used in loads.
465 if (!MI.mayLoad())
466 continue;
467
468 LLVM_DEBUG(dbgs() << "About to harden: " << MI);
469
470 // For general purpose register loads, harden the registers loaded into.
471 // For other loads, harden the address loaded from.
472 // Masking the loaded value is expected to result in less performance
473 // overhead, as the load can still execute speculatively in comparison to
474 // when the address loaded from gets masked. However, masking is only
475 // easy to do efficiently on GPR registers, so for loads into non-GPR
476 // registers (e.g. floating point loads), mask the address loaded from.
477 bool AllDefsAreGPR = llvm::all_of(Range: MI.defs(), P: [&](MachineOperand &Op) {
478 return Op.isReg() && (AArch64::GPR32allRegClass.contains(Reg: Op.getReg()) ||
479 AArch64::GPR64allRegClass.contains(Reg: Op.getReg()));
480 });
481 // FIXME: it might be a worthwhile optimization to not mask loaded
482 // values if all the registers involved in address calculation are already
483 // hardened, leading to this load not able to execute on a miss-speculated
484 // path.
485 bool HardenLoadedData = AllDefsAreGPR;
486 bool HardenAddressLoadedFrom = !HardenLoadedData;
487
488 // First remove registers from AlreadyMaskedRegisters if their value is
489 // updated by this instruction - it makes them contain a new value that is
490 // not guaranteed to already have been masked.
491 for (MachineOperand Op : MI.defs())
492 for (MCRegAliasIterator AI(Op.getReg(), TRI, true); AI.isValid(); ++AI)
493 RegsAlreadyMasked.reset(Idx: *AI);
494
495 // FIXME: loads from the stack with an immediate offset from the stack
496 // pointer probably shouldn't be hardened, which could result in a
497 // significant optimization. See section "Don’t check loads from
498 // compile-time constant stack offsets", in
499 // https://llvm.org/docs/SpeculativeLoadHardening.html
500
501 if (HardenLoadedData)
502 for (auto Def : MI.defs()) {
503 if (Def.isDead())
504 // Do not mask a register that is not used further.
505 continue;
506 // FIXME: For pre/post-increment addressing modes, the base register
507 // used in address calculation is also defined by this instruction.
508 // It might be a worthwhile optimization to not harden that
509 // base register increment/decrement when the increment/decrement is
510 // an immediate.
511 Modified |= makeGPRSpeculationSafe(MBB, MBBI: NextMBBI, MI, Reg: Def.getReg());
512 }
513
514 if (HardenAddressLoadedFrom)
515 for (auto Use : MI.uses()) {
516 if (!Use.isReg())
517 continue;
518 Register Reg = Use.getReg();
519 // Some loads of floating point data have implicit defs/uses on a
520 // super register of that floating point data. Some examples:
521 // $s0 = LDRSui $sp, 22, implicit-def $q0
522 // $q0 = LD1i64 $q0, 1, renamable $x0
523 // We need to filter out these uses for non-GPR register which occur
524 // because the load partially fills a non-GPR register with the loaded
525 // data. Just skipping all non-GPR registers is safe (for now) as all
526 // AArch64 load instructions only use GPR registers to perform the
527 // address calculation. FIXME: However that might change once we can
528 // produce SVE gather instructions.
529 if (!(AArch64::GPR32allRegClass.contains(Reg) ||
530 AArch64::GPR64allRegClass.contains(Reg)))
531 continue;
532 Modified |= makeGPRSpeculationSafe(MBB, MBBI, MI, Reg);
533 }
534 }
535 return Modified;
536}
537
538/// \brief If MBBI references a pseudo instruction that should be expanded
539/// here, do the expansion and return true. Otherwise return false.
540bool AArch64SpeculationHardening::expandSpeculationSafeValue(
541 MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
542 bool UsesFullSpeculationBarrier) {
543 MachineInstr &MI = *MBBI;
544 unsigned Opcode = MI.getOpcode();
545 bool Is64Bit = true;
546
547 switch (Opcode) {
548 default:
549 break;
550 case AArch64::SpeculationSafeValueW:
551 Is64Bit = false;
552 [[fallthrough]];
553 case AArch64::SpeculationSafeValueX:
554 // Just remove the SpeculationSafe pseudo's if control flow
555 // miss-speculation isn't happening because we're already inserting barriers
556 // to guarantee that.
557 if (!UseControlFlowSpeculationBarrier && !UsesFullSpeculationBarrier) {
558 Register DstReg = MI.getOperand(i: 0).getReg();
559 Register SrcReg = MI.getOperand(i: 1).getReg();
560 // Mark this register and all its aliasing registers as needing to be
561 // value speculation hardened before its next use, by using a CSDB
562 // barrier instruction.
563 for (MachineOperand Op : MI.defs())
564 for (MCRegAliasIterator AI(Op.getReg(), TRI, true); AI.isValid(); ++AI)
565 RegsNeedingCSDBBeforeUse.set(*AI);
566
567 // Mask off with taint state.
568 BuildMI(BB&: MBB, I: MBBI, MIMD: MI.getDebugLoc(),
569 MCID: Is64Bit ? TII->get(Opcode: AArch64::ANDXrs) : TII->get(Opcode: AArch64::ANDWrs))
570 .addDef(RegNo: DstReg)
571 .addUse(RegNo: SrcReg, Flags: RegState::Kill)
572 .addUse(RegNo: Is64Bit ? MisspeculatingTaintReg
573 : MisspeculatingTaintReg32Bit)
574 .addImm(Val: 0);
575 }
576 MI.eraseFromParent();
577 return true;
578 }
579 return false;
580}
581
582bool AArch64SpeculationHardening::insertCSDB(MachineBasicBlock &MBB,
583 MachineBasicBlock::iterator MBBI,
584 DebugLoc DL) {
585 assert(!UseControlFlowSpeculationBarrier && "No need to insert CSDBs when "
586 "control flow miss-speculation "
587 "is already blocked");
588 // insert data value speculation barrier (CSDB)
589 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::HINT)).addImm(Val: 0x14);
590 RegsNeedingCSDBBeforeUse.reset();
591 return true;
592}
593
594bool AArch64SpeculationHardening::lowerSpeculationSafeValuePseudos(
595 MachineBasicBlock &MBB, bool UsesFullSpeculationBarrier) {
596 bool Modified = false;
597
598 RegsNeedingCSDBBeforeUse.reset();
599
600 // The following loop iterates over all instructions in the basic block,
601 // and performs 2 operations:
602 // 1. Insert a CSDB at this location if needed.
603 // 2. Expand the SpeculationSafeValuePseudo if the current instruction is
604 // one.
605 //
606 // The insertion of the CSDB is done as late as possible (i.e. just before
607 // the use of a masked register), in the hope that that will reduce the
608 // total number of CSDBs in a block when there are multiple masked registers
609 // in the block.
610 MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
611 DebugLoc DL;
612 while (MBBI != E) {
613 MachineInstr &MI = *MBBI;
614 DL = MI.getDebugLoc();
615 MachineBasicBlock::iterator NMBBI = std::next(x: MBBI);
616
617 // First check if a CSDB needs to be inserted due to earlier registers
618 // that were masked and that are used by the next instruction.
619 // Also emit the barrier on any potential control flow changes.
620 bool NeedToEmitBarrier = false;
621 if (RegsNeedingCSDBBeforeUse.any() && (MI.isCall() || MI.isTerminator()))
622 NeedToEmitBarrier = true;
623 if (!NeedToEmitBarrier)
624 for (MachineOperand Op : MI.uses())
625 if (Op.isReg() && RegsNeedingCSDBBeforeUse[Op.getReg()]) {
626 NeedToEmitBarrier = true;
627 break;
628 }
629
630 if (NeedToEmitBarrier && !UsesFullSpeculationBarrier)
631 Modified |= insertCSDB(MBB, MBBI, DL);
632
633 Modified |=
634 expandSpeculationSafeValue(MBB, MBBI, UsesFullSpeculationBarrier);
635
636 MBBI = NMBBI;
637 }
638
639 if (RegsNeedingCSDBBeforeUse.any() && !UsesFullSpeculationBarrier)
640 Modified |= insertCSDB(MBB, MBBI, DL);
641
642 return Modified;
643}
644
645bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
646 if (!MF.getFunction().hasFnAttribute(Kind: Attribute::SpeculativeLoadHardening))
647 return false;
648
649 MisspeculatingTaintReg = AArch64::X16;
650 MisspeculatingTaintReg32Bit = AArch64::W16;
651 TII = MF.getSubtarget().getInstrInfo();
652 TRI = MF.getSubtarget().getRegisterInfo();
653 RegsNeedingCSDBBeforeUse.resize(N: TRI->getNumRegs());
654 RegsAlreadyMasked.resize(N: TRI->getNumRegs());
655 UseControlFlowSpeculationBarrier = functionUsesHardeningRegister(MF);
656
657 bool Modified = false;
658
659 // Step 1: Enable automatic insertion of SpeculationSafeValue.
660 if (MF.getSubtarget<AArch64Subtarget>().getCLOpts().slh_loads) {
661 LLVM_DEBUG(
662 dbgs() << "***** AArch64SpeculationHardening - automatic insertion of "
663 "SpeculationSafeValue intrinsics *****\n");
664 for (auto &MBB : MF)
665 Modified |= slhLoads(MBB);
666 }
667
668 // 2. Add instrumentation code to function entry and exits.
669 LLVM_DEBUG(
670 dbgs()
671 << "***** AArch64SpeculationHardening - track control flow *****\n");
672
673 SmallVector<MachineBasicBlock *, 2> EntryBlocks;
674 EntryBlocks.push_back(Elt: &MF.front());
675 for (const LandingPadInfo &LPI : MF.getLandingPads())
676 EntryBlocks.push_back(Elt: LPI.LandingPadBlock);
677 for (auto *Entry : EntryBlocks)
678 insertSPToRegTaintPropagation(
679 MBB&: *Entry, MBBI: Entry->SkipPHIsLabelsAndDebug(I: Entry->begin()));
680
681 // 3. Add instrumentation code to every basic block.
682 for (auto &MBB : MF) {
683 bool UsesFullSpeculationBarrier = false;
684 Modified |= instrumentControlFlow(MBB, UsesFullSpeculationBarrier);
685 Modified |=
686 lowerSpeculationSafeValuePseudos(MBB, UsesFullSpeculationBarrier);
687 }
688
689 return Modified;
690}
691
692/// \brief Returns an instance of the pseudo instruction expansion pass.
693FunctionPass *llvm::createAArch64SpeculationHardeningPass() {
694 return new AArch64SpeculationHardening();
695}
696