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 | |
33 | using namespace llvm; |
34 | |
35 | #define DEBUG_TYPE "aarch64-sls-hardening" |
36 | |
37 | #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass" |
38 | |
39 | namespace { |
40 | |
41 | class AArch64SLSHardening : public MachineFunctionPass { |
42 | public: |
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 | |
57 | private: |
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 | |
66 | char AArch64SLSHardening::ID = 0; |
67 | |
68 | INITIALIZE_PASS(AArch64SLSHardening, "aarch64-sls-hardening" , |
69 | AARCH64_SLS_HARDENING_NAME, false, false) |
70 | |
71 | static 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 | |
93 | bool 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 | |
107 | static 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 | |
123 | bool 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 | |
141 | static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_" ; |
142 | |
143 | static 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 | |
183 | namespace { |
184 | struct 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 | |
197 | private: |
198 | bool ComdatThunks = true; |
199 | }; |
200 | } // namespace |
201 | |
202 | bool 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 | |
212 | void 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 | |
259 | MachineBasicBlock &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 | |
391 | bool 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 | |
409 | FunctionPass *llvm::createAArch64SLSHardeningPass() { |
410 | return new AArch64SLSHardening(); |
411 | } |
412 | |
413 | namespace { |
414 | class AArch64IndirectThunks : public MachineFunctionPass { |
415 | public: |
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 | |
425 | private: |
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 | |
442 | char AArch64IndirectThunks::ID = 0; |
443 | |
444 | FunctionPass *llvm::createAArch64IndirectThunks() { |
445 | return new AArch64IndirectThunks(); |
446 | } |
447 | |
448 | bool AArch64IndirectThunks::doInitialization(Module &M) { |
449 | initTIs(M, ThunkInserters&: TIs); |
450 | return false; |
451 | } |
452 | |
453 | bool 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 | |