Skip to content

[AArch64] When hardening against SLS, only create called thunks #97472

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 5, 2024
Merged

[AArch64] When hardening against SLS, only create called thunks #97472

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion llvm/lib/Target/AArch64/AArch64.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ FunctionPass *createAArch64ISelDag(AArch64TargetMachine &TM,
FunctionPass *createAArch64StorePairSuppressPass();
FunctionPass *createAArch64ExpandPseudoPass();
FunctionPass *createAArch64SLSHardeningPass();
FunctionPass *createAArch64IndirectThunks();
FunctionPass *createAArch64SpeculationHardeningPass();
FunctionPass *createAArch64LoadStoreOptimizationPass();
ModulePass *createAArch64LowerHomogeneousPrologEpilogPass();
Expand Down
256 changes: 105 additions & 151 deletions llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,20 +13,16 @@

#include "AArch64InstrInfo.h"
#include "AArch64Subtarget.h"
#include "Utils/AArch64BaseInfo.h"
#include "llvm/CodeGen/IndirectThunks.h"
#include "llvm/CodeGen/MachineBasicBlock.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineOperand.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/RegisterScavenging.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/Pass.h"
#include "llvm/Support/CodeGen.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Target/TargetMachine.h"
#include <cassert>

Expand All @@ -36,38 +32,42 @@ using namespace llvm;

#define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"

static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";

namespace {

class AArch64SLSHardening : public MachineFunctionPass {
public:
const TargetInstrInfo *TII;
const TargetRegisterInfo *TRI;
const AArch64Subtarget *ST;
// Set of inserted thunks: bitmask with bits corresponding to
// indexes in SLSBLRThunks array.
typedef uint32_t ThunksSet;

static char ID;

AArch64SLSHardening() : MachineFunctionPass(ID) {
initializeAArch64SLSHardeningPass(*PassRegistry::getPassRegistry());
struct SLSHardeningInserter : ThunkInserter<SLSHardeningInserter, ThunksSet> {
public:
const char *getThunkPrefix() { return SLSBLRNamePrefix; }
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);

bool runOnMachineFunction(MachineFunction &Fn) override;
private:
bool ComdatThunks = true;

StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }
bool hardenReturnsAndBRs(MachineModuleInfo &MMI, MachineBasicBlock &MBB);
bool hardenBLRs(MachineModuleInfo &MMI, MachineBasicBlock &MBB,
ThunksSet &Thunks);

private:
bool hardenReturnsAndBRs(MachineBasicBlock &MBB) const;
bool hardenBLRs(MachineBasicBlock &MBB) const;
MachineBasicBlock &ConvertBLRToBL(MachineBasicBlock &MBB,
MachineBasicBlock::instr_iterator) const;
void convertBLRToBL(MachineModuleInfo &MMI, MachineBasicBlock &MBB,
MachineBasicBlock::instr_iterator MBBI,
ThunksSet &Thunks);
};

} // end anonymous namespace

char AArch64SLSHardening::ID = 0;

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

static void insertSpeculationBarrier(const AArch64Subtarget *ST,
MachineBasicBlock &MBB,
MachineBasicBlock::iterator MBBI,
Expand All @@ -90,18 +90,18 @@ static void insertSpeculationBarrier(const AArch64Subtarget *ST,
BuildMI(MBB, MBBI, DL, TII->get(BarrierOpc));
}

bool AArch64SLSHardening::runOnMachineFunction(MachineFunction &MF) {
ST = &MF.getSubtarget<AArch64Subtarget>();
TII = MF.getSubtarget().getInstrInfo();
TRI = MF.getSubtarget().getRegisterInfo();
ThunksSet SLSHardeningInserter::insertThunks(MachineModuleInfo &MMI,
MachineFunction &MF,
ThunksSet ExistingThunks) {
const AArch64Subtarget *ST = &MF.getSubtarget<AArch64Subtarget>();

bool Modified = false;
for (auto &MBB : MF) {
Modified |= hardenReturnsAndBRs(MBB);
Modified |= hardenBLRs(MBB);
if (ST->hardenSlsRetBr())
hardenReturnsAndBRs(MMI, MBB);
if (ST->hardenSlsBlr())
hardenBLRs(MMI, MBB, ExistingThunks);
}

return Modified;
return ExistingThunks;
}

static bool isBLR(const MachineInstr &MI) {
Expand All @@ -120,9 +120,10 @@ static bool isBLR(const MachineInstr &MI) {
return false;
}

bool AArch64SLSHardening::hardenReturnsAndBRs(MachineBasicBlock &MBB) const {
if (!ST->hardenSlsRetBr())
return false;
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;
Expand All @@ -138,78 +139,55 @@ bool AArch64SLSHardening::hardenReturnsAndBRs(MachineBasicBlock &MBB) const {
return Modified;
}

static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";

static const unsigned NumPermittedRegs = 29;
static const struct ThunkNameAndReg {
const char* Name;
Register Reg;
} SLSBLRThunks[] = {
{ "__llvm_slsblr_thunk_x0", AArch64::X0},
{ "__llvm_slsblr_thunk_x1", AArch64::X1},
{ "__llvm_slsblr_thunk_x2", AArch64::X2},
{ "__llvm_slsblr_thunk_x3", AArch64::X3},
{ "__llvm_slsblr_thunk_x4", AArch64::X4},
{ "__llvm_slsblr_thunk_x5", AArch64::X5},
{ "__llvm_slsblr_thunk_x6", AArch64::X6},
{ "__llvm_slsblr_thunk_x7", AArch64::X7},
{ "__llvm_slsblr_thunk_x8", AArch64::X8},
{ "__llvm_slsblr_thunk_x9", AArch64::X9},
{ "__llvm_slsblr_thunk_x10", AArch64::X10},
{ "__llvm_slsblr_thunk_x11", AArch64::X11},
{ "__llvm_slsblr_thunk_x12", AArch64::X12},
{ "__llvm_slsblr_thunk_x13", AArch64::X13},
{ "__llvm_slsblr_thunk_x14", AArch64::X14},
{ "__llvm_slsblr_thunk_x15", AArch64::X15},
// X16 and X17 are deliberately missing, as the mitigation requires those
// register to not be used in BLR. See comment in ConvertBLRToBL for more
// details.
{ "__llvm_slsblr_thunk_x18", AArch64::X18},
{ "__llvm_slsblr_thunk_x19", AArch64::X19},
{ "__llvm_slsblr_thunk_x20", AArch64::X20},
{ "__llvm_slsblr_thunk_x21", AArch64::X21},
{ "__llvm_slsblr_thunk_x22", AArch64::X22},
{ "__llvm_slsblr_thunk_x23", AArch64::X23},
{ "__llvm_slsblr_thunk_x24", AArch64::X24},
{ "__llvm_slsblr_thunk_x25", AArch64::X25},
{ "__llvm_slsblr_thunk_x26", AArch64::X26},
{ "__llvm_slsblr_thunk_x27", AArch64::X27},
{ "__llvm_slsblr_thunk_x28", AArch64::X28},
{ "__llvm_slsblr_thunk_x29", AArch64::FP},
// X30 is deliberately missing, for similar reasons as X16 and X17 are
// missing.
{ "__llvm_slsblr_thunk_x31", AArch64::XZR},
} SLSBLRThunks[NumPermittedRegs] = {
{"__llvm_slsblr_thunk_x0", AArch64::X0},
{"__llvm_slsblr_thunk_x1", AArch64::X1},
{"__llvm_slsblr_thunk_x2", AArch64::X2},
{"__llvm_slsblr_thunk_x3", AArch64::X3},
{"__llvm_slsblr_thunk_x4", AArch64::X4},
{"__llvm_slsblr_thunk_x5", AArch64::X5},
{"__llvm_slsblr_thunk_x6", AArch64::X6},
{"__llvm_slsblr_thunk_x7", AArch64::X7},
{"__llvm_slsblr_thunk_x8", AArch64::X8},
{"__llvm_slsblr_thunk_x9", AArch64::X9},
{"__llvm_slsblr_thunk_x10", AArch64::X10},
{"__llvm_slsblr_thunk_x11", AArch64::X11},
{"__llvm_slsblr_thunk_x12", AArch64::X12},
{"__llvm_slsblr_thunk_x13", AArch64::X13},
{"__llvm_slsblr_thunk_x14", AArch64::X14},
{"__llvm_slsblr_thunk_x15", AArch64::X15},
// X16 and X17 are deliberately missing, as the mitigation requires those
// register to not be used in BLR. See comment in ConvertBLRToBL for more
// details.
{"__llvm_slsblr_thunk_x18", AArch64::X18},
{"__llvm_slsblr_thunk_x19", AArch64::X19},
{"__llvm_slsblr_thunk_x20", AArch64::X20},
{"__llvm_slsblr_thunk_x21", AArch64::X21},
{"__llvm_slsblr_thunk_x22", AArch64::X22},
{"__llvm_slsblr_thunk_x23", AArch64::X23},
{"__llvm_slsblr_thunk_x24", AArch64::X24},
{"__llvm_slsblr_thunk_x25", AArch64::X25},
{"__llvm_slsblr_thunk_x26", AArch64::X26},
{"__llvm_slsblr_thunk_x27", AArch64::X27},
{"__llvm_slsblr_thunk_x28", AArch64::X28},
{"__llvm_slsblr_thunk_x29", AArch64::FP},
// X30 is deliberately missing, for similar reasons as X16 and X17 are
// missing.
{"__llvm_slsblr_thunk_x31", AArch64::XZR},
};

namespace {
struct SLSBLRThunkInserter : ThunkInserter<SLSBLRThunkInserter> {
const char *getThunkPrefix() { return SLSBLRNamePrefix; }
bool mayUseThunk(const MachineFunction &MF) {
ComdatThunks &= !MF.getSubtarget<AArch64Subtarget>().hardenSlsNoComdat();
return MF.getSubtarget<AArch64Subtarget>().hardenSlsBlr();
}
bool insertThunks(MachineModuleInfo &MMI, MachineFunction &MF,
bool ExistingThunks);
void populateThunk(MachineFunction &MF);

private:
bool ComdatThunks = true;
};
} // namespace

bool SLSBLRThunkInserter::insertThunks(MachineModuleInfo &MMI,
MachineFunction &MF,
bool ExistingThunks) {
if (ExistingThunks)
return false;
// FIXME: It probably would be possible to filter which thunks to produce
// based on which registers are actually used in BLR instructions in this
// function. But would that be a worthwhile optimization?
for (auto T : SLSBLRThunks)
createThunkFunction(MMI, T.Name, ComdatThunks);
return true;
unsigned getThunkIndex(Register Reg) {
for (unsigned I = 0; I < NumPermittedRegs; ++I)
if (SLSBLRThunks[I].Reg == Reg)
return I;
llvm_unreachable("Unexpected register");
}

void SLSBLRThunkInserter::populateThunk(MachineFunction &MF) {
void SLSHardeningInserter::populateThunk(MachineFunction &MF) {
assert(MF.getFunction().hasComdat() == ComdatThunks &&
"ComdatThunks value changed since MF creation");
// FIXME: How to better communicate Register number, rather than through
Expand Down Expand Up @@ -258,8 +236,9 @@ void SLSBLRThunkInserter::populateThunk(MachineFunction &MF) {
Entry->end(), DebugLoc(), true /*AlwaysUseISBDSB*/);
}

MachineBasicBlock &AArch64SLSHardening::ConvertBLRToBL(
MachineBasicBlock &MBB, MachineBasicBlock::instr_iterator MBBI) const {
void SLSHardeningInserter::convertBLRToBL(
MachineModuleInfo &MMI, MachineBasicBlock &MBB,
MachineBasicBlock::instr_iterator MBBI, ThunksSet &Thunks) {
// Transform a BLR to a BL as follows:
// Before:
// |-----------------------------|
Expand All @@ -285,7 +264,6 @@ MachineBasicBlock &AArch64SLSHardening::ConvertBLRToBL(
// | barrierInsts |
// |-----------------------------|
//
// The __llvm_slsblr_thunk_xN thunks are created by the SLSBLRThunkInserter.
// This function merely needs to transform BLR xN into BL
// __llvm_slsblr_thunk_xN.
//
Expand Down Expand Up @@ -318,37 +296,16 @@ MachineBasicBlock &AArch64SLSHardening::ConvertBLRToBL(
}
DebugLoc DL = BLR.getDebugLoc();

// If we'd like to support also BLRAA and BLRAB instructions, we'd need
// a lot more different kind of thunks.
// For example, a
//
// BLRAA xN, xM
//
// instruction probably would need to be transformed to something like:
//
// BL __llvm_slsblraa_thunk_x<N>_x<M>
//
// __llvm_slsblraa_thunk_x<N>_x<M>:
// BRAA x<N>, x<M>
// barrierInsts
//
// Given that about 30 different values of N are possible and about 30
// different values of M are possible in the above, with the current way
// of producing indirect thunks, we'd be producing about 30 times 30, i.e.
// about 900 thunks (where most might not be actually called). This would
// multiply further by two to support both BLRAA and BLRAB variants of those
// instructions.
// If we'd want to support this, we'd probably need to look into a different
// way to produce thunk functions, based on which variants are actually
// needed, rather than producing all possible variants.
// So far, LLVM does never produce BLRA* instructions, so let's leave this
// for the future when LLVM can start producing BLRA* instructions.
MachineFunction &MF = *MBBI->getMF();
MCContext &Context = MBB.getParent()->getContext();
auto ThunkIt =
llvm::find_if(SLSBLRThunks, [Reg](auto T) { return T.Reg == Reg; });
assert (ThunkIt != std::end(SLSBLRThunks));
MCSymbol *Sym = Context.getOrCreateSymbol(ThunkIt->Name);
const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
unsigned ThunkIndex = getThunkIndex(Reg);
StringRef ThunkName = SLSBLRThunks[ThunkIndex].Name;
MCSymbol *Sym = Context.getOrCreateSymbol(ThunkName);
if (!(Thunks & (1u << ThunkIndex))) {
Thunks |= 1u << ThunkIndex;
createThunkFunction(MMI, ThunkName, ComdatThunks);
}

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

Expand Down Expand Up @@ -386,13 +343,11 @@ MachineBasicBlock &AArch64SLSHardening::ConvertBLRToBL(
RegIsKilled /*isKill*/));
// Remove BLR instruction
MBB.erase(MBBI);

return MBB;
}

bool AArch64SLSHardening::hardenBLRs(MachineBasicBlock &MBB) const {
if (!ST->hardenSlsBlr())
return false;
bool SLSHardeningInserter::hardenBLRs(MachineModuleInfo &MMI,
MachineBasicBlock &MBB,
ThunksSet &Thunks) {
bool Modified = false;
MachineBasicBlock::instr_iterator MBBI = MBB.instr_begin(),
E = MBB.instr_end();
Expand All @@ -401,31 +356,30 @@ bool AArch64SLSHardening::hardenBLRs(MachineBasicBlock &MBB) const {
MachineInstr &MI = *MBBI;
NextMBBI = std::next(MBBI);
if (isBLR(MI)) {
ConvertBLRToBL(MBB, MBBI);
convertBLRToBL(MMI, MBB, MBBI, Thunks);
Modified = true;
}
}
return Modified;
}

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

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

AArch64IndirectThunks() : ThunkInserterPass(ID) {}
AArch64SLSHardening() : ThunkInserterPass(ID) {}

StringRef getPassName() const override { return "AArch64 Indirect Thunks"; }
StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }
};

} // end anonymous namespace

char AArch64IndirectThunks::ID = 0;
char AArch64SLSHardening::ID = 0;

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

FunctionPass *llvm::createAArch64IndirectThunks() {
return new AArch64IndirectThunks();
FunctionPass *llvm::createAArch64SLSHardeningPass() {
return new AArch64SLSHardening();
}
1 change: 0 additions & 1 deletion llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -861,7 +861,6 @@ void AArch64PassConfig::addPreEmitPass() {
}

void AArch64PassConfig::addPostBBSections() {
addPass(createAArch64IndirectThunks());
addPass(createAArch64SLSHardeningPass());
addPass(createAArch64PointerAuthPass());
if (EnableBranchTargets)
Expand Down
1 change: 0 additions & 1 deletion llvm/test/CodeGen/AArch64/O0-pipeline.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,6 @@
; CHECK-NEXT: StackMap Liveness Analysis
; CHECK-NEXT: Live DEBUG_VALUE analysis
; CHECK-NEXT: Machine Sanitizer Binary Metadata
; CHECK-NEXT: AArch64 Indirect Thunks
; CHECK-NEXT: AArch64 sls hardening pass
; CHECK-NEXT: AArch64 Pointer Authentication
; CHECK-NEXT: AArch64 Branch Targets
Expand Down
Loading
Loading