llvm/lib/Target/AArch64/AArch64SLSHardening.cpp - llvm-project.git - Git at Google

 //===- AArch64SLSHardening.cpp - Harden Straight Line Missspeculation -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains a pass to insert code to mitigate against side channel
 // vulnerabilities that may happen under straight line miss-speculation.
 //
 //===----------------------------------------------------------------------===//

 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/CodeGen/IndirectThunks.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Target/TargetMachine.h"
 #include <cassert>
 #include <climits>
 #include <tuple>

 using namespace llvm;

 #define DEBUG_TYPE "aarch64-sls-hardening"

 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"

 // Common name prefix of all thunks generated by this pass.
 //
 // The generic form is
 // __llvm_slsblr_thunk_xN            for BLR thunks
 // __llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
 // __llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks
 static constexpr StringRef CommonNamePrefix = "__llvm_slsblr_thunk_";

 namespace {

 struct ThunkKind {
   enum ThunkKindId {
     ThunkBR,
     ThunkBRAA,
     ThunkBRAB,
     ThunkBRAAZ,
     ThunkBRABZ,
   };

   ThunkKindId Id;
   StringRef NameInfix;
   bool HasXmOperand;
   bool NeedsPAuth;

   // Opcode to perform indirect jump from inside the thunk.
   unsigned BROpcode;

   static const ThunkKind BR;
   static const ThunkKind BRAA;
   static const ThunkKind BRAB;
   static const ThunkKind BRAAZ;
   static const ThunkKind BRABZ;
 };

 // Set of inserted thunks.
 class ThunksSet {
 public:
   static constexpr unsigned NumXRegisters = 32;

   // Given Xn register, returns n.
   static unsigned indexOfXReg(Register Xn);
   // Given n, returns Xn register.
   static Register xRegByIndex(unsigned N);

   ThunksSet &operator|=(const ThunksSet &Other) {
     BLRThunks |= Other.BLRThunks;
     BLRAAZThunks |= Other.BLRAAZThunks;
     BLRABZThunks |= Other.BLRABZThunks;
     for (unsigned I = 0; I < NumXRegisters; ++I)
       BLRAAThunks[I] |= Other.BLRAAThunks[I];
     for (unsigned I = 0; I < NumXRegisters; ++I)
       BLRABThunks[I] |= Other.BLRABThunks[I];

     return *this;
   }

   bool get(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
     reg_bitmask_t XnBit = reg_bitmask_t(1) << indexOfXReg(Xn);
     return getBitmask(Kind, Xm) & XnBit;
   }

   void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
     reg_bitmask_t XnBit = reg_bitmask_t(1) << indexOfXReg(Xn);
     getBitmask(Kind, Xm) |= XnBit;
   }

 private:
   typedef uint32_t reg_bitmask_t;
   static_assert(NumXRegisters <= sizeof(reg_bitmask_t) * CHAR_BIT,
                 "Bitmask is not wide enough to hold all Xn registers");

   // Bitmasks representing operands used, with n-th bit corresponding to Xn
   // register operand. If the instruction has a second operand (Xm), an array
   // of bitmasks is used, indexed by m.
   // Indexes corresponding to the forbidden x16, x17 and x30 registers are
   // always unset, for simplicity there are no holes.
   reg_bitmask_t BLRThunks = 0;
   reg_bitmask_t BLRAAZThunks = 0;
   reg_bitmask_t BLRABZThunks = 0;
   reg_bitmask_t BLRAAThunks[NumXRegisters] = {};
   reg_bitmask_t BLRABThunks[NumXRegisters] = {};

   reg_bitmask_t &getBitmask(ThunkKind::ThunkKindId Kind, Register Xm) {
     switch (Kind) {
     case ThunkKind::ThunkBR:
       return BLRThunks;
     case ThunkKind::ThunkBRAAZ:
       return BLRAAZThunks;
     case ThunkKind::ThunkBRABZ:
       return BLRABZThunks;
     case ThunkKind::ThunkBRAA:
       return BLRAAThunks[indexOfXReg(Xm)];
     case ThunkKind::ThunkBRAB:
       return BLRABThunks[indexOfXReg(Xm)];
     }
     llvm_unreachable("Unknown ThunkKindId enum");
   }
 };

 struct SLSHardeningInserter : ThunkInserter<SLSHardeningInserter, ThunksSet> {
 public:
   const char *getThunkPrefix() { return CommonNamePrefix.data(); }
   bool mayUseThunk(const MachineFunction &MF) {
     ComdatThunks &= !MF.getSubtarget<AArch64Subtarget>().hardenSlsNoComdat();
     // We are inserting barriers aside from thunk calls, so
     // check hardenSlsRetBr() as well.
     return MF.getSubtarget<AArch64Subtarget>().hardenSlsBlr() ||
            MF.getSubtarget<AArch64Subtarget>().hardenSlsRetBr();
   }
   ThunksSet insertThunks(MachineModuleInfo &MMI, MachineFunction &MF,
                          ThunksSet ExistingThunks);
   void populateThunk(MachineFunction &MF);

 private:
   bool ComdatThunks = true;

   bool hardenReturnsAndBRs(MachineModuleInfo &MMI, MachineBasicBlock &MBB);
   bool hardenBLRs(MachineModuleInfo &MMI, MachineBasicBlock &MBB,
                   ThunksSet &Thunks);

   void convertBLRToBL(MachineModuleInfo &MMI, MachineBasicBlock &MBB,
                       MachineBasicBlock::instr_iterator MBBI,
                       ThunksSet &Thunks);
 };

 } // end anonymous namespace

 const ThunkKind ThunkKind::BR = {ThunkBR, "", /*HasXmOperand=*/false,
                                  /*NeedsPAuth=*/false, AArch64::BR};
 const ThunkKind ThunkKind::BRAA = {ThunkBRAA, "aa_", /*HasXmOperand=*/true,
                                    /*NeedsPAuth=*/true, AArch64::BRAA};
 const ThunkKind ThunkKind::BRAB = {ThunkBRAB, "ab_", /*HasXmOperand=*/true,
                                    /*NeedsPAuth=*/true, AArch64::BRAB};
 const ThunkKind ThunkKind::BRAAZ = {ThunkBRAAZ, "aaz_", /*HasXmOperand=*/false,
                                     /*NeedsPAuth=*/true, AArch64::BRAAZ};
 const ThunkKind ThunkKind::BRABZ = {ThunkBRABZ, "abz_", /*HasXmOperand=*/false,
                                     /*NeedsPAuth=*/true, AArch64::BRABZ};

 // Returns thunk kind to emit, or nullptr if not a BLR* instruction.
 static const ThunkKind *getThunkKind(unsigned OriginalOpcode) {
   switch (OriginalOpcode) {
   case AArch64::BLR:
   case AArch64::BLRNoIP:
     return &ThunkKind::BR;
   case AArch64::BLRAA:
     return &ThunkKind::BRAA;
   case AArch64::BLRAB:
     return &ThunkKind::BRAB;
   case AArch64::BLRAAZ:
     return &ThunkKind::BRAAZ;
   case AArch64::BLRABZ:
     return &ThunkKind::BRABZ;
   }
   return nullptr;
 }

 static bool isBLR(const MachineInstr &MI) {
   return getThunkKind(MI.getOpcode()) != nullptr;
 }

 unsigned ThunksSet::indexOfXReg(Register Reg) {
   assert(AArch64::GPR64RegClass.contains(Reg));
   assert(Reg != AArch64::X16 && Reg != AArch64::X17 && Reg != AArch64::LR);

   // Most Xn registers have consecutive ids, except for FP and XZR.
   unsigned Result = (unsigned)Reg - (unsigned)AArch64::X0;
   if (Reg == AArch64::FP)
     Result = 29;
   else if (Reg == AArch64::XZR)
     Result = 31;

   assert(Result < NumXRegisters && "Internal register numbering changed");
   assert(AArch64::GPR64RegClass.getRegister(Result).id() == Reg &&
          "Internal register numbering changed");

   return Result;
 }

 Register ThunksSet::xRegByIndex(unsigned N) {
   return AArch64::GPR64RegClass.getRegister(N);
 }

 static void insertSpeculationBarrier(const AArch64Subtarget *ST,
                                      MachineBasicBlock &MBB,
                                      MachineBasicBlock::iterator MBBI,
                                      DebugLoc DL,
                                      bool AlwaysUseISBDSB = false) {
   assert(MBBI != MBB.begin() &&
          "Must not insert SpeculationBarrierEndBB as only instruction in MBB.");
   assert(std::prev(MBBI)->isBarrier() &&
          "SpeculationBarrierEndBB must only follow unconditional control flow "
          "instructions.");
   assert(std::prev(MBBI)->isTerminator() &&
          "SpeculationBarrierEndBB must only follow terminators.");
   const TargetInstrInfo *TII = ST->getInstrInfo();
   unsigned BarrierOpc = ST->hasSB() && !AlwaysUseISBDSB
                             ? AArch64::SpeculationBarrierSBEndBB
                             : AArch64::SpeculationBarrierISBDSBEndBB;
   if (MBBI == MBB.end() ||
       (MBBI->getOpcode() != AArch64::SpeculationBarrierSBEndBB &&
        MBBI->getOpcode() != AArch64::SpeculationBarrierISBDSBEndBB))
     BuildMI(MBB, MBBI, DL, TII->get(BarrierOpc));
 }

 ThunksSet SLSHardeningInserter::insertThunks(MachineModuleInfo &MMI,
                                              MachineFunction &MF,
                                              ThunksSet ExistingThunks) {
   const AArch64Subtarget *ST = &MF.getSubtarget<AArch64Subtarget>();

   for (auto &MBB : MF) {
     if (ST->hardenSlsRetBr())
       hardenReturnsAndBRs(MMI, MBB);
     if (ST->hardenSlsBlr())
       hardenBLRs(MMI, MBB, ExistingThunks);
   }
   return ExistingThunks;
 }

 bool SLSHardeningInserter::hardenReturnsAndBRs(MachineModuleInfo &MMI,
                                                MachineBasicBlock &MBB) {
   const AArch64Subtarget *ST =
       &MBB.getParent()->getSubtarget<AArch64Subtarget>();
   bool Modified = false;
   MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator(), E = MBB.end();
   MachineBasicBlock::iterator NextMBBI;
   for (; MBBI != E; MBBI = NextMBBI) {
     MachineInstr &MI = *MBBI;
     NextMBBI = std::next(MBBI);
     if (MI.isReturn() || isIndirectBranchOpcode(MI.getOpcode())) {
       assert(MI.isTerminator());
       insertSpeculationBarrier(ST, MBB, std::next(MBBI), MI.getDebugLoc());
       Modified = true;
     }
   }
   return Modified;
 }

 // Currently, the longest possible thunk name is
 //   __llvm_slsblr_thunk_aa_xNN_xMM
 // which is 31 characters (without the '\0' character).
 static SmallString<32> createThunkName(const ThunkKind &Kind, Register Xn,
                                        Register Xm) {
   unsigned N = ThunksSet::indexOfXReg(Xn);
   if (!Kind.HasXmOperand)
     return formatv("{0}{1}x{2}", CommonNamePrefix, Kind.NameInfix, N);

   unsigned M = ThunksSet::indexOfXReg(Xm);
   return formatv("{0}{1}x{2}_x{3}", CommonNamePrefix, Kind.NameInfix, N, M);
 }

 static std::tuple<const ThunkKind &, Register, Register>
 parseThunkName(StringRef ThunkName) {
   assert(ThunkName.starts_with(CommonNamePrefix) &&
          "Should be filtered out by ThunkInserter");
   // Thunk name suffix, such as "x1" or "aa_x2_x3".
   StringRef NameSuffix = ThunkName.drop_front(CommonNamePrefix.size());

   // Parse thunk kind based on thunk name infix.
   const ThunkKind &Kind = *StringSwitch<const ThunkKind *>(NameSuffix)
                                .StartsWith("aa_", &ThunkKind::BRAA)
                                .StartsWith("ab_", &ThunkKind::BRAB)
                                .StartsWith("aaz_", &ThunkKind::BRAAZ)
                                .StartsWith("abz_", &ThunkKind::BRABZ)
                                .Default(&ThunkKind::BR);

   auto ParseRegName = [](StringRef Name) {
     unsigned N;

     assert(Name.starts_with("x") && "xN register name expected");
     bool Fail = Name.drop_front(1).getAsInteger(/*Radix=*/10, N);
     assert(!Fail && N < ThunksSet::NumXRegisters && "Unexpected register");
     (void)Fail;

     return ThunksSet::xRegByIndex(N);
   };

   // For example, "x1" or "x2_x3".
   StringRef RegsStr = NameSuffix.drop_front(Kind.NameInfix.size());
   StringRef XnStr, XmStr;
   std::tie(XnStr, XmStr) = RegsStr.split('_');

   // Parse register operands.
   Register Xn = ParseRegName(XnStr);
   Register Xm = Kind.HasXmOperand ? ParseRegName(XmStr) : AArch64::NoRegister;

   return std::make_tuple(std::ref(Kind), Xn, Xm);
 }

 void SLSHardeningInserter::populateThunk(MachineFunction &MF) {
   assert(MF.getFunction().hasComdat() == ComdatThunks &&
          "ComdatThunks value changed since MF creation");
   Register Xn, Xm;
   auto KindAndRegs = parseThunkName(MF.getName());
   const ThunkKind &Kind = std::get<0>(KindAndRegs);
   std::tie(std::ignore, Xn, Xm) = KindAndRegs;

   const TargetInstrInfo *TII =
       MF.getSubtarget<AArch64Subtarget>().getInstrInfo();

   // Depending on whether this pass is in the same FunctionPassManager as the
   // IR->MIR conversion, the thunk may be completely empty, or contain a single
   // basic block with a single return instruction. Normalise it to contain a
   // single empty basic block.
   if (MF.size() == 1) {
     assert(MF.front().size() == 1);
     assert(MF.front().front().getOpcode() == AArch64::RET);
     MF.front().erase(MF.front().begin());
   } else {
     assert(MF.size() == 0);
     MF.push_back(MF.CreateMachineBasicBlock());
   }

   MachineBasicBlock *Entry = &MF.front();
   Entry->clear();

   //  These thunks need to consist of the following instructions:
   //  __llvm_slsblr_thunk_...:
   //      MOV x16, xN     ; BR* instructions are not compatible with "BTI c"
   //                      ; branch target unless xN is x16 or x17.
   //      BR* ...         ; One of: BR        x16
   //                      ;         BRA(A|B)  x16, xM
   //                      ;         BRA(A|B)Z x16
   //      barrierInsts
   Entry->addLiveIn(Xn);
   // MOV X16, Reg == ORR X16, XZR, Reg, LSL #0
   BuildMI(Entry, DebugLoc(), TII->get(AArch64::ORRXrs), AArch64::X16)
       .addReg(AArch64::XZR)
       .addReg(Xn)
       .addImm(0);
   MachineInstrBuilder Builder =
       BuildMI(Entry, DebugLoc(), TII->get(Kind.BROpcode)).addReg(AArch64::X16);
   if (Xm != AArch64::NoRegister) {
     Entry->addLiveIn(Xm);
     Builder.addReg(Xm);
   }

   // Make sure the thunks do not make use of the SB extension in case there is
   // a function somewhere that will call to it that for some reason disabled
   // the SB extension locally on that function, even though it's enabled for
   // the module otherwise. Therefore set AlwaysUseISBSDB to true.
   insertSpeculationBarrier(&MF.getSubtarget<AArch64Subtarget>(), *Entry,
                            Entry->end(), DebugLoc(), true /*AlwaysUseISBDSB*/);
 }

 void SLSHardeningInserter::convertBLRToBL(
     MachineModuleInfo &MMI, MachineBasicBlock &MBB,
     MachineBasicBlock::instr_iterator MBBI, ThunksSet &Thunks) {
   // Transform a BLR* instruction (one of BLR, BLRAA/BLRAB or BLRAAZ/BLRABZ) to
   // a BL to the thunk containing BR, BRAA/BRAB or BRAAZ/BRABZ, respectively.
   //
   // Before:
   //   |-----------------------------|
   //   |      ...                    |
   //   |  instI                      |
   //   |  BLR* xN or BLR* xN, xM     |
   //   |  instJ                      |
   //   |      ...                    |
   //   |-----------------------------|
   //
   // After:
   //   |-----------------------------|
   //   |      ...                    |
   //   |  instI                      |
   //   |  BL __llvm_slsblr_thunk_... |
   //   |  instJ                      |
   //   |      ...                    |
   //   |-----------------------------|
   //
   //   __llvm_slsblr_thunk_...:
   //   |-----------------------------|
   //   |  MOV x16, xN                |
   //   |  BR* x16 or BR* x16, xM     |
   //   |  barrierInsts               |
   //   |-----------------------------|
   //
   // This function needs to transform BLR* instruction into BL with the correct
   // thunk name and lazily create the thunk if it does not exist yet.
   //
   // Since linkers are allowed to clobber X16 and X17 on function calls, the
   // above mitigation only works if the original BLR* instruction had neither
   // X16 nor X17 as one of its operands. Code generation before must make sure
   // that no such BLR* instruction was produced if the mitigation is enabled.

   MachineInstr &BLR = *MBBI;
   assert(isBLR(BLR));
   const ThunkKind &Kind = *getThunkKind(BLR.getOpcode());

   unsigned NumRegOperands = Kind.HasXmOperand ? 2 : 1;
   assert(BLR.getNumExplicitOperands() == NumRegOperands &&
          "Expected one or two register inputs");
   Register Xn = BLR.getOperand(0).getReg();
   Register Xm =
       Kind.HasXmOperand ? BLR.getOperand(1).getReg() : AArch64::NoRegister;

   DebugLoc DL = BLR.getDebugLoc();

   MachineFunction &MF = *MBBI->getMF();
   MCContext &Context = MBB.getParent()->getContext();
   const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();

   auto ThunkName = createThunkName(Kind, Xn, Xm);
   MCSymbol *Sym = Context.getOrCreateSymbol(ThunkName);

   if (!Thunks.get(Kind.Id, Xn, Xm)) {
     StringRef TargetAttrs = Kind.NeedsPAuth ? "+pauth" : "";
     Thunks.set(Kind.Id, Xn, Xm);
     createThunkFunction(MMI, ThunkName, ComdatThunks, TargetAttrs);
   }

   MachineInstr *BL = BuildMI(MBB, MBBI, DL, TII->get(AArch64::BL)).addSym(Sym);

   // Now copy the implicit operands from BLR to BL and copy other necessary
   // info.
   // However, both BLR and BL instructions implictly use SP and implicitly
   // define LR. Blindly copying implicit operands would result in SP and LR
   // operands to be present multiple times. While this may not be too much of
   // an issue, let's avoid that for cleanliness, by removing those implicit
   // operands from the BL created above before we copy over all implicit
   // operands from the BLR.
   int ImpLROpIdx = -1;
   int ImpSPOpIdx = -1;
   for (unsigned OpIdx = BL->getNumExplicitOperands();
        OpIdx < BL->getNumOperands(); OpIdx++) {
     MachineOperand Op = BL->getOperand(OpIdx);
     if (!Op.isReg())
       continue;
     if (Op.getReg() == AArch64::LR && Op.isDef())
       ImpLROpIdx = OpIdx;
     if (Op.getReg() == AArch64::SP && !Op.isDef())
       ImpSPOpIdx = OpIdx;
   }
   assert(ImpLROpIdx != -1);
   assert(ImpSPOpIdx != -1);
   int FirstOpIdxToRemove = std::max(ImpLROpIdx, ImpSPOpIdx);
   int SecondOpIdxToRemove = std::min(ImpLROpIdx, ImpSPOpIdx);
   BL->removeOperand(FirstOpIdxToRemove);
   BL->removeOperand(SecondOpIdxToRemove);
   // Now copy over the implicit operands from the original BLR
   BL->copyImplicitOps(MF, BLR);
   MF.moveAdditionalCallInfo(&BLR, BL);
   // Also add the register operands of the original BLR* instruction
   // as being used in the called thunk.
   for (unsigned OpIdx = 0; OpIdx < NumRegOperands; ++OpIdx) {
     MachineOperand &Op = BLR.getOperand(OpIdx);
     BL->addOperand(MachineOperand::CreateReg(Op.getReg(), /*isDef=*/false,
                                              /*isImp=*/true, Op.isKill()));
   }
   // Remove BLR instruction
   MBB.erase(MBBI);
 }

 bool SLSHardeningInserter::hardenBLRs(MachineModuleInfo &MMI,
                                       MachineBasicBlock &MBB,
                                       ThunksSet &Thunks) {
   bool Modified = false;
   MachineBasicBlock::instr_iterator MBBI = MBB.instr_begin(),
                                     E = MBB.instr_end();
   MachineBasicBlock::instr_iterator NextMBBI;
   for (; MBBI != E; MBBI = NextMBBI) {
     MachineInstr &MI = *MBBI;
     NextMBBI = std::next(MBBI);
     if (isBLR(MI)) {
       convertBLRToBL(MMI, MBB, MBBI, Thunks);
       Modified = true;
     }
   }
   return Modified;
 }

 namespace {
 class AArch64SLSHardening : public ThunkInserterPass<SLSHardeningInserter> {
 public:
   static char ID;

   AArch64SLSHardening() : ThunkInserterPass(ID) {}

   StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }
 };

 } // end anonymous namespace

 char AArch64SLSHardening::ID = 0;

 INITIALIZE_PASS(AArch64SLSHardening, "aarch64-sls-hardening",
                 AARCH64_SLS_HARDENING_NAME, false, false)

 FunctionPass *llvm::createAArch64SLSHardeningPass() {
   return new AArch64SLSHardening();
 }
	//===- AArch64SLSHardening.cpp - Harden Straight Line Missspeculation -----===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//
	//
	// This file contains a pass to insert code to mitigate against side channel
	// vulnerabilities that may happen under straight line miss-speculation.
	//
	//===----------------------------------------------------------------------===//

	#include "AArch64InstrInfo.h"
	#include "AArch64Subtarget.h"
	#include "llvm/ADT/StringSwitch.h"
	#include "llvm/CodeGen/IndirectThunks.h"
	#include "llvm/CodeGen/MachineBasicBlock.h"
	#include "llvm/CodeGen/MachineFunction.h"
	#include "llvm/CodeGen/MachineInstr.h"
	#include "llvm/CodeGen/MachineInstrBuilder.h"
	#include "llvm/CodeGen/MachineOperand.h"
	#include "llvm/CodeGen/RegisterScavenging.h"
	#include "llvm/IR/DebugLoc.h"
	#include "llvm/Pass.h"
	#include "llvm/Support/ErrorHandling.h"
	#include "llvm/Support/FormatVariadic.h"
	#include "llvm/Target/TargetMachine.h"
	#include <cassert>
	#include <climits>
	#include <tuple>

	using namespace llvm;

	#define DEBUG_TYPE "aarch64-sls-hardening"

	#define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"

	// Common name prefix of all thunks generated by this pass.
	//
	// The generic form is
	// __llvm_slsblr_thunk_xN for BLR thunks
	// __llvm_slsblr_thunk_(aaz\|abz)_xN for BLRAAZ and BLRABZ thunks
	// __llvm_slsblr_thunk_(aa\|ab)_xN_xM for BLRAA and BLRAB thunks
	static constexpr StringRef CommonNamePrefix = "__llvm_slsblr_thunk_";

	namespace {

	struct ThunkKind {
	enum ThunkKindId {
	ThunkBR,
	ThunkBRAA,
	ThunkBRAB,
	ThunkBRAAZ,
	ThunkBRABZ,
	};

	ThunkKindId Id;
	StringRef NameInfix;
	bool HasXmOperand;
	bool NeedsPAuth;

	// Opcode to perform indirect jump from inside the thunk.
	unsigned BROpcode;

	static const ThunkKind BR;
	static const ThunkKind BRAA;
	static const ThunkKind BRAB;
	static const ThunkKind BRAAZ;
	static const ThunkKind BRABZ;
	};

	// Set of inserted thunks.
	class ThunksSet {
	public:
	static constexpr unsigned NumXRegisters = 32;

	// Given Xn register, returns n.
	static unsigned indexOfXReg(Register Xn);
	// Given n, returns Xn register.
	static Register xRegByIndex(unsigned N);

	ThunksSet &operator\|=(const ThunksSet &Other) {
	BLRThunks \|= Other.BLRThunks;
	BLRAAZThunks \|= Other.BLRAAZThunks;
	BLRABZThunks \|= Other.BLRABZThunks;
	for (unsigned I = 0; I < NumXRegisters; ++I)
	BLRAAThunks[I] \|= Other.BLRAAThunks[I];
	for (unsigned I = 0; I < NumXRegisters; ++I)
	BLRABThunks[I] \|= Other.BLRABThunks[I];

	return *this;
	}

	bool get(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
	reg_bitmask_t XnBit = reg_bitmask_t(1) << indexOfXReg(Xn);
	return getBitmask(Kind, Xm) & XnBit;
	}

	void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
	reg_bitmask_t XnBit = reg_bitmask_t(1) << indexOfXReg(Xn);
	getBitmask(Kind, Xm) \|= XnBit;
	}

	private:
	typedef uint32_t reg_bitmask_t;
	static_assert(NumXRegisters <= sizeof(reg_bitmask_t) * CHAR_BIT,
	"Bitmask is not wide enough to hold all Xn registers");

	// Bitmasks representing operands used, with n-th bit corresponding to Xn
	// register operand. If the instruction has a second operand (Xm), an array
	// of bitmasks is used, indexed by m.
	// Indexes corresponding to the forbidden x16, x17 and x30 registers are
	// always unset, for simplicity there are no holes.
	reg_bitmask_t BLRThunks = 0;
	reg_bitmask_t BLRAAZThunks = 0;
	reg_bitmask_t BLRABZThunks = 0;
	reg_bitmask_t BLRAAThunks[NumXRegisters] = {};
	reg_bitmask_t BLRABThunks[NumXRegisters] = {};

	reg_bitmask_t &getBitmask(ThunkKind::ThunkKindId Kind, Register Xm) {
	switch (Kind) {
	case ThunkKind::ThunkBR:
	return BLRThunks;
	case ThunkKind::ThunkBRAAZ:
	return BLRAAZThunks;
	case ThunkKind::ThunkBRABZ:
	return BLRABZThunks;
	case ThunkKind::ThunkBRAA:
	return BLRAAThunks[indexOfXReg(Xm)];
	case ThunkKind::ThunkBRAB:
	return BLRABThunks[indexOfXReg(Xm)];
	}
	llvm_unreachable("Unknown ThunkKindId enum");
	}
	};

	struct SLSHardeningInserter : ThunkInserter<SLSHardeningInserter, ThunksSet> {
	public:
	const char *getThunkPrefix() { return CommonNamePrefix.data(); }
	bool mayUseThunk(const MachineFunction &MF) {
	ComdatThunks &= !MF.getSubtarget<AArch64Subtarget>().hardenSlsNoComdat();
	// We are inserting barriers aside from thunk calls, so
	// check hardenSlsRetBr() as well.
	return MF.getSubtarget<AArch64Subtarget>().hardenSlsBlr() \|\|
	MF.getSubtarget<AArch64Subtarget>().hardenSlsRetBr();
	}
	ThunksSet insertThunks(MachineModuleInfo &MMI, MachineFunction &MF,
	ThunksSet ExistingThunks);
	void populateThunk(MachineFunction &MF);

	private:
	bool ComdatThunks = true;

	bool hardenReturnsAndBRs(MachineModuleInfo &MMI, MachineBasicBlock &MBB);
	bool hardenBLRs(MachineModuleInfo &MMI, MachineBasicBlock &MBB,
	ThunksSet &Thunks);

	void convertBLRToBL(MachineModuleInfo &MMI, MachineBasicBlock &MBB,
	MachineBasicBlock::instr_iterator MBBI,
	ThunksSet &Thunks);
	};

	} // end anonymous namespace

	const ThunkKind ThunkKind::BR = {ThunkBR, "", /HasXmOperand=/false,
	/NeedsPAuth=/false, AArch64::BR};
	const ThunkKind ThunkKind::BRAA = {ThunkBRAA, "aa_", /HasXmOperand=/true,
	/NeedsPAuth=/true, AArch64::BRAA};
	const ThunkKind ThunkKind::BRAB = {ThunkBRAB, "ab_", /HasXmOperand=/true,
	/NeedsPAuth=/true, AArch64::BRAB};
	const ThunkKind ThunkKind::BRAAZ = {ThunkBRAAZ, "aaz_", /HasXmOperand=/false,
	/NeedsPAuth=/true, AArch64::BRAAZ};
	const ThunkKind ThunkKind::BRABZ = {ThunkBRABZ, "abz_", /HasXmOperand=/false,
	/NeedsPAuth=/true, AArch64::BRABZ};

	// Returns thunk kind to emit, or nullptr if not a BLR* instruction.
	static const ThunkKind *getThunkKind(unsigned OriginalOpcode) {
	switch (OriginalOpcode) {
	case AArch64::BLR:
	case AArch64::BLRNoIP:
	return &ThunkKind::BR;
	case AArch64::BLRAA:
	return &ThunkKind::BRAA;
	case AArch64::BLRAB:
	return &ThunkKind::BRAB;
	case AArch64::BLRAAZ:
	return &ThunkKind::BRAAZ;
	case AArch64::BLRABZ:
	return &ThunkKind::BRABZ;
	}
	return nullptr;
	}

	static bool isBLR(const MachineInstr &MI) {
	return getThunkKind(MI.getOpcode()) != nullptr;
	}

	unsigned ThunksSet::indexOfXReg(Register Reg) {
	assert(AArch64::GPR64RegClass.contains(Reg));
	assert(Reg != AArch64::X16 && Reg != AArch64::X17 && Reg != AArch64::LR);

	// Most Xn registers have consecutive ids, except for FP and XZR.
	unsigned Result = (unsigned)Reg - (unsigned)AArch64::X0;
	if (Reg == AArch64::FP)
	Result = 29;
	else if (Reg == AArch64::XZR)
	Result = 31;

	assert(Result < NumXRegisters && "Internal register numbering changed");
	assert(AArch64::GPR64RegClass.getRegister(Result).id() == Reg &&
	"Internal register numbering changed");

	return Result;
	}

	Register ThunksSet::xRegByIndex(unsigned N) {
	return AArch64::GPR64RegClass.getRegister(N);
	}

	static void insertSpeculationBarrier(const AArch64Subtarget *ST,
	MachineBasicBlock &MBB,
	MachineBasicBlock::iterator MBBI,
	DebugLoc DL,
	bool AlwaysUseISBDSB = false) {
	assert(MBBI != MBB.begin() &&
	"Must not insert SpeculationBarrierEndBB as only instruction in MBB.");
	assert(std::prev(MBBI)->isBarrier() &&
	"SpeculationBarrierEndBB must only follow unconditional control flow "
	"instructions.");
	assert(std::prev(MBBI)->isTerminator() &&
	"SpeculationBarrierEndBB must only follow terminators.");
	const TargetInstrInfo *TII = ST->getInstrInfo();
	unsigned BarrierOpc = ST->hasSB() && !AlwaysUseISBDSB
	? AArch64::SpeculationBarrierSBEndBB
	: AArch64::SpeculationBarrierISBDSBEndBB;
	if (MBBI == MBB.end() \|\|
	(MBBI->getOpcode() != AArch64::SpeculationBarrierSBEndBB &&
	MBBI->getOpcode() != AArch64::SpeculationBarrierISBDSBEndBB))
	BuildMI(MBB, MBBI, DL, TII->get(BarrierOpc));
	}

	ThunksSet SLSHardeningInserter::insertThunks(MachineModuleInfo &MMI,
	MachineFunction &MF,
	ThunksSet ExistingThunks) {
	const AArch64Subtarget *ST = &MF.getSubtarget<AArch64Subtarget>();

	for (auto &MBB : MF) {
	if (ST->hardenSlsRetBr())
	hardenReturnsAndBRs(MMI, MBB);
	if (ST->hardenSlsBlr())
	hardenBLRs(MMI, MBB, ExistingThunks);
	}
	return ExistingThunks;
	}

	bool SLSHardeningInserter::hardenReturnsAndBRs(MachineModuleInfo &MMI,
	MachineBasicBlock &MBB) {
	const AArch64Subtarget *ST =
	&MBB.getParent()->getSubtarget<AArch64Subtarget>();
	bool Modified = false;
	MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator(), E = MBB.end();
	MachineBasicBlock::iterator NextMBBI;
	for (; MBBI != E; MBBI = NextMBBI) {
	MachineInstr &MI = *MBBI;
	NextMBBI = std::next(MBBI);
	if (MI.isReturn() \|\| isIndirectBranchOpcode(MI.getOpcode())) {
	assert(MI.isTerminator());
	insertSpeculationBarrier(ST, MBB, std::next(MBBI), MI.getDebugLoc());
	Modified = true;
	}
	}
	return Modified;
	}

	// Currently, the longest possible thunk name is
	// __llvm_slsblr_thunk_aa_xNN_xMM
	// which is 31 characters (without the '\0' character).
	static SmallString<32> createThunkName(const ThunkKind &Kind, Register Xn,
	Register Xm) {
	unsigned N = ThunksSet::indexOfXReg(Xn);
	if (!Kind.HasXmOperand)
	return formatv("{0}{1}x{2}", CommonNamePrefix, Kind.NameInfix, N);

	unsigned M = ThunksSet::indexOfXReg(Xm);
	return formatv("{0}{1}x{2}_x{3}", CommonNamePrefix, Kind.NameInfix, N, M);
	}

	static std::tuple<const ThunkKind &, Register, Register>
	parseThunkName(StringRef ThunkName) {
	assert(ThunkName.starts_with(CommonNamePrefix) &&
	"Should be filtered out by ThunkInserter");
	// Thunk name suffix, such as "x1" or "aa_x2_x3".
	StringRef NameSuffix = ThunkName.drop_front(CommonNamePrefix.size());

	// Parse thunk kind based on thunk name infix.
	const ThunkKind &Kind = StringSwitch<const ThunkKind >(NameSuffix)
	.StartsWith("aa_", &ThunkKind::BRAA)
	.StartsWith("ab_", &ThunkKind::BRAB)
	.StartsWith("aaz_", &ThunkKind::BRAAZ)
	.StartsWith("abz_", &ThunkKind::BRABZ)
	.Default(&ThunkKind::BR);

	auto ParseRegName = [](StringRef Name) {
	unsigned N;

	assert(Name.starts_with("x") && "xN register name expected");
	bool Fail = Name.drop_front(1).getAsInteger(/Radix=/10, N);
	assert(!Fail && N < ThunksSet::NumXRegisters && "Unexpected register");
	(void)Fail;

	return ThunksSet::xRegByIndex(N);
	};

	// For example, "x1" or "x2_x3".
	StringRef RegsStr = NameSuffix.drop_front(Kind.NameInfix.size());
	StringRef XnStr, XmStr;
	std::tie(XnStr, XmStr) = RegsStr.split('_');

	// Parse register operands.
	Register Xn = ParseRegName(XnStr);
	Register Xm = Kind.HasXmOperand ? ParseRegName(XmStr) : AArch64::NoRegister;

	return std::make_tuple(std::ref(Kind), Xn, Xm);
	}

	void SLSHardeningInserter::populateThunk(MachineFunction &MF) {
	assert(MF.getFunction().hasComdat() == ComdatThunks &&
	"ComdatThunks value changed since MF creation");
	Register Xn, Xm;
	auto KindAndRegs = parseThunkName(MF.getName());
	const ThunkKind &Kind = std::get<0>(KindAndRegs);
	std::tie(std::ignore, Xn, Xm) = KindAndRegs;

	const TargetInstrInfo *TII =
	MF.getSubtarget<AArch64Subtarget>().getInstrInfo();

	// Depending on whether this pass is in the same FunctionPassManager as the
	// IR->MIR conversion, the thunk may be completely empty, or contain a single
	// basic block with a single return instruction. Normalise it to contain a
	// single empty basic block.
	if (MF.size() == 1) {
	assert(MF.front().size() == 1);
	assert(MF.front().front().getOpcode() == AArch64::RET);
	MF.front().erase(MF.front().begin());
	} else {
	assert(MF.size() == 0);
	MF.push_back(MF.CreateMachineBasicBlock());
	}

	MachineBasicBlock *Entry = &MF.front();
	Entry->clear();

	// These thunks need to consist of the following instructions:
	// __llvm_slsblr_thunk_...:
	// MOV x16, xN ; BR* instructions are not compatible with "BTI c"
	// ; branch target unless xN is x16 or x17.
	// BR* ... ; One of: BR x16
	// ; BRA(A\|B) x16, xM
	// ; BRA(A\|B)Z x16
	// barrierInsts
	Entry->addLiveIn(Xn);
	// MOV X16, Reg == ORR X16, XZR, Reg, LSL #0
	BuildMI(Entry, DebugLoc(), TII->get(AArch64::ORRXrs), AArch64::X16)
	.addReg(AArch64::XZR)
	.addReg(Xn)
	.addImm(0);
	MachineInstrBuilder Builder =
	BuildMI(Entry, DebugLoc(), TII->get(Kind.BROpcode)).addReg(AArch64::X16);
	if (Xm != AArch64::NoRegister) {
	Entry->addLiveIn(Xm);
	Builder.addReg(Xm);
	}

	// Make sure the thunks do not make use of the SB extension in case there is
	// a function somewhere that will call to it that for some reason disabled
	// the SB extension locally on that function, even though it's enabled for
	// the module otherwise. Therefore set AlwaysUseISBSDB to true.
	insertSpeculationBarrier(&MF.getSubtarget<AArch64Subtarget>(), *Entry,
	Entry->end(), DebugLoc(), true /AlwaysUseISBDSB/);
	}

	void SLSHardeningInserter::convertBLRToBL(
	MachineModuleInfo &MMI, MachineBasicBlock &MBB,
	MachineBasicBlock::instr_iterator MBBI, ThunksSet &Thunks) {
	// Transform a BLR* instruction (one of BLR, BLRAA/BLRAB or BLRAAZ/BLRABZ) to
	// a BL to the thunk containing BR, BRAA/BRAB or BRAAZ/BRABZ, respectively.
	//
	// Before:
	// \|-----------------------------\|
	// \| ... \|
	// \| instI \|
	// \| BLR* xN or BLR* xN, xM \|
	// \| instJ \|
	// \| ... \|
	// \|-----------------------------\|
	//
	// After:
	// \|-----------------------------\|
	// \| ... \|
	// \| instI \|
	// \| BL __llvm_slsblr_thunk_... \|
	// \| instJ \|
	// \| ... \|
	// \|-----------------------------\|
	//
	// __llvm_slsblr_thunk_...:
	// \|-----------------------------\|
	// \| MOV x16, xN \|
	// \| BR* x16 or BR* x16, xM \|
	// \| barrierInsts \|
	// \|-----------------------------\|
	//
	// This function needs to transform BLR* instruction into BL with the correct
	// thunk name and lazily create the thunk if it does not exist yet.
	//
	// Since linkers are allowed to clobber X16 and X17 on function calls, the
	// above mitigation only works if the original BLR* instruction had neither
	// X16 nor X17 as one of its operands. Code generation before must make sure
	// that no such BLR* instruction was produced if the mitigation is enabled.

	MachineInstr &BLR = *MBBI;
	assert(isBLR(BLR));
	const ThunkKind &Kind = *getThunkKind(BLR.getOpcode());

	unsigned NumRegOperands = Kind.HasXmOperand ? 2 : 1;
	assert(BLR.getNumExplicitOperands() == NumRegOperands &&
	"Expected one or two register inputs");
	Register Xn = BLR.getOperand(0).getReg();
	Register Xm =
	Kind.HasXmOperand ? BLR.getOperand(1).getReg() : AArch64::NoRegister;

	DebugLoc DL = BLR.getDebugLoc();

	MachineFunction &MF = *MBBI->getMF();
	MCContext &Context = MBB.getParent()->getContext();
	const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();

	auto ThunkName = createThunkName(Kind, Xn, Xm);
	MCSymbol *Sym = Context.getOrCreateSymbol(ThunkName);

	if (!Thunks.get(Kind.Id, Xn, Xm)) {
	StringRef TargetAttrs = Kind.NeedsPAuth ? "+pauth" : "";
	Thunks.set(Kind.Id, Xn, Xm);
	createThunkFunction(MMI, ThunkName, ComdatThunks, TargetAttrs);
	}

	MachineInstr *BL = BuildMI(MBB, MBBI, DL, TII->get(AArch64::BL)).addSym(Sym);

	// Now copy the implicit operands from BLR to BL and copy other necessary
	// info.
	// However, both BLR and BL instructions implictly use SP and implicitly
	// define LR. Blindly copying implicit operands would result in SP and LR
	// operands to be present multiple times. While this may not be too much of
	// an issue, let's avoid that for cleanliness, by removing those implicit
	// operands from the BL created above before we copy over all implicit
	// operands from the BLR.
	int ImpLROpIdx = -1;
	int ImpSPOpIdx = -1;
	for (unsigned OpIdx = BL->getNumExplicitOperands();
	OpIdx < BL->getNumOperands(); OpIdx++) {
	MachineOperand Op = BL->getOperand(OpIdx);
	if (!Op.isReg())
	continue;
	if (Op.getReg() == AArch64::LR && Op.isDef())
	ImpLROpIdx = OpIdx;
	if (Op.getReg() == AArch64::SP && !Op.isDef())
	ImpSPOpIdx = OpIdx;
	}
	assert(ImpLROpIdx != -1);
	assert(ImpSPOpIdx != -1);
	int FirstOpIdxToRemove = std::max(ImpLROpIdx, ImpSPOpIdx);
	int SecondOpIdxToRemove = std::min(ImpLROpIdx, ImpSPOpIdx);
	BL->removeOperand(FirstOpIdxToRemove);
	BL->removeOperand(SecondOpIdxToRemove);
	// Now copy over the implicit operands from the original BLR
	BL->copyImplicitOps(MF, BLR);
	MF.moveAdditionalCallInfo(&BLR, BL);
	// Also add the register operands of the original BLR* instruction
	// as being used in the called thunk.
	for (unsigned OpIdx = 0; OpIdx < NumRegOperands; ++OpIdx) {
	MachineOperand &Op = BLR.getOperand(OpIdx);
	BL->addOperand(MachineOperand::CreateReg(Op.getReg(), /isDef=/false,
	/isImp=/true, Op.isKill()));
	}
	// Remove BLR instruction
	MBB.erase(MBBI);
	}

	bool SLSHardeningInserter::hardenBLRs(MachineModuleInfo &MMI,
	MachineBasicBlock &MBB,
	ThunksSet &Thunks) {
	bool Modified = false;
	MachineBasicBlock::instr_iterator MBBI = MBB.instr_begin(),
	E = MBB.instr_end();
	MachineBasicBlock::instr_iterator NextMBBI;
	for (; MBBI != E; MBBI = NextMBBI) {
	MachineInstr &MI = *MBBI;
	NextMBBI = std::next(MBBI);
	if (isBLR(MI)) {
	convertBLRToBL(MMI, MBB, MBBI, Thunks);
	Modified = true;
	}
	}
	return Modified;
	}

	namespace {
	class AArch64SLSHardening : public ThunkInserterPass<SLSHardeningInserter> {
	public:
	static char ID;

	AArch64SLSHardening() : ThunkInserterPass(ID) {}

	StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }
	};

	} // end anonymous namespace

	char AArch64SLSHardening::ID = 0;

	INITIALIZE_PASS(AArch64SLSHardening, "aarch64-sls-hardening",
	AARCH64_SLS_HARDENING_NAME, false, false)

	FunctionPass *llvm::createAArch64SLSHardeningPass() {
	return new AArch64SLSHardening();
	}