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