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