llvm/lib/Target/AArch64/AArch64SLSHardening.cpp

   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::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) {
 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
 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
 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 &
 245 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
 377 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
 394 FunctionPass *llvm::createAArch64SLSHardeningPass() {
 395   return new AArch64SLSHardening();
 396 }
 397
 398 namespace {
 399 class AArch64IndirectThunks : public MachineFunctionPass {
 400 public:
 401   static char ID;
 402
 403   AArch64IndirectThunks() : MachineFunctionPass(ID) {}
 404
 405   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>
 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>
 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
 434 FunctionPass *llvm::createAArch64IndirectThunks() {
 435   return new AArch64IndirectThunks();
 436 }
 437
 438 bool AArch64IndirectThunks::doInitialization(Module &M) {
 439   initTIs(M, TIs);
 440   return false;
 441 }
 442
 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 }