[llvm] [AArch64][PAC] Protect the entire function if pac-ret+leaf is passed (PR #140895)
via llvm-commits
llvm-commits at lists.llvm.org
Wed May 21 06:18:55 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-aarch64
Author: Anatoly Trosinenko (atrosinenko)
<details>
<summary>Changes</summary>
Normally, pac-ret hardening is emitted as part of function prologues and epilogues, thus it is affected by the shrink-wrapping optimization.
As protecting LR when it is spilled to the stack is already handled by regular -mbranch-protection=pac-ret option, it is reasonable to assume that pac-ret+leaf option means the user wants to apply pac-ret hardening to as much code as possible. For that reason, if pac-ret+leaf hardening mode is requested, this patch moves the emission of PAUTH_PROLOGUE (or PAUTH_EPILOGUE) pseudos from emitPrologue (emitEpilogue) methods of the AArch64FrameLowering class to processFunctionBeforeFrameIndicesReplaced.
This commit only affects where LR is signed and authenticated, but does not otherwise prevents the shrink-wrapping optimization. Moreover, without "+leaf" modifier PAUTH_(PROLOGUE|EPILOGUE) pseudos respect the shrink-wrapping optimization just as any other prologue/epilogue code.
---
Full diff: https://github.com/llvm/llvm-project/pull/140895.diff
7 Files Affected:
- (modified) llvm/lib/Target/AArch64/AArch64FrameLowering.cpp (+55-5)
- (modified) llvm/lib/Target/AArch64/AArch64FrameLowering.h (+8)
- (modified) llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp (+4)
- (modified) llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h (+1)
- (modified) llvm/test/CodeGen/AArch64/sign-return-address-cfi-negate-ra-state.ll (+3-3)
- (modified) llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll (+1-1)
- (modified) llvm/test/CodeGen/AArch64/sign-return-address.ll (+86)
``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
index 0f33e77d4eecc..705c23384c0ef 100644
--- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -1679,6 +1679,39 @@ static void getLivePhysRegsUpTo(MachineInstr &MI, const TargetRegisterInfo &TRI,
}
#endif
+void AArch64FrameLowering::emitPacRetPlusLeafHardening(
+ MachineFunction &MF) const {
+ const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
+ const TargetInstrInfo *TII = Subtarget.getInstrInfo();
+
+ auto EmitSignRA = [&](MachineBasicBlock &MBB) {
+ DebugLoc DL; // Set debug location to unknown.
+ MachineBasicBlock::iterator MBBI = MBB.begin();
+
+ BuildMI(MBB, MBBI, DL, TII->get(AArch64::PAUTH_PROLOGUE))
+ .setMIFlag(MachineInstr::FrameSetup);
+ };
+
+ auto EmitAuthRA = [&](MachineBasicBlock &MBB) {
+ DebugLoc DL;
+ MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
+ if (MBBI != MBB.end())
+ DL = MBBI->getDebugLoc();
+
+ BuildMI(MBB, MBBI, DL, TII->get(AArch64::PAUTH_EPILOGUE))
+ .setMIFlag(MachineInstr::FrameDestroy);
+ };
+
+ // This should be in sync with PEIImpl::calculateSaveRestoreBlocks.
+ EmitSignRA(MF.front());
+ for (MachineBasicBlock &MBB : MF) {
+ if (MBB.isEHFuncletEntry())
+ EmitSignRA(MBB);
+ if (MBB.isReturnBlock())
+ EmitAuthRA(MBB);
+ }
+}
+
void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
MachineBasicBlock &MBB) const {
MachineBasicBlock::iterator MBBI = MBB.begin();
@@ -1753,8 +1786,12 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
MFnI.needsDwarfUnwindInfo(MF));
if (MFnI.shouldSignReturnAddress(MF)) {
- BuildMI(MBB, MBBI, DL, TII->get(AArch64::PAUTH_PROLOGUE))
- .setMIFlag(MachineInstr::FrameSetup);
+ // If pac-ret+leaf is in effect, PAUTH_PROLOGUE pseudo instructions
+ // are inserted by emitPacRetPlusLeafHardening().
+ if (!MFnI.shouldSignReturnAddressEverywhere()) {
+ BuildMI(MBB, MBBI, DL, TII->get(AArch64::PAUTH_PROLOGUE))
+ .setMIFlag(MachineInstr::FrameSetup);
+ }
if (NeedsWinCFI)
HasWinCFI = true; // AArch64PointerAuth pass will insert SEH_PACSignLR
}
@@ -2237,9 +2274,13 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
auto FinishingTouches = make_scope_exit([&]() {
if (AFI->shouldSignReturnAddress(MF)) {
- BuildMI(MBB, MBB.getFirstTerminator(), DL,
- TII->get(AArch64::PAUTH_EPILOGUE))
- .setMIFlag(MachineInstr::FrameDestroy);
+ // If pac-ret+leaf is in effect, PAUTH_EPILOGUE pseudo instructions
+ // are inserted by emitPacRetPlusLeafHardening().
+ if (!AFI->shouldSignReturnAddressEverywhere()) {
+ BuildMI(MBB, MBB.getFirstTerminator(), DL,
+ TII->get(AArch64::PAUTH_EPILOGUE))
+ .setMIFlag(MachineInstr::FrameDestroy);
+ }
if (NeedsWinCFI)
HasWinCFI = true; // AArch64PointerAuth pass will insert SEH_PACSignLR
}
@@ -4934,6 +4975,8 @@ static void emitVGSaveRestore(MachineBasicBlock::iterator II,
void AArch64FrameLowering::processFunctionBeforeFrameIndicesReplaced(
MachineFunction &MF, RegScavenger *RS = nullptr) const {
+ AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+
for (auto &BB : MF)
for (MachineBasicBlock::iterator II = BB.begin(); II != BB.end();) {
if (requiresSaveVG(MF))
@@ -4941,6 +4984,13 @@ void AArch64FrameLowering::processFunctionBeforeFrameIndicesReplaced(
else if (StackTaggingMergeSetTag)
II = tryMergeAdjacentSTG(II, this, RS);
}
+
+ // By the time this method is called, most of the prologue/epilogue code is
+ // already emitted, whether its location was affected by the shrink-wrapping
+ // optimization or not.
+ if (!MF.getFunction().hasFnAttribute(Attribute::Naked) &&
+ AFI->shouldSignReturnAddressEverywhere())
+ emitPacRetPlusLeafHardening(MF);
}
/// For Win64 AArch64 EH, the offset to the Unwind object is from the SP
diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.h b/llvm/lib/Target/AArch64/AArch64FrameLowering.h
index e7d52bb350f13..dd44588ab8b26 100644
--- a/llvm/lib/Target/AArch64/AArch64FrameLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.h
@@ -36,6 +36,14 @@ class AArch64FrameLowering : public TargetFrameLowering {
void emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
void emitEpilogue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
+ /// Harden the entire function with pac-ret.
+ ///
+ /// If pac-ret+leaf is requested, we want to harden as much code as possible.
+ /// This function inserts pac-ret hardening at the points where prologue and
+ /// epilogue are traditionally inserted, ignoring possible shrink-wrapping
+ /// optimization.
+ void emitPacRetPlusLeafHardening(MachineFunction &MF) const;
+
bool enableCFIFixup(const MachineFunction &MF) const override;
bool enableFullCFIFixup(const MachineFunction &MF) const override;
diff --git a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp
index 5bcff61cef4b1..bd3f42828ddd3 100644
--- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp
@@ -159,6 +159,10 @@ bool AArch64FunctionInfo::shouldSignReturnAddress(
return shouldSignReturnAddress(isLRSpilled(MF));
}
+bool AArch64FunctionInfo::shouldSignReturnAddressEverywhere() const {
+ return SignReturnAddressAll;
+}
+
bool AArch64FunctionInfo::needsShadowCallStackPrologueEpilogue(
MachineFunction &MF) const {
if (!(isLRSpilled(MF) &&
diff --git a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h
index d3026ca45c349..cb985b48bd69d 100644
--- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h
+++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h
@@ -523,6 +523,7 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
bool shouldSignReturnAddress(const MachineFunction &MF) const;
bool shouldSignReturnAddress(bool SpillsLR) const;
+ bool shouldSignReturnAddressEverywhere() const;
bool needsShadowCallStackPrologueEpilogue(MachineFunction &MF) const;
diff --git a/llvm/test/CodeGen/AArch64/sign-return-address-cfi-negate-ra-state.ll b/llvm/test/CodeGen/AArch64/sign-return-address-cfi-negate-ra-state.ll
index 4d4b7c215b978..bf70bf30534ec 100644
--- a/llvm/test/CodeGen/AArch64/sign-return-address-cfi-negate-ra-state.ll
+++ b/llvm/test/CodeGen/AArch64/sign-return-address-cfi-negate-ra-state.ll
@@ -74,9 +74,9 @@ define hidden noundef i32 @baz_async(i32 noundef %a) #0 uwtable(async) {
; CHECK-V8A-NEXT: bl _Z3bari
; CHECK-V8A-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
; CHECK-V8A-NEXT: .cfi_def_cfa_offset 0
+; CHECK-V8A-NEXT: .cfi_restore w30
; CHECK-V8A-NEXT: hint #29
; CHECK-V8A-NEXT: .cfi_negate_ra_state
-; CHECK-V8A-NEXT: .cfi_restore w30
; CHECK-V8A-NEXT: b _Z3bari
; CHECK-V8A-NEXT: .LBB1_2: // %if.else
; CHECK-V8A-NEXT: .cfi_restore_state
@@ -84,9 +84,9 @@ define hidden noundef i32 @baz_async(i32 noundef %a) #0 uwtable(async) {
; CHECK-V8A-NEXT: add w0, w0, #1
; CHECK-V8A-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
; CHECK-V8A-NEXT: .cfi_def_cfa_offset 0
+; CHECK-V8A-NEXT: .cfi_restore w30
; CHECK-V8A-NEXT: hint #29
; CHECK-V8A-NEXT: .cfi_negate_ra_state
-; CHECK-V8A-NEXT: .cfi_restore w30
; CHECK-V8A-NEXT: ret
;
; CHECK-V83A-LABEL: baz_async:
@@ -103,9 +103,9 @@ define hidden noundef i32 @baz_async(i32 noundef %a) #0 uwtable(async) {
; CHECK-V83A-NEXT: bl _Z3bari
; CHECK-V83A-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
; CHECK-V83A-NEXT: .cfi_def_cfa_offset 0
+; CHECK-V83A-NEXT: .cfi_restore w30
; CHECK-V83A-NEXT: autiasp
; CHECK-V83A-NEXT: .cfi_negate_ra_state
-; CHECK-V83A-NEXT: .cfi_restore w30
; CHECK-V83A-NEXT: b _Z3bari
; CHECK-V83A-NEXT: .LBB1_2: // %if.else
; CHECK-V83A-NEXT: .cfi_restore_state
diff --git a/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll b/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll
index 032d3cc05961f..90e48fe0ed042 100644
--- a/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll
+++ b/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll
@@ -156,8 +156,8 @@ define i32 @tailcall_two_branches(i1 %0) "sign-return-address"="all" {
; COMMON: str x30, [sp, #-16]!
; COMMON: bl callee2
; COMMON: ldr x30, [sp], #16
-; COMMON-NEXT: [[AUTIASP]]
; COMMON-NEXT: .[[ELSE]]:
+; COMMON-NEXT: [[AUTIASP]]
; LDR-NEXT: ldr w16, [x30]
;
diff --git a/llvm/test/CodeGen/AArch64/sign-return-address.ll b/llvm/test/CodeGen/AArch64/sign-return-address.ll
index dafe0d71ceb5f..61b72d806e082 100644
--- a/llvm/test/CodeGen/AArch64/sign-return-address.ll
+++ b/llvm/test/CodeGen/AArch64/sign-return-address.ll
@@ -149,6 +149,92 @@ define i32 @non_leaf_scs(i32 %x) "sign-return-address"="non-leaf" shadowcallstac
ret i32 %call
}
+; By default, pac-ret hardening respects shrink-wrapping optimization.
+define void @shrink_wrap_sign_non_leaf(i32 %x, i32 %cond) "sign-return-address"="non-leaf" {
+; COMPAT-LABEL: shrink_wrap_sign_non_leaf:
+; COMPAT: // %bb.0: // %entry
+; COMPAT-NEXT: cbz w1, .LBB8_2
+; COMPAT-NEXT: // %bb.1: // %common.ret
+; COMPAT-NEXT: ret
+; COMPAT-NEXT: .LBB8_2: // %if.then
+; COMPAT-NEXT: hint #25
+; COMPAT-NEXT: .cfi_negate_ra_state
+; COMPAT-NEXT: str x30, [sp, #-16]! // 8-byte Folded Spill
+; COMPAT-NEXT: .cfi_def_cfa_offset 16
+; COMPAT-NEXT: .cfi_offset w30, -16
+; COMPAT-NEXT: bl foo
+; COMPAT-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
+; COMPAT-NEXT: hint #29
+; COMPAT-NEXT: ret
+;
+; V83A-LABEL: shrink_wrap_sign_non_leaf:
+; V83A: // %bb.0: // %entry
+; V83A-NEXT: cbz w1, .LBB8_2
+; V83A-NEXT: // %bb.1: // %common.ret
+; V83A-NEXT: ret
+; V83A-NEXT: .LBB8_2: // %if.then
+; V83A-NEXT: paciasp
+; V83A-NEXT: .cfi_negate_ra_state
+; V83A-NEXT: str x30, [sp, #-16]! // 8-byte Folded Spill
+; V83A-NEXT: .cfi_def_cfa_offset 16
+; V83A-NEXT: .cfi_offset w30, -16
+; V83A-NEXT: bl foo
+; V83A-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
+; V83A-NEXT: retaa
+entry:
+ %cond.bool = icmp eq i32 %cond, 0
+ br i1 %cond.bool, label %if.then, label %if.else
+if.then:
+ %call = call i32 @foo(i32 %x)
+ ret void
+if.else:
+ ret void
+}
+
+; When "+leaf" is specified to harden everything, pac-ret hardens the entire
+; function, ignoring shrink-wrapping.
+define void @shrink_wrap_sign_all(i32 %x, i32 %cond) "sign-return-address"="all" {
+; COMPAT-LABEL: shrink_wrap_sign_all:
+; COMPAT: // %bb.0: // %entry
+; COMPAT-NEXT: hint #25
+; COMPAT-NEXT: .cfi_negate_ra_state
+; COMPAT-NEXT: cbz w1, .LBB9_2
+; COMPAT-NEXT: // %bb.1: // %common.ret
+; COMPAT-NEXT: hint #29
+; COMPAT-NEXT: ret
+; COMPAT-NEXT: .LBB9_2: // %if.then
+; COMPAT-NEXT: str x30, [sp, #-16]! // 8-byte Folded Spill
+; COMPAT-NEXT: .cfi_def_cfa_offset 16
+; COMPAT-NEXT: .cfi_offset w30, -16
+; COMPAT-NEXT: bl foo
+; COMPAT-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
+; COMPAT-NEXT: hint #29
+; COMPAT-NEXT: ret
+;
+; V83A-LABEL: shrink_wrap_sign_all:
+; V83A: // %bb.0: // %entry
+; V83A-NEXT: paciasp
+; V83A-NEXT: .cfi_negate_ra_state
+; V83A-NEXT: cbz w1, .LBB9_2
+; V83A-NEXT: // %bb.1: // %common.ret
+; V83A-NEXT: retaa
+; V83A-NEXT: .LBB9_2: // %if.then
+; V83A-NEXT: str x30, [sp, #-16]! // 8-byte Folded Spill
+; V83A-NEXT: .cfi_def_cfa_offset 16
+; V83A-NEXT: .cfi_offset w30, -16
+; V83A-NEXT: bl foo
+; V83A-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
+; V83A-NEXT: retaa
+entry:
+ %cond.bool = icmp eq i32 %cond, 0
+ br i1 %cond.bool, label %if.then, label %if.else
+if.then:
+ %call = call i32 @foo(i32 %x)
+ ret void
+if.else:
+ ret void
+}
+
define i32 @leaf_sign_all_v83(i32 %x) "sign-return-address"="all" "target-features"="+v8.3a" {
; CHECK-LABEL: leaf_sign_all_v83:
; CHECK: // %bb.0:
``````````
</details>
https://github.com/llvm/llvm-project/pull/140895
More information about the llvm-commits
mailing list