diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp index 6d9390235baf1..7290b3f67c2e3 100644 --- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp +++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp @@ -973,8 +973,7 @@ bool AArch64FrameLowering::shouldSignReturnAddressEverywhere( if (MF.getTarget().getMCAsmInfo()->usesWindowsCFI()) return false; const AArch64FunctionInfo *AFI = MF.getInfo(); - bool SignReturnAddressAll = AFI->shouldSignReturnAddress(/*SpillsLR=*/false); - return SignReturnAddressAll; + return AFI->getSignReturnAddressCondition() == SignReturnAddress::All; } // Given a load or a store instruction, generate an appropriate unwinding SEH diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp index 221812f1ebc7b..cdee602deb43c 100644 --- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp +++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp @@ -9771,8 +9771,8 @@ outliningCandidatesSigningScopeConsensus(const outliner::Candidate &a, const auto &MFIa = a.getMF()->getInfo(); const auto &MFIb = b.getMF()->getInfo(); - return MFIa->shouldSignReturnAddress(false) == MFIb->shouldSignReturnAddress(false) && - MFIa->shouldSignReturnAddress(true) == MFIb->shouldSignReturnAddress(true); + return MFIa->getSignReturnAddressCondition() == + MFIb->getSignReturnAddressCondition(); } static bool @@ -9863,10 +9863,11 @@ AArch64InstrInfo::getOutliningCandidateInfo( // Performing a tail call may require extra checks when PAuth is enabled. // If PAuth is disabled, set it to zero for uniformity. unsigned NumBytesToCheckLRInTCEpilogue = 0; - if (RepeatedSequenceLocs[0] - .getMF() - ->getInfo() - ->shouldSignReturnAddress(true)) { + const auto RASignCondition = RepeatedSequenceLocs[0] + .getMF() + ->getInfo() + ->getSignReturnAddressCondition(); + if (RASignCondition != SignReturnAddress::None) { // One PAC and one AUT instructions NumBytesToCreateFrame += 8; @@ -10670,7 +10671,9 @@ void AArch64InstrInfo::buildOutlinedFrame( Et = MBB.insert(Et, LDRXpost); } - bool ShouldSignReturnAddr = FI->shouldSignReturnAddress(!IsLeafFunction); + auto RASignCondition = FI->getSignReturnAddressCondition(); + bool ShouldSignReturnAddr = AArch64FunctionInfo::shouldSignReturnAddress( + RASignCondition, !IsLeafFunction); // If this is a tail call outlined function, then there's already a return. if (OF.FrameConstructionID == MachineOutlinerTailCall || diff --git a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp index 343fd81ace0a5..d6d63f3b6198d 100644 --- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp +++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp @@ -16,6 +16,7 @@ #include "AArch64MachineFunctionInfo.h" #include "AArch64InstrInfo.h" #include "AArch64Subtarget.h" +#include "llvm/ADT/StringSwitch.h" #include "llvm/IR/Constants.h" #include "llvm/IR/Metadata.h" #include "llvm/IR/Module.h" @@ -63,24 +64,21 @@ void AArch64FunctionInfo::initializeBaseYamlFields( setHasStreamingModeChanges(*YamlMFI.HasStreamingModeChanges); } -static std::pair GetSignReturnAddress(const Function &F) { +static SignReturnAddress GetSignReturnAddress(const Function &F) { if (F.hasFnAttribute("ptrauth-returns")) - return {true, false}; // non-leaf + return SignReturnAddress::NonLeaf; + // The function should be signed in the following situations: // - sign-return-address=all // - sign-return-address=non-leaf and the functions spills the LR if (!F.hasFnAttribute("sign-return-address")) - return {false, false}; + return SignReturnAddress::None; StringRef Scope = F.getFnAttribute("sign-return-address").getValueAsString(); - if (Scope == "none") - return {false, false}; - - if (Scope == "all") - return {true, true}; - - assert(Scope == "non-leaf"); - return {true, false}; + return StringSwitch(Scope) + .Case("none", SignReturnAddress::None) + .Case("non-leaf", SignReturnAddress::NonLeaf) + .Case("all", SignReturnAddress::All); } static bool ShouldSignWithBKey(const Function &F, const AArch64Subtarget &STI) { @@ -116,7 +114,7 @@ AArch64FunctionInfo::AArch64FunctionInfo(const Function &F, // HasRedZone here. if (F.hasFnAttribute(Attribute::NoRedZone)) HasRedZone = false; - std::tie(SignReturnAddress, SignReturnAddressAll) = GetSignReturnAddress(F); + SignCondition = GetSignReturnAddress(F); SignWithBKey = ShouldSignWithBKey(F, *STI); HasELFSignedGOT = hasELFSignedGOTHelper(F, STI); // TODO: skip functions that have no instrumented allocas for optimization @@ -169,23 +167,27 @@ MachineFunctionInfo *AArch64FunctionInfo::clone( return DestMF.cloneInfo(*this); } -bool AArch64FunctionInfo::shouldSignReturnAddress(bool SpillsLR) const { - if (!SignReturnAddress) - return false; - if (SignReturnAddressAll) - return true; - return SpillsLR; -} - static bool isLRSpilled(const MachineFunction &MF) { return llvm::any_of( MF.getFrameInfo().getCalleeSavedInfo(), [](const auto &Info) { return Info.getReg() == AArch64::LR; }); } +bool AArch64FunctionInfo::shouldSignReturnAddress(SignReturnAddress Condition, + bool IsLRSpilled) { + switch (Condition) { + case SignReturnAddress::None: + return false; + case SignReturnAddress::NonLeaf: + return IsLRSpilled; + case SignReturnAddress::All: + return true; + } +} + bool AArch64FunctionInfo::shouldSignReturnAddress( const MachineFunction &MF) const { - return shouldSignReturnAddress(isLRSpilled(MF)); + return shouldSignReturnAddress(SignCondition, isLRSpilled(MF)); } bool AArch64FunctionInfo::needsShadowCallStackPrologueEpilogue( diff --git a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h index c137a5458b047..00e0c2511aaf0 100644 --- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h +++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h @@ -42,6 +42,15 @@ struct TPIDR2Object { unsigned Uses = 0; }; +/// Condition of signing the return address in a function. +/// +/// Corresponds to possible values of "sign-return-address" function attribute. +enum class SignReturnAddress { + None, + NonLeaf, + All, +}; + /// AArch64FunctionInfo - This class is derived from MachineFunctionInfo and /// contains private AArch64-specific information for each MachineFunction. class AArch64FunctionInfo final : public MachineFunctionInfo { @@ -170,13 +179,8 @@ class AArch64FunctionInfo final : public MachineFunctionInfo { // CalleeSavedStackSize) to the address of the frame record. int CalleeSaveBaseToFrameRecordOffset = 0; - /// SignReturnAddress is true if PAC-RET is enabled for the function with - /// defaults being sign non-leaf functions only, with the B key. - bool SignReturnAddress = false; - - /// SignReturnAddressAll modifies the default PAC-RET mode to signing leaf - /// functions as well. - bool SignReturnAddressAll = false; + /// SignCondition controls when PAC-RET protection should be used. + SignReturnAddress SignCondition = SignReturnAddress::None; /// SignWithBKey modifies the default PAC-RET mode to signing with the B key. bool SignWithBKey = false; @@ -591,8 +595,14 @@ class AArch64FunctionInfo final : public MachineFunctionInfo { CalleeSaveBaseToFrameRecordOffset = Offset; } + static bool shouldSignReturnAddress(SignReturnAddress Condition, + bool IsLRSpilled); + bool shouldSignReturnAddress(const MachineFunction &MF) const; - bool shouldSignReturnAddress(bool SpillsLR) const; + + SignReturnAddress getSignReturnAddressCondition() const { + return SignCondition; + } bool needsShadowCallStackPrologueEpilogue(MachineFunction &MF) const;