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