1//===- AArch64SLSHardening.cpp - Harden Straight Line 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 straight line miss-speculation.
11//
12//===----------------------------------------------------------------------===//
13
14#include "AArch64InstrInfo.h"
15#include "AArch64Subtarget.h"
16#include "Utils/AArch64BaseInfo.h"
17#include "llvm/CodeGen/IndirectThunks.h"
18#include "llvm/CodeGen/MachineBasicBlock.h"
19#include "llvm/CodeGen/MachineFunction.h"
20#include "llvm/CodeGen/MachineFunctionPass.h"
21#include "llvm/CodeGen/MachineInstr.h"
22#include "llvm/CodeGen/MachineInstrBuilder.h"
23#include "llvm/CodeGen/MachineOperand.h"
24#include "llvm/CodeGen/MachineRegisterInfo.h"
25#include "llvm/CodeGen/RegisterScavenging.h"
26#include "llvm/IR/DebugLoc.h"
27#include "llvm/Pass.h"
28#include "llvm/Support/CodeGen.h"
29#include "llvm/Support/Debug.h"
30#include "llvm/Target/TargetMachine.h"
31#include <cassert>
32
33using namespace llvm;
34
35#define DEBUG_TYPE "aarch64-sls-hardening"
36
37#define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
38
39namespace {
40
41class AArch64SLSHardening : public MachineFunctionPass {
42public:
43 const TargetInstrInfo *TII;
44 const TargetRegisterInfo *TRI;
45 const AArch64Subtarget *ST;
46
47 static char ID;
48
49 AArch64SLSHardening() : MachineFunctionPass(ID) {
50 initializeAArch64SLSHardeningPass(*PassRegistry::getPassRegistry());
51 }
52
53 bool runOnMachineFunction(MachineFunction &Fn) override;
54
55 StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }
56
57private:
58 bool hardenReturnsAndBRs(MachineBasicBlock &MBB) const;
59 bool hardenBLRs(MachineBasicBlock &MBB) const;
60 MachineBasicBlock &ConvertBLRToBL(MachineBasicBlock &MBB,
61 MachineBasicBlock::instr_iterator) const;
62};
63
64} // end anonymous namespace
65
66char AArch64SLSHardening::ID = 0;
67
68INITIALIZE_PASS(AArch64SLSHardening, "aarch64-sls-hardening",
69 AARCH64_SLS_HARDENING_NAME, false, false)
70
71static void insertSpeculationBarrier(const AArch64Subtarget *ST,
72 MachineBasicBlock &MBB,
73 MachineBasicBlock::iterator MBBI,
74 DebugLoc DL,
75 bool AlwaysUseISBDSB = false) {
76 assert(MBBI != MBB.begin() &&
77 "Must not insert SpeculationBarrierEndBB as only instruction in MBB.");
78 assert(std::prev(MBBI)->isBarrier() &&
79 "SpeculationBarrierEndBB must only follow unconditional control flow "
80 "instructions.");
81 assert(std::prev(MBBI)->isTerminator() &&
82 "SpeculationBarrierEndBB must only follow terminators.");
83 const TargetInstrInfo *TII = ST->getInstrInfo();
84 unsigned BarrierOpc = ST->hasSB() && !AlwaysUseISBDSB
85 ? AArch64::SpeculationBarrierSBEndBB
86 : AArch64::SpeculationBarrierISBDSBEndBB;
87 if (MBBI == MBB.end() ||
88 (MBBI->getOpcode() != AArch64::SpeculationBarrierSBEndBB &&
89 MBBI->getOpcode() != AArch64::SpeculationBarrierISBDSBEndBB))
90 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: BarrierOpc));
91}
92
93bool AArch64SLSHardening::runOnMachineFunction(MachineFunction &MF) {
94 ST = &MF.getSubtarget<AArch64Subtarget>();
95 TII = MF.getSubtarget().getInstrInfo();
96 TRI = MF.getSubtarget().getRegisterInfo();
97
98 bool Modified = false;
99 for (auto &MBB : MF) {
100 Modified |= hardenReturnsAndBRs(MBB);
101 Modified |= hardenBLRs(MBB);
102 }
103
104 return Modified;
105}
106
107static bool isBLR(const MachineInstr &MI) {
108 switch (MI.getOpcode()) {
109 case AArch64::BLR:
110 case AArch64::BLRNoIP:
111 return true;
112 case AArch64::BLRAA:
113 case AArch64::BLRAB:
114 case AArch64::BLRAAZ:
115 case AArch64::BLRABZ:
116 llvm_unreachable("Currently, LLVM's code generator does not support "
117 "producing BLRA* instructions. Therefore, there's no "
118 "support in this pass for those instructions.");
119 }
120 return false;
121}
122
123bool AArch64SLSHardening::hardenReturnsAndBRs(MachineBasicBlock &MBB) const {
124 if (!ST->hardenSlsRetBr())
125 return false;
126 bool Modified = false;
127 MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator(), E = MBB.end();
128 MachineBasicBlock::iterator NextMBBI;
129 for (; MBBI != E; MBBI = NextMBBI) {
130 MachineInstr &MI = *MBBI;
131 NextMBBI = std::next(x: MBBI);
132 if (MI.isReturn() || isIndirectBranchOpcode(Opc: MI.getOpcode())) {
133 assert(MI.isTerminator());
134 insertSpeculationBarrier(ST, MBB, MBBI: std::next(x: MBBI), DL: MI.getDebugLoc());
135 Modified = true;
136 }
137 }
138 return Modified;
139}
140
141static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
142
143static const struct ThunkNameAndReg {
144 const char* Name;
145 Register Reg;
146} SLSBLRThunks[] = {
147 { "__llvm_slsblr_thunk_x0", AArch64::X0},
148 { "__llvm_slsblr_thunk_x1", AArch64::X1},
149 { "__llvm_slsblr_thunk_x2", AArch64::X2},
150 { "__llvm_slsblr_thunk_x3", AArch64::X3},
151 { "__llvm_slsblr_thunk_x4", AArch64::X4},
152 { "__llvm_slsblr_thunk_x5", AArch64::X5},
153 { "__llvm_slsblr_thunk_x6", AArch64::X6},
154 { "__llvm_slsblr_thunk_x7", AArch64::X7},
155 { "__llvm_slsblr_thunk_x8", AArch64::X8},
156 { "__llvm_slsblr_thunk_x9", AArch64::X9},
157 { "__llvm_slsblr_thunk_x10", AArch64::X10},
158 { "__llvm_slsblr_thunk_x11", AArch64::X11},
159 { "__llvm_slsblr_thunk_x12", AArch64::X12},
160 { "__llvm_slsblr_thunk_x13", AArch64::X13},
161 { "__llvm_slsblr_thunk_x14", AArch64::X14},
162 { "__llvm_slsblr_thunk_x15", AArch64::X15},
163 // X16 and X17 are deliberately missing, as the mitigation requires those
164 // register to not be used in BLR. See comment in ConvertBLRToBL for more
165 // details.
166 { "__llvm_slsblr_thunk_x18", AArch64::X18},
167 { "__llvm_slsblr_thunk_x19", AArch64::X19},
168 { "__llvm_slsblr_thunk_x20", AArch64::X20},
169 { "__llvm_slsblr_thunk_x21", AArch64::X21},
170 { "__llvm_slsblr_thunk_x22", AArch64::X22},
171 { "__llvm_slsblr_thunk_x23", AArch64::X23},
172 { "__llvm_slsblr_thunk_x24", AArch64::X24},
173 { "__llvm_slsblr_thunk_x25", AArch64::X25},
174 { "__llvm_slsblr_thunk_x26", AArch64::X26},
175 { "__llvm_slsblr_thunk_x27", AArch64::X27},
176 { "__llvm_slsblr_thunk_x28", AArch64::X28},
177 { "__llvm_slsblr_thunk_x29", AArch64::FP},
178 // X30 is deliberately missing, for similar reasons as X16 and X17 are
179 // missing.
180 { "__llvm_slsblr_thunk_x31", AArch64::XZR},
181};
182
183namespace {
184struct SLSBLRThunkInserter : ThunkInserter<SLSBLRThunkInserter> {
185 const char *getThunkPrefix() { return SLSBLRNamePrefix; }
186 bool mayUseThunk(const MachineFunction &MF, bool InsertedThunks) {
187 if (InsertedThunks)
188 return false;
189 ComdatThunks &= !MF.getSubtarget<AArch64Subtarget>().hardenSlsNoComdat();
190 // FIXME: This could also check if there are any BLRs in the function
191 // to more accurately reflect if a thunk will be needed.
192 return MF.getSubtarget<AArch64Subtarget>().hardenSlsBlr();
193 }
194 bool insertThunks(MachineModuleInfo &MMI, MachineFunction &MF);
195 void populateThunk(MachineFunction &MF);
196
197private:
198 bool ComdatThunks = true;
199};
200} // namespace
201
202bool SLSBLRThunkInserter::insertThunks(MachineModuleInfo &MMI,
203 MachineFunction &MF) {
204 // FIXME: It probably would be possible to filter which thunks to produce
205 // based on which registers are actually used in BLR instructions in this
206 // function. But would that be a worthwhile optimization?
207 for (auto T : SLSBLRThunks)
208 createThunkFunction(MMI, T.Name, ComdatThunks);
209 return true;
210}
211
212void SLSBLRThunkInserter::populateThunk(MachineFunction &MF) {
213 // FIXME: How to better communicate Register number, rather than through
214 // name and lookup table?
215 assert(MF.getName().starts_with(getThunkPrefix()));
216 auto ThunkIt = llvm::find_if(
217 SLSBLRThunks, [&MF](auto T) { return T.Name == MF.getName(); });
218 assert(ThunkIt != std::end(SLSBLRThunks));
219 Register ThunkReg = ThunkIt->Reg;
220
221 const TargetInstrInfo *TII =
222 MF.getSubtarget<AArch64Subtarget>().getInstrInfo();
223
224 // Depending on whether this pass is in the same FunctionPassManager as the
225 // IR->MIR conversion, the thunk may be completely empty, or contain a single
226 // basic block with a single return instruction. Normalise it to contain a
227 // single empty basic block.
228 if (MF.size() == 1) {
229 assert(MF.front().size() == 1);
230 assert(MF.front().front().getOpcode() == AArch64::RET);
231 MF.front().erase(I: MF.front().begin());
232 } else {
233 assert(MF.size() == 0);
234 MF.push_back(MBB: MF.CreateMachineBasicBlock());
235 }
236
237 MachineBasicBlock *Entry = &MF.front();
238 Entry->clear();
239
240 // These thunks need to consist of the following instructions:
241 // __llvm_slsblr_thunk_xN:
242 // BR xN
243 // barrierInsts
244 Entry->addLiveIn(PhysReg: ThunkReg);
245 // MOV X16, ThunkReg == ORR X16, XZR, ThunkReg, LSL #0
246 BuildMI(Entry, DebugLoc(), TII->get(AArch64::ORRXrs), AArch64::X16)
247 .addReg(AArch64::XZR)
248 .addReg(ThunkReg)
249 .addImm(0);
250 BuildMI(Entry, DebugLoc(), TII->get(AArch64::BR)).addReg(AArch64::X16);
251 // Make sure the thunks do not make use of the SB extension in case there is
252 // a function somewhere that will call to it that for some reason disabled
253 // the SB extension locally on that function, even though it's enabled for
254 // the module otherwise. Therefore set AlwaysUseISBSDB to true.
255 insertSpeculationBarrier(ST: &MF.getSubtarget<AArch64Subtarget>(), MBB&: *Entry,
256 MBBI: Entry->end(), DL: DebugLoc(), AlwaysUseISBDSB: true /*AlwaysUseISBDSB*/);
257}
258
259MachineBasicBlock &AArch64SLSHardening::ConvertBLRToBL(
260 MachineBasicBlock &MBB, MachineBasicBlock::instr_iterator MBBI) const {
261 // Transform a BLR to a BL as follows:
262 // Before:
263 // |-----------------------------|
264 // | ... |
265 // | instI |
266 // | BLR xN |
267 // | instJ |
268 // | ... |
269 // |-----------------------------|
270 //
271 // After:
272 // |-----------------------------|
273 // | ... |
274 // | instI |
275 // | BL __llvm_slsblr_thunk_xN |
276 // | instJ |
277 // | ... |
278 // |-----------------------------|
279 //
280 // __llvm_slsblr_thunk_xN:
281 // |-----------------------------|
282 // | BR xN |
283 // | barrierInsts |
284 // |-----------------------------|
285 //
286 // The __llvm_slsblr_thunk_xN thunks are created by the SLSBLRThunkInserter.
287 // This function merely needs to transform BLR xN into BL
288 // __llvm_slsblr_thunk_xN.
289 //
290 // Since linkers are allowed to clobber X16 and X17 on function calls, the
291 // above mitigation only works if the original BLR instruction was not
292 // BLR X16 nor BLR X17. Code generation before must make sure that no BLR
293 // X16|X17 was produced if the mitigation is enabled.
294
295 MachineInstr &BLR = *MBBI;
296 assert(isBLR(BLR));
297 unsigned BLOpcode;
298 Register Reg;
299 bool RegIsKilled;
300 switch (BLR.getOpcode()) {
301 case AArch64::BLR:
302 case AArch64::BLRNoIP:
303 BLOpcode = AArch64::BL;
304 Reg = BLR.getOperand(i: 0).getReg();
305 assert(Reg != AArch64::X16 && Reg != AArch64::X17 && Reg != AArch64::LR);
306 RegIsKilled = BLR.getOperand(i: 0).isKill();
307 break;
308 case AArch64::BLRAA:
309 case AArch64::BLRAB:
310 case AArch64::BLRAAZ:
311 case AArch64::BLRABZ:
312 llvm_unreachable("BLRA instructions cannot yet be produced by LLVM, "
313 "therefore there is no need to support them for now.");
314 default:
315 llvm_unreachable("unhandled BLR");
316 }
317 DebugLoc DL = BLR.getDebugLoc();
318
319 // If we'd like to support also BLRAA and BLRAB instructions, we'd need
320 // a lot more different kind of thunks.
321 // For example, a
322 //
323 // BLRAA xN, xM
324 //
325 // instruction probably would need to be transformed to something like:
326 //
327 // BL __llvm_slsblraa_thunk_x<N>_x<M>
328 //
329 // __llvm_slsblraa_thunk_x<N>_x<M>:
330 // BRAA x<N>, x<M>
331 // barrierInsts
332 //
333 // Given that about 30 different values of N are possible and about 30
334 // different values of M are possible in the above, with the current way
335 // of producing indirect thunks, we'd be producing about 30 times 30, i.e.
336 // about 900 thunks (where most might not be actually called). This would
337 // multiply further by two to support both BLRAA and BLRAB variants of those
338 // instructions.
339 // If we'd want to support this, we'd probably need to look into a different
340 // way to produce thunk functions, based on which variants are actually
341 // needed, rather than producing all possible variants.
342 // So far, LLVM does never produce BLRA* instructions, so let's leave this
343 // for the future when LLVM can start producing BLRA* instructions.
344 MachineFunction &MF = *MBBI->getMF();
345 MCContext &Context = MBB.getParent()->getContext();
346 auto ThunkIt =
347 llvm::find_if(SLSBLRThunks, [Reg](auto T) { return T.Reg == Reg; });
348 assert (ThunkIt != std::end(SLSBLRThunks));
349 MCSymbol *Sym = Context.getOrCreateSymbol(Name: ThunkIt->Name);
350
351 MachineInstr *BL = BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: BLOpcode)).addSym(Sym);
352
353 // Now copy the implicit operands from BLR to BL and copy other necessary
354 // info.
355 // However, both BLR and BL instructions implictly use SP and implicitly
356 // define LR. Blindly copying implicit operands would result in SP and LR
357 // operands to be present multiple times. While this may not be too much of
358 // an issue, let's avoid that for cleanliness, by removing those implicit
359 // operands from the BL created above before we copy over all implicit
360 // operands from the BLR.
361 int ImpLROpIdx = -1;
362 int ImpSPOpIdx = -1;
363 for (unsigned OpIdx = BL->getNumExplicitOperands();
364 OpIdx < BL->getNumOperands(); OpIdx++) {
365 MachineOperand Op = BL->getOperand(i: OpIdx);
366 if (!Op.isReg())
367 continue;
368 if (Op.getReg() == AArch64::LR && Op.isDef())
369 ImpLROpIdx = OpIdx;
370 if (Op.getReg() == AArch64::SP && !Op.isDef())
371 ImpSPOpIdx = OpIdx;
372 }
373 assert(ImpLROpIdx != -1);
374 assert(ImpSPOpIdx != -1);
375 int FirstOpIdxToRemove = std::max(a: ImpLROpIdx, b: ImpSPOpIdx);
376 int SecondOpIdxToRemove = std::min(a: ImpLROpIdx, b: ImpSPOpIdx);
377 BL->removeOperand(OpNo: FirstOpIdxToRemove);
378 BL->removeOperand(OpNo: SecondOpIdxToRemove);
379 // Now copy over the implicit operands from the original BLR
380 BL->copyImplicitOps(MF, MI: BLR);
381 MF.moveCallSiteInfo(Old: &BLR, New: BL);
382 // Also add the register called in the BLR as being used in the called thunk.
383 BL->addOperand(Op: MachineOperand::CreateReg(Reg, isDef: false /*isDef*/, isImp: true /*isImp*/,
384 isKill: RegIsKilled /*isKill*/));
385 // Remove BLR instruction
386 MBB.erase(I: MBBI);
387
388 return MBB;
389}
390
391bool AArch64SLSHardening::hardenBLRs(MachineBasicBlock &MBB) const {
392 if (!ST->hardenSlsBlr())
393 return false;
394 bool Modified = false;
395 MachineBasicBlock::instr_iterator MBBI = MBB.instr_begin(),
396 E = MBB.instr_end();
397 MachineBasicBlock::instr_iterator NextMBBI;
398 for (; MBBI != E; MBBI = NextMBBI) {
399 MachineInstr &MI = *MBBI;
400 NextMBBI = std::next(x: MBBI);
401 if (isBLR(MI)) {
402 ConvertBLRToBL(MBB, MBBI);
403 Modified = true;
404 }
405 }
406 return Modified;
407}
408
409FunctionPass *llvm::createAArch64SLSHardeningPass() {
410 return new AArch64SLSHardening();
411}
412
413namespace {
414class AArch64IndirectThunks : public MachineFunctionPass {
415public:
416 static char ID;
417
418 AArch64IndirectThunks() : MachineFunctionPass(ID) {}
419
420 StringRef getPassName() const override { return "AArch64 Indirect Thunks"; }
421
422 bool doInitialization(Module &M) override;
423 bool runOnMachineFunction(MachineFunction &MF) override;
424
425private:
426 std::tuple<SLSBLRThunkInserter> TIs;
427
428 template <typename... ThunkInserterT>
429 static void initTIs(Module &M,
430 std::tuple<ThunkInserterT...> &ThunkInserters) {
431 (..., std::get<ThunkInserterT>(ThunkInserters).init(M));
432 }
433 template <typename... ThunkInserterT>
434 static bool runTIs(MachineModuleInfo &MMI, MachineFunction &MF,
435 std::tuple<ThunkInserterT...> &ThunkInserters) {
436 return (0 | ... | std::get<ThunkInserterT>(ThunkInserters).run(MMI, MF));
437 }
438};
439
440} // end anonymous namespace
441
442char AArch64IndirectThunks::ID = 0;
443
444FunctionPass *llvm::createAArch64IndirectThunks() {
445 return new AArch64IndirectThunks();
446}
447
448bool AArch64IndirectThunks::doInitialization(Module &M) {
449 initTIs(M, ThunkInserters&: TIs);
450 return false;
451}
452
453bool AArch64IndirectThunks::runOnMachineFunction(MachineFunction &MF) {
454 LLVM_DEBUG(dbgs() << getPassName() << '\n');
455 auto &MMI = getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
456 return runTIs(MMI, MF, ThunkInserters&: TIs);
457}
458

source code of llvm/lib/Target/AArch64/AArch64SLSHardening.cpp