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 | assert (MF.size() == 1); |
224 | MachineBasicBlock *Entry = &MF.front(); |
225 | Entry->clear(); |
226 | |
227 | // These thunks need to consist of the following instructions: |
228 | // __llvm_slsblr_thunk_xN: |
229 | // BR xN |
230 | // barrierInsts |
231 | Entry->addLiveIn(PhysReg: ThunkReg); |
232 | // MOV X16, ThunkReg == ORR X16, XZR, ThunkReg, LSL #0 |
233 | BuildMI(Entry, DebugLoc(), TII->get(AArch64::ORRXrs), AArch64::X16) |
234 | .addReg(AArch64::XZR) |
235 | .addReg(ThunkReg) |
236 | .addImm(0); |
237 | BuildMI(Entry, DebugLoc(), TII->get(AArch64::BR)).addReg(AArch64::X16); |
238 | // Make sure the thunks do not make use of the SB extension in case there is |
239 | // a function somewhere that will call to it that for some reason disabled |
240 | // the SB extension locally on that function, even though it's enabled for |
241 | // the module otherwise. Therefore set AlwaysUseISBSDB to true. |
242 | insertSpeculationBarrier(ST: &MF.getSubtarget<AArch64Subtarget>(), MBB&: *Entry, |
243 | MBBI: Entry->end(), DL: DebugLoc(), AlwaysUseISBDSB: true /*AlwaysUseISBDSB*/); |
244 | } |
245 | |
246 | MachineBasicBlock &AArch64SLSHardening::ConvertBLRToBL( |
247 | MachineBasicBlock &MBB, MachineBasicBlock::instr_iterator MBBI) const { |
248 | // Transform a BLR to a BL as follows: |
249 | // Before: |
250 | // |-----------------------------| |
251 | // | ... | |
252 | // | instI | |
253 | // | BLR xN | |
254 | // | instJ | |
255 | // | ... | |
256 | // |-----------------------------| |
257 | // |
258 | // After: |
259 | // |-----------------------------| |
260 | // | ... | |
261 | // | instI | |
262 | // | BL __llvm_slsblr_thunk_xN | |
263 | // | instJ | |
264 | // | ... | |
265 | // |-----------------------------| |
266 | // |
267 | // __llvm_slsblr_thunk_xN: |
268 | // |-----------------------------| |
269 | // | BR xN | |
270 | // | barrierInsts | |
271 | // |-----------------------------| |
272 | // |
273 | // The __llvm_slsblr_thunk_xN thunks are created by the SLSBLRThunkInserter. |
274 | // This function merely needs to transform BLR xN into BL |
275 | // __llvm_slsblr_thunk_xN. |
276 | // |
277 | // Since linkers are allowed to clobber X16 and X17 on function calls, the |
278 | // above mitigation only works if the original BLR instruction was not |
279 | // BLR X16 nor BLR X17. Code generation before must make sure that no BLR |
280 | // X16|X17 was produced if the mitigation is enabled. |
281 | |
282 | MachineInstr &BLR = *MBBI; |
283 | assert(isBLR(BLR)); |
284 | unsigned BLOpcode; |
285 | Register Reg; |
286 | bool RegIsKilled; |
287 | switch (BLR.getOpcode()) { |
288 | case AArch64::BLR: |
289 | case AArch64::BLRNoIP: |
290 | BLOpcode = AArch64::BL; |
291 | Reg = BLR.getOperand(i: 0).getReg(); |
292 | assert(Reg != AArch64::X16 && Reg != AArch64::X17 && Reg != AArch64::LR); |
293 | RegIsKilled = BLR.getOperand(i: 0).isKill(); |
294 | break; |
295 | case AArch64::BLRAA: |
296 | case AArch64::BLRAB: |
297 | case AArch64::BLRAAZ: |
298 | case AArch64::BLRABZ: |
299 | llvm_unreachable("BLRA instructions cannot yet be produced by LLVM, " |
300 | "therefore there is no need to support them for now." ); |
301 | default: |
302 | llvm_unreachable("unhandled BLR" ); |
303 | } |
304 | DebugLoc DL = BLR.getDebugLoc(); |
305 | |
306 | // If we'd like to support also BLRAA and BLRAB instructions, we'd need |
307 | // a lot more different kind of thunks. |
308 | // For example, a |
309 | // |
310 | // BLRAA xN, xM |
311 | // |
312 | // instruction probably would need to be transformed to something like: |
313 | // |
314 | // BL __llvm_slsblraa_thunk_x<N>_x<M> |
315 | // |
316 | // __llvm_slsblraa_thunk_x<N>_x<M>: |
317 | // BRAA x<N>, x<M> |
318 | // barrierInsts |
319 | // |
320 | // Given that about 30 different values of N are possible and about 30 |
321 | // different values of M are possible in the above, with the current way |
322 | // of producing indirect thunks, we'd be producing about 30 times 30, i.e. |
323 | // about 900 thunks (where most might not be actually called). This would |
324 | // multiply further by two to support both BLRAA and BLRAB variants of those |
325 | // instructions. |
326 | // If we'd want to support this, we'd probably need to look into a different |
327 | // way to produce thunk functions, based on which variants are actually |
328 | // needed, rather than producing all possible variants. |
329 | // So far, LLVM does never produce BLRA* instructions, so let's leave this |
330 | // for the future when LLVM can start producing BLRA* instructions. |
331 | MachineFunction &MF = *MBBI->getMF(); |
332 | MCContext &Context = MBB.getParent()->getContext(); |
333 | auto ThunkIt = |
334 | llvm::find_if(SLSBLRThunks, [Reg](auto T) { return T.Reg == Reg; }); |
335 | assert (ThunkIt != std::end(SLSBLRThunks)); |
336 | MCSymbol *Sym = Context.getOrCreateSymbol(Name: ThunkIt->Name); |
337 | |
338 | MachineInstr *BL = BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: BLOpcode)).addSym(Sym); |
339 | |
340 | // Now copy the implicit operands from BLR to BL and copy other necessary |
341 | // info. |
342 | // However, both BLR and BL instructions implictly use SP and implicitly |
343 | // define LR. Blindly copying implicit operands would result in SP and LR |
344 | // operands to be present multiple times. While this may not be too much of |
345 | // an issue, let's avoid that for cleanliness, by removing those implicit |
346 | // operands from the BL created above before we copy over all implicit |
347 | // operands from the BLR. |
348 | int ImpLROpIdx = -1; |
349 | int ImpSPOpIdx = -1; |
350 | for (unsigned OpIdx = BL->getNumExplicitOperands(); |
351 | OpIdx < BL->getNumOperands(); OpIdx++) { |
352 | MachineOperand Op = BL->getOperand(i: OpIdx); |
353 | if (!Op.isReg()) |
354 | continue; |
355 | if (Op.getReg() == AArch64::LR && Op.isDef()) |
356 | ImpLROpIdx = OpIdx; |
357 | if (Op.getReg() == AArch64::SP && !Op.isDef()) |
358 | ImpSPOpIdx = OpIdx; |
359 | } |
360 | assert(ImpLROpIdx != -1); |
361 | assert(ImpSPOpIdx != -1); |
362 | int FirstOpIdxToRemove = std::max(a: ImpLROpIdx, b: ImpSPOpIdx); |
363 | int SecondOpIdxToRemove = std::min(a: ImpLROpIdx, b: ImpSPOpIdx); |
364 | BL->removeOperand(OpNo: FirstOpIdxToRemove); |
365 | BL->removeOperand(OpNo: SecondOpIdxToRemove); |
366 | // Now copy over the implicit operands from the original BLR |
367 | BL->copyImplicitOps(MF, MI: BLR); |
368 | MF.moveCallSiteInfo(Old: &BLR, New: BL); |
369 | // Also add the register called in the BLR as being used in the called thunk. |
370 | BL->addOperand(Op: MachineOperand::CreateReg(Reg, isDef: false /*isDef*/, isImp: true /*isImp*/, |
371 | isKill: RegIsKilled /*isKill*/)); |
372 | // Remove BLR instruction |
373 | MBB.erase(I: MBBI); |
374 | |
375 | return MBB; |
376 | } |
377 | |
378 | bool AArch64SLSHardening::hardenBLRs(MachineBasicBlock &MBB) const { |
379 | if (!ST->hardenSlsBlr()) |
380 | return false; |
381 | bool Modified = false; |
382 | MachineBasicBlock::instr_iterator MBBI = MBB.instr_begin(), |
383 | E = MBB.instr_end(); |
384 | MachineBasicBlock::instr_iterator NextMBBI; |
385 | for (; MBBI != E; MBBI = NextMBBI) { |
386 | MachineInstr &MI = *MBBI; |
387 | NextMBBI = std::next(x: MBBI); |
388 | if (isBLR(MI)) { |
389 | ConvertBLRToBL(MBB, MBBI); |
390 | Modified = true; |
391 | } |
392 | } |
393 | return Modified; |
394 | } |
395 | |
396 | FunctionPass *llvm::createAArch64SLSHardeningPass() { |
397 | return new AArch64SLSHardening(); |
398 | } |
399 | |
400 | namespace { |
401 | class AArch64IndirectThunks : public MachineFunctionPass { |
402 | public: |
403 | static char ID; |
404 | |
405 | AArch64IndirectThunks() : MachineFunctionPass(ID) {} |
406 | |
407 | StringRef getPassName() const override { return "AArch64 Indirect Thunks" ; } |
408 | |
409 | bool doInitialization(Module &M) override; |
410 | bool runOnMachineFunction(MachineFunction &MF) override; |
411 | |
412 | private: |
413 | std::tuple<SLSBLRThunkInserter> TIs; |
414 | |
415 | // FIXME: When LLVM moves to C++17, these can become folds |
416 | template <typename... ThunkInserterT> |
417 | static void initTIs(Module &M, |
418 | std::tuple<ThunkInserterT...> &ThunkInserters) { |
419 | (void)std::initializer_list<int>{ |
420 | (std::get<ThunkInserterT>(ThunkInserters).init(M), 0)...}; |
421 | } |
422 | template <typename... ThunkInserterT> |
423 | static bool runTIs(MachineModuleInfo &MMI, MachineFunction &MF, |
424 | std::tuple<ThunkInserterT...> &ThunkInserters) { |
425 | bool Modified = false; |
426 | (void)std::initializer_list<int>{ |
427 | Modified |= std::get<ThunkInserterT>(ThunkInserters).run(MMI, MF)...}; |
428 | return Modified; |
429 | } |
430 | }; |
431 | |
432 | } // end anonymous namespace |
433 | |
434 | char AArch64IndirectThunks::ID = 0; |
435 | |
436 | FunctionPass *llvm::createAArch64IndirectThunks() { |
437 | return new AArch64IndirectThunks(); |
438 | } |
439 | |
440 | bool AArch64IndirectThunks::doInitialization(Module &M) { |
441 | initTIs(M, ThunkInserters&: TIs); |
442 | return false; |
443 | } |
444 | |
445 | bool AArch64IndirectThunks::runOnMachineFunction(MachineFunction &MF) { |
446 | LLVM_DEBUG(dbgs() << getPassName() << '\n'); |
447 | auto &MMI = getAnalysis<MachineModuleInfoWrapperPass>().getMMI(); |
448 | return runTIs(MMI, MF, ThunkInserters&: TIs); |
449 | } |
450 | |