528 lines · 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-exception6//7//===----------------------------------------------------------------------===//8//9// This file contains a pass to insert code to mitigate against side channel10// vulnerabilities that may happen under straight line miss-speculation.11//12//===----------------------------------------------------------------------===//13 14#include "AArch64InstrInfo.h"15#include "AArch64Subtarget.h"16#include "llvm/ADT/StringSwitch.h"17#include "llvm/CodeGen/IndirectThunks.h"18#include "llvm/CodeGen/MachineBasicBlock.h"19#include "llvm/CodeGen/MachineFunction.h"20#include "llvm/CodeGen/MachineInstr.h"21#include "llvm/CodeGen/MachineInstrBuilder.h"22#include "llvm/CodeGen/MachineOperand.h"23#include "llvm/CodeGen/RegisterScavenging.h"24#include "llvm/IR/DebugLoc.h"25#include "llvm/Pass.h"26#include "llvm/Support/ErrorHandling.h"27#include "llvm/Support/FormatVariadic.h"28#include "llvm/Target/TargetMachine.h"29#include <cassert>30#include <climits>31#include <tuple>32 33using namespace llvm;34 35#define DEBUG_TYPE "aarch64-sls-hardening"36 37#define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"38 39// Common name prefix of all thunks generated by this pass.40//41// The generic form is42// __llvm_slsblr_thunk_xN for BLR thunks43// __llvm_slsblr_thunk_(aaz|abz)_xN for BLRAAZ and BLRABZ thunks44// __llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks45static constexpr StringRef CommonNamePrefix = "__llvm_slsblr_thunk_";46 47namespace {48 49struct ThunkKind {50 enum ThunkKindId {51 ThunkBR,52 ThunkBRAA,53 ThunkBRAB,54 ThunkBRAAZ,55 ThunkBRABZ,56 };57 58 ThunkKindId Id;59 StringRef NameInfix;60 bool HasXmOperand;61 bool NeedsPAuth;62 63 // Opcode to perform indirect jump from inside the thunk.64 unsigned BROpcode;65 66 static const ThunkKind BR;67 static const ThunkKind BRAA;68 static const ThunkKind BRAB;69 static const ThunkKind BRAAZ;70 static const ThunkKind BRABZ;71};72 73// Set of inserted thunks.74class ThunksSet {75public:76 static constexpr unsigned NumXRegisters = 32;77 78 // Given Xn register, returns n.79 static unsigned indexOfXReg(Register Xn);80 // Given n, returns Xn register.81 static Register xRegByIndex(unsigned N);82 83 ThunksSet &operator|=(const ThunksSet &Other) {84 BLRThunks |= Other.BLRThunks;85 BLRAAZThunks |= Other.BLRAAZThunks;86 BLRABZThunks |= Other.BLRABZThunks;87 for (unsigned I = 0; I < NumXRegisters; ++I)88 BLRAAThunks[I] |= Other.BLRAAThunks[I];89 for (unsigned I = 0; I < NumXRegisters; ++I)90 BLRABThunks[I] |= Other.BLRABThunks[I];91 92 return *this;93 }94 95 bool get(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {96 reg_bitmask_t XnBit = reg_bitmask_t(1) << indexOfXReg(Xn);97 return getBitmask(Kind, Xm) & XnBit;98 }99 100 void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {101 reg_bitmask_t XnBit = reg_bitmask_t(1) << indexOfXReg(Xn);102 getBitmask(Kind, Xm) |= XnBit;103 }104 105private:106 typedef uint32_t reg_bitmask_t;107 static_assert(NumXRegisters <= sizeof(reg_bitmask_t) * CHAR_BIT,108 "Bitmask is not wide enough to hold all Xn registers");109 110 // Bitmasks representing operands used, with n-th bit corresponding to Xn111 // register operand. If the instruction has a second operand (Xm), an array112 // of bitmasks is used, indexed by m.113 // Indexes corresponding to the forbidden x16, x17 and x30 registers are114 // always unset, for simplicity there are no holes.115 reg_bitmask_t BLRThunks = 0;116 reg_bitmask_t BLRAAZThunks = 0;117 reg_bitmask_t BLRABZThunks = 0;118 reg_bitmask_t BLRAAThunks[NumXRegisters] = {};119 reg_bitmask_t BLRABThunks[NumXRegisters] = {};120 121 reg_bitmask_t &getBitmask(ThunkKind::ThunkKindId Kind, Register Xm) {122 switch (Kind) {123 case ThunkKind::ThunkBR:124 return BLRThunks;125 case ThunkKind::ThunkBRAAZ:126 return BLRAAZThunks;127 case ThunkKind::ThunkBRABZ:128 return BLRABZThunks;129 case ThunkKind::ThunkBRAA:130 return BLRAAThunks[indexOfXReg(Xm)];131 case ThunkKind::ThunkBRAB:132 return BLRABThunks[indexOfXReg(Xm)];133 }134 llvm_unreachable("Unknown ThunkKindId enum");135 }136};137 138struct SLSHardeningInserter : ThunkInserter<SLSHardeningInserter, ThunksSet> {139public:140 const char *getThunkPrefix() { return CommonNamePrefix.data(); }141 bool mayUseThunk(const MachineFunction &MF) {142 ComdatThunks &= !MF.getSubtarget<AArch64Subtarget>().hardenSlsNoComdat();143 // We are inserting barriers aside from thunk calls, so144 // check hardenSlsRetBr() as well.145 return MF.getSubtarget<AArch64Subtarget>().hardenSlsBlr() ||146 MF.getSubtarget<AArch64Subtarget>().hardenSlsRetBr();147 }148 ThunksSet insertThunks(MachineModuleInfo &MMI, MachineFunction &MF,149 ThunksSet ExistingThunks);150 void populateThunk(MachineFunction &MF);151 152private:153 bool ComdatThunks = true;154 155 bool hardenReturnsAndBRs(MachineModuleInfo &MMI, MachineBasicBlock &MBB);156 bool hardenBLRs(MachineModuleInfo &MMI, MachineBasicBlock &MBB,157 ThunksSet &Thunks);158 159 void convertBLRToBL(MachineModuleInfo &MMI, MachineBasicBlock &MBB,160 MachineBasicBlock::instr_iterator MBBI,161 ThunksSet &Thunks);162};163 164} // end anonymous namespace165 166const ThunkKind ThunkKind::BR = {ThunkBR, "", /*HasXmOperand=*/false,167 /*NeedsPAuth=*/false, AArch64::BR};168const ThunkKind ThunkKind::BRAA = {ThunkBRAA, "aa_", /*HasXmOperand=*/true,169 /*NeedsPAuth=*/true, AArch64::BRAA};170const ThunkKind ThunkKind::BRAB = {ThunkBRAB, "ab_", /*HasXmOperand=*/true,171 /*NeedsPAuth=*/true, AArch64::BRAB};172const ThunkKind ThunkKind::BRAAZ = {ThunkBRAAZ, "aaz_", /*HasXmOperand=*/false,173 /*NeedsPAuth=*/true, AArch64::BRAAZ};174const ThunkKind ThunkKind::BRABZ = {ThunkBRABZ, "abz_", /*HasXmOperand=*/false,175 /*NeedsPAuth=*/true, AArch64::BRABZ};176 177// Returns thunk kind to emit, or nullptr if not a BLR* instruction.178static const ThunkKind *getThunkKind(unsigned OriginalOpcode) {179 switch (OriginalOpcode) {180 case AArch64::BLR:181 case AArch64::BLRNoIP:182 return &ThunkKind::BR;183 case AArch64::BLRAA:184 return &ThunkKind::BRAA;185 case AArch64::BLRAB:186 return &ThunkKind::BRAB;187 case AArch64::BLRAAZ:188 return &ThunkKind::BRAAZ;189 case AArch64::BLRABZ:190 return &ThunkKind::BRABZ;191 }192 return nullptr;193}194 195static bool isBLR(const MachineInstr &MI) {196 return getThunkKind(MI.getOpcode()) != nullptr;197}198 199unsigned ThunksSet::indexOfXReg(Register Reg) {200 assert(AArch64::GPR64RegClass.contains(Reg));201 assert(Reg != AArch64::X16 && Reg != AArch64::X17 && Reg != AArch64::LR);202 203 // Most Xn registers have consecutive ids, except for FP and XZR.204 unsigned Result = (unsigned)Reg - (unsigned)AArch64::X0;205 if (Reg == AArch64::FP)206 Result = 29;207 else if (Reg == AArch64::XZR)208 Result = 31;209 210 assert(Result < NumXRegisters && "Internal register numbering changed");211 assert(AArch64::GPR64RegClass.getRegister(Result).id() == Reg &&212 "Internal register numbering changed");213 214 return Result;215}216 217Register ThunksSet::xRegByIndex(unsigned N) {218 return AArch64::GPR64RegClass.getRegister(N);219}220 221static void insertSpeculationBarrier(const AArch64Subtarget *ST,222 MachineBasicBlock &MBB,223 MachineBasicBlock::iterator MBBI,224 DebugLoc DL,225 bool AlwaysUseISBDSB = false) {226 assert(MBBI != MBB.begin() &&227 "Must not insert SpeculationBarrierEndBB as only instruction in MBB.");228 assert(std::prev(MBBI)->isBarrier() &&229 "SpeculationBarrierEndBB must only follow unconditional control flow "230 "instructions.");231 assert(std::prev(MBBI)->isTerminator() &&232 "SpeculationBarrierEndBB must only follow terminators.");233 const TargetInstrInfo *TII = ST->getInstrInfo();234 unsigned BarrierOpc = ST->hasSB() && !AlwaysUseISBDSB235 ? AArch64::SpeculationBarrierSBEndBB236 : AArch64::SpeculationBarrierISBDSBEndBB;237 if (MBBI == MBB.end() ||238 (MBBI->getOpcode() != AArch64::SpeculationBarrierSBEndBB &&239 MBBI->getOpcode() != AArch64::SpeculationBarrierISBDSBEndBB))240 BuildMI(MBB, MBBI, DL, TII->get(BarrierOpc));241}242 243ThunksSet SLSHardeningInserter::insertThunks(MachineModuleInfo &MMI,244 MachineFunction &MF,245 ThunksSet ExistingThunks) {246 const AArch64Subtarget *ST = &MF.getSubtarget<AArch64Subtarget>();247 248 for (auto &MBB : MF) {249 if (ST->hardenSlsRetBr())250 hardenReturnsAndBRs(MMI, MBB);251 if (ST->hardenSlsBlr())252 hardenBLRs(MMI, MBB, ExistingThunks);253 }254 return ExistingThunks;255}256 257bool SLSHardeningInserter::hardenReturnsAndBRs(MachineModuleInfo &MMI,258 MachineBasicBlock &MBB) {259 const AArch64Subtarget *ST =260 &MBB.getParent()->getSubtarget<AArch64Subtarget>();261 bool Modified = false;262 MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator(), E = MBB.end();263 MachineBasicBlock::iterator NextMBBI;264 for (; MBBI != E; MBBI = NextMBBI) {265 MachineInstr &MI = *MBBI;266 NextMBBI = std::next(MBBI);267 if (MI.isReturn() || isIndirectBranchOpcode(MI.getOpcode())) {268 assert(MI.isTerminator());269 insertSpeculationBarrier(ST, MBB, std::next(MBBI), MI.getDebugLoc());270 Modified = true;271 }272 }273 return Modified;274}275 276// Currently, the longest possible thunk name is277// __llvm_slsblr_thunk_aa_xNN_xMM278// which is 31 characters (without the '\0' character).279static SmallString<32> createThunkName(const ThunkKind &Kind, Register Xn,280 Register Xm) {281 unsigned N = ThunksSet::indexOfXReg(Xn);282 if (!Kind.HasXmOperand)283 return formatv("{0}{1}x{2}", CommonNamePrefix, Kind.NameInfix, N);284 285 unsigned M = ThunksSet::indexOfXReg(Xm);286 return formatv("{0}{1}x{2}_x{3}", CommonNamePrefix, Kind.NameInfix, N, M);287}288 289static std::tuple<const ThunkKind &, Register, Register>290parseThunkName(StringRef ThunkName) {291 assert(ThunkName.starts_with(CommonNamePrefix) &&292 "Should be filtered out by ThunkInserter");293 // Thunk name suffix, such as "x1" or "aa_x2_x3".294 StringRef NameSuffix = ThunkName.drop_front(CommonNamePrefix.size());295 296 // Parse thunk kind based on thunk name infix.297 const ThunkKind &Kind = *StringSwitch<const ThunkKind *>(NameSuffix)298 .StartsWith("aa_", &ThunkKind::BRAA)299 .StartsWith("ab_", &ThunkKind::BRAB)300 .StartsWith("aaz_", &ThunkKind::BRAAZ)301 .StartsWith("abz_", &ThunkKind::BRABZ)302 .Default(&ThunkKind::BR);303 304 auto ParseRegName = [](StringRef Name) {305 unsigned N;306 307 assert(Name.starts_with("x") && "xN register name expected");308 bool Fail = Name.drop_front(1).getAsInteger(/*Radix=*/10, N);309 assert(!Fail && N < ThunksSet::NumXRegisters && "Unexpected register");310 (void)Fail;311 312 return ThunksSet::xRegByIndex(N);313 };314 315 // For example, "x1" or "x2_x3".316 StringRef RegsStr = NameSuffix.drop_front(Kind.NameInfix.size());317 StringRef XnStr, XmStr;318 std::tie(XnStr, XmStr) = RegsStr.split('_');319 320 // Parse register operands.321 Register Xn = ParseRegName(XnStr);322 Register Xm = Kind.HasXmOperand ? ParseRegName(XmStr) : AArch64::NoRegister;323 324 return std::make_tuple(std::ref(Kind), Xn, Xm);325}326 327void SLSHardeningInserter::populateThunk(MachineFunction &MF) {328 assert(MF.getFunction().hasComdat() == ComdatThunks &&329 "ComdatThunks value changed since MF creation");330 Register Xn, Xm;331 auto KindAndRegs = parseThunkName(MF.getName());332 const ThunkKind &Kind = std::get<0>(KindAndRegs);333 std::tie(std::ignore, Xn, Xm) = KindAndRegs;334 335 const TargetInstrInfo *TII =336 MF.getSubtarget<AArch64Subtarget>().getInstrInfo();337 338 // Depending on whether this pass is in the same FunctionPassManager as the339 // IR->MIR conversion, the thunk may be completely empty, or contain a single340 // basic block with a single return instruction. Normalise it to contain a341 // single empty basic block.342 if (MF.size() == 1) {343 assert(MF.front().size() == 1);344 assert(MF.front().front().getOpcode() == AArch64::RET);345 MF.front().erase(MF.front().begin());346 } else {347 assert(MF.size() == 0);348 MF.push_back(MF.CreateMachineBasicBlock());349 }350 351 MachineBasicBlock *Entry = &MF.front();352 Entry->clear();353 354 // These thunks need to consist of the following instructions:355 // __llvm_slsblr_thunk_...:356 // MOV x16, xN ; BR* instructions are not compatible with "BTI c"357 // ; branch target unless xN is x16 or x17.358 // BR* ... ; One of: BR x16359 // ; BRA(A|B) x16, xM360 // ; BRA(A|B)Z x16361 // barrierInsts362 Entry->addLiveIn(Xn);363 // MOV X16, Reg == ORR X16, XZR, Reg, LSL #0364 BuildMI(Entry, DebugLoc(), TII->get(AArch64::ORRXrs), AArch64::X16)365 .addReg(AArch64::XZR)366 .addReg(Xn)367 .addImm(0);368 MachineInstrBuilder Builder =369 BuildMI(Entry, DebugLoc(), TII->get(Kind.BROpcode)).addReg(AArch64::X16);370 if (Xm != AArch64::NoRegister) {371 Entry->addLiveIn(Xm);372 Builder.addReg(Xm);373 }374 375 // Make sure the thunks do not make use of the SB extension in case there is376 // a function somewhere that will call to it that for some reason disabled377 // the SB extension locally on that function, even though it's enabled for378 // the module otherwise. Therefore set AlwaysUseISBSDB to true.379 insertSpeculationBarrier(&MF.getSubtarget<AArch64Subtarget>(), *Entry,380 Entry->end(), DebugLoc(), true /*AlwaysUseISBDSB*/);381}382 383void SLSHardeningInserter::convertBLRToBL(384 MachineModuleInfo &MMI, MachineBasicBlock &MBB,385 MachineBasicBlock::instr_iterator MBBI, ThunksSet &Thunks) {386 // Transform a BLR* instruction (one of BLR, BLRAA/BLRAB or BLRAAZ/BLRABZ) to387 // a BL to the thunk containing BR, BRAA/BRAB or BRAAZ/BRABZ, respectively.388 //389 // Before:390 // |-----------------------------|391 // | ... |392 // | instI |393 // | BLR* xN or BLR* xN, xM |394 // | instJ |395 // | ... |396 // |-----------------------------|397 //398 // After:399 // |-----------------------------|400 // | ... |401 // | instI |402 // | BL __llvm_slsblr_thunk_... |403 // | instJ |404 // | ... |405 // |-----------------------------|406 //407 // __llvm_slsblr_thunk_...:408 // |-----------------------------|409 // | MOV x16, xN |410 // | BR* x16 or BR* x16, xM |411 // | barrierInsts |412 // |-----------------------------|413 //414 // This function needs to transform BLR* instruction into BL with the correct415 // thunk name and lazily create the thunk if it does not exist yet.416 //417 // Since linkers are allowed to clobber X16 and X17 on function calls, the418 // above mitigation only works if the original BLR* instruction had neither419 // X16 nor X17 as one of its operands. Code generation before must make sure420 // that no such BLR* instruction was produced if the mitigation is enabled.421 422 MachineInstr &BLR = *MBBI;423 assert(isBLR(BLR));424 const ThunkKind &Kind = *getThunkKind(BLR.getOpcode());425 426 unsigned NumRegOperands = Kind.HasXmOperand ? 2 : 1;427 assert(BLR.getNumExplicitOperands() == NumRegOperands &&428 "Expected one or two register inputs");429 Register Xn = BLR.getOperand(0).getReg();430 Register Xm =431 Kind.HasXmOperand ? BLR.getOperand(1).getReg() : AArch64::NoRegister;432 433 DebugLoc DL = BLR.getDebugLoc();434 435 MachineFunction &MF = *MBBI->getMF();436 MCContext &Context = MBB.getParent()->getContext();437 const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();438 439 auto ThunkName = createThunkName(Kind, Xn, Xm);440 MCSymbol *Sym = Context.getOrCreateSymbol(ThunkName);441 442 if (!Thunks.get(Kind.Id, Xn, Xm)) {443 StringRef TargetAttrs = Kind.NeedsPAuth ? "+pauth" : "";444 Thunks.set(Kind.Id, Xn, Xm);445 createThunkFunction(MMI, ThunkName, ComdatThunks, TargetAttrs);446 }447 448 MachineInstr *BL = BuildMI(MBB, MBBI, DL, TII->get(AArch64::BL)).addSym(Sym);449 450 // Now copy the implicit operands from BLR to BL and copy other necessary451 // info.452 // However, both BLR and BL instructions implicitly use SP and implicitly453 // define LR. Blindly copying implicit operands would result in SP and LR454 // operands to be present multiple times. While this may not be too much of455 // an issue, let's avoid that for cleanliness, by removing those implicit456 // operands from the BL created above before we copy over all implicit457 // operands from the BLR.458 int ImpLROpIdx = -1;459 int ImpSPOpIdx = -1;460 for (unsigned OpIdx = BL->getNumExplicitOperands();461 OpIdx < BL->getNumOperands(); OpIdx++) {462 MachineOperand Op = BL->getOperand(OpIdx);463 if (!Op.isReg())464 continue;465 if (Op.getReg() == AArch64::LR && Op.isDef())466 ImpLROpIdx = OpIdx;467 if (Op.getReg() == AArch64::SP && !Op.isDef())468 ImpSPOpIdx = OpIdx;469 }470 assert(ImpLROpIdx != -1);471 assert(ImpSPOpIdx != -1);472 int FirstOpIdxToRemove = std::max(ImpLROpIdx, ImpSPOpIdx);473 int SecondOpIdxToRemove = std::min(ImpLROpIdx, ImpSPOpIdx);474 BL->removeOperand(FirstOpIdxToRemove);475 BL->removeOperand(SecondOpIdxToRemove);476 // Now copy over the implicit operands from the original BLR477 BL->copyImplicitOps(MF, BLR);478 MF.moveAdditionalCallInfo(&BLR, BL);479 // Also add the register operands of the original BLR* instruction480 // as being used in the called thunk.481 for (unsigned OpIdx = 0; OpIdx < NumRegOperands; ++OpIdx) {482 MachineOperand &Op = BLR.getOperand(OpIdx);483 BL->addOperand(MachineOperand::CreateReg(Op.getReg(), /*isDef=*/false,484 /*isImp=*/true, Op.isKill()));485 }486 // Remove BLR instruction487 MBB.erase(MBBI);488}489 490bool SLSHardeningInserter::hardenBLRs(MachineModuleInfo &MMI,491 MachineBasicBlock &MBB,492 ThunksSet &Thunks) {493 bool Modified = false;494 MachineBasicBlock::instr_iterator MBBI = MBB.instr_begin(),495 E = MBB.instr_end();496 MachineBasicBlock::instr_iterator NextMBBI;497 for (; MBBI != E; MBBI = NextMBBI) {498 MachineInstr &MI = *MBBI;499 NextMBBI = std::next(MBBI);500 if (isBLR(MI)) {501 convertBLRToBL(MMI, MBB, MBBI, Thunks);502 Modified = true;503 }504 }505 return Modified;506}507 508namespace {509class AArch64SLSHardening : public ThunkInserterPass<SLSHardeningInserter> {510public:511 static char ID;512 513 AArch64SLSHardening() : ThunkInserterPass(ID) {}514 515 StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }516};517 518} // end anonymous namespace519 520char AArch64SLSHardening::ID = 0;521 522INITIALIZE_PASS(AArch64SLSHardening, "aarch64-sls-hardening",523 AARCH64_SLS_HARDENING_NAME, false, false)524 525FunctionPass *llvm::createAArch64SLSHardeningPass() {526 return new AArch64SLSHardening();527}528