[llvm] [PAC] Sign LR with B key for non-leaf functions with ptrauth-returns attr (PR #100552)
Daniil Kovalev via llvm-commits
llvm-commits at lists.llvm.org
Thu Jul 25 11:33:48 PDT 2024
https://github.com/kovdan01 updated https://github.com/llvm/llvm-project/pull/100552
>From 086d761d73925af982748d17ff2756d0d7d4cb93 Mon Sep 17 00:00:00 2001
From: Daniil Kovalev <dkovalev at accesssoftek.com>
Date: Thu, 25 Jul 2024 10:42:22 +0300
Subject: [PATCH] [PAC] Sign LR with B key for non-leaf functions with
ptrauth-returns attr
For pauthtest ABI and, there is a bunch of ptrauth-* options, including
ptrauth-returns. Use "ptrauth-returns" function attribute to indicate need
for LR signing with B key for non-leaf function to avoid using
"sign-return-address" and "sign-return-address-key" which were
originally designed for pac-ret.
Co-authored-by: Ahmed Bougacha <ahmed at bougacha.org>
Co-authored-by: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
---
llvm/lib/Target/AArch64/AArch64InstrInfo.cpp | 7 +-
.../AArch64/AArch64MachineFunctionInfo.cpp | 4 +
.../lib/Target/AArch64/AArch64PointerAuth.cpp | 3 +-
llvm/lib/Target/AArch64/AArch64Subtarget.cpp | 9 +-
llvm/lib/Target/AArch64/AArch64Subtarget.h | 3 +-
llvm/test/CodeGen/AArch64/ptrauth-ret-trap.ll | 98 ++++++++
llvm/test/CodeGen/AArch64/ptrauth-ret.ll | 225 ++++++++++++++++++
7 files changed, 344 insertions(+), 5 deletions(-)
create mode 100644 llvm/test/CodeGen/AArch64/ptrauth-ret-trap.ll
create mode 100644 llvm/test/CodeGen/AArch64/ptrauth-ret.ll
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 702c2831e013f..6cd9a1a817086 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -8340,7 +8340,8 @@ AArch64InstrInfo::getOutliningCandidateInfo(
NumBytesToCreateFrame += 8;
// PAuth is enabled - set extra tail call cost, if any.
- auto LRCheckMethod = Subtarget.getAuthenticatedLRCheckMethod();
+ auto LRCheckMethod = Subtarget.getAuthenticatedLRCheckMethod(
+ *RepeatedSequenceLocs[0].getMF());
NumBytesToCheckLRInTCEpilogue =
AArch64PAuth::getCheckerSizeInBytes(LRCheckMethod);
// Checking the authenticated LR value may significantly impact
@@ -8701,6 +8702,10 @@ void AArch64InstrInfo::mergeOutliningCandidateAttributes(
// behaviour of one of them
const auto &CFn = Candidates.front().getMF()->getFunction();
+ if (CFn.hasFnAttribute("ptrauth-returns"))
+ F.addFnAttr(CFn.getFnAttribute("ptrauth-returns"));
+ if (CFn.hasFnAttribute("ptrauth-auth-traps"))
+ F.addFnAttr(CFn.getFnAttribute("ptrauth-auth-traps"));
// Since all candidates belong to the same module, just copy the
// function-level attributes of an arbitrary function.
if (CFn.hasFnAttribute("sign-return-address"))
diff --git a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp
index 201e8047b3686..e96c5a953ff2b 100644
--- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp
@@ -38,6 +38,8 @@ void AArch64FunctionInfo::initializeBaseYamlFields(
}
static std::pair<bool, bool> GetSignReturnAddress(const Function &F) {
+ if (F.hasFnAttribute("ptrauth-returns"))
+ return {true, false}; // non-leaf
// The function should be signed in the following situations:
// - sign-return-address=all
// - sign-return-address=non-leaf and the functions spills the LR
@@ -56,6 +58,8 @@ static std::pair<bool, bool> GetSignReturnAddress(const Function &F) {
}
static bool ShouldSignWithBKey(const Function &F, const AArch64Subtarget &STI) {
+ if (F.hasFnAttribute("ptrauth-returns"))
+ return true;
if (!F.hasFnAttribute("sign-return-address-key")) {
if (STI.getTargetTriple().isOSWindows())
return true;
diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 465e689d4a7a5..92ab4b5c3d251 100644
--- a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
+++ b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
@@ -341,7 +341,8 @@ bool AArch64PointerAuth::checkAuthenticatedLR(
AArch64PACKey::ID KeyId =
MFnI->shouldSignWithBKey() ? AArch64PACKey::IB : AArch64PACKey::IA;
- AuthCheckMethod Method = Subtarget->getAuthenticatedLRCheckMethod();
+ AuthCheckMethod Method =
+ Subtarget->getAuthenticatedLRCheckMethod(*TI->getMF());
if (Method == AuthCheckMethod::None)
return false;
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
index 32a355fe38f1c..642006e706c13 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
@@ -565,8 +565,13 @@ bool AArch64Subtarget::useAA() const { return UseAA; }
// exception on its own. Later, if the callee spills the signed LR value and
// neither FEAT_PAuth2 nor FEAT_EPAC are implemented, the valid PAC replaces
// the higher bits of LR thus hiding the authentication failure.
-AArch64PAuth::AuthCheckMethod
-AArch64Subtarget::getAuthenticatedLRCheckMethod() const {
+AArch64PAuth::AuthCheckMethod AArch64Subtarget::getAuthenticatedLRCheckMethod(
+ const MachineFunction &MF) const {
+ // TODO: Check subtarget for the scheme. Present variant is a default for
+ // pauthtest ABI.
+ if (MF.getFunction().hasFnAttribute("ptrauth-returns") &&
+ MF.getFunction().hasFnAttribute("ptrauth-auth-traps"))
+ return AArch64PAuth::AuthCheckMethod::HighBitsNoTBI;
if (AuthenticatedLRCheckMethod.getNumOccurrences())
return AuthenticatedLRCheckMethod;
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.h b/llvm/lib/Target/AArch64/AArch64Subtarget.h
index e585aad2f7a68..0f3a637f98fbe 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.h
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.h
@@ -413,7 +413,8 @@ class AArch64Subtarget final : public AArch64GenSubtargetInfo {
}
/// Choose a method of checking LR before performing a tail call.
- AArch64PAuth::AuthCheckMethod getAuthenticatedLRCheckMethod() const;
+ AArch64PAuth::AuthCheckMethod
+ getAuthenticatedLRCheckMethod(const MachineFunction &MF) const;
/// Compute the integer discriminator for a given BlockAddress constant, if
/// blockaddress signing is enabled, or std::nullopt otherwise.
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-ret-trap.ll b/llvm/test/CodeGen/AArch64/ptrauth-ret-trap.ll
new file mode 100644
index 0000000000000..42a3050eda112
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/ptrauth-ret-trap.ll
@@ -0,0 +1,98 @@
+; RUN: llc -mtriple aarch64-linux-gnu -mattr=+pauth -asm-verbose=false -disable-post-ra -o - %s | FileCheck %s
+
+; CHECK-LABEL: test_tailcall:
+; CHECK-NEXT: pacibsp
+; CHECK-NEXT: str x30, [sp, #-16]!
+; CHECK-NEXT: bl bar
+; CHECK-NEXT: ldr x30, [sp], #16
+; CHECK-NEXT: autibsp
+; CHECK-NEXT: eor x16, x30, x30, lsl #1
+; CHECK-NEXT: tbnz x16, #62, [[BAD:.L.*]]
+; CHECK-NEXT: b bar
+; CHECK-NEXT: [[BAD]]:
+; CHECK-NEXT: brk #0xc471
+define i32 @test_tailcall() #0 {
+ call i32 @bar()
+ %c = tail call i32 @bar()
+ ret i32 %c
+}
+
+; CHECK-LABEL: test_tailcall_noframe:
+; CHECK-NEXT: b bar
+define i32 @test_tailcall_noframe() #0 {
+ %c = tail call i32 @bar()
+ ret i32 %c
+}
+
+; CHECK-LABEL: test_tailcall_indirect:
+; CHECK: autibsp
+; CHECK: eor x16, x30, x30, lsl #1
+; CHECK: tbnz x16, #62, [[BAD:.L.*]]
+; CHECK: br x0
+; CHECK: [[BAD]]:
+; CHECK: brk #0xc471
+define void @test_tailcall_indirect(ptr %fptr) #0 {
+ call i32 @test_tailcall()
+ tail call void %fptr()
+ ret void
+}
+
+; CHECK-LABEL: test_tailcall_indirect_in_x9:
+; CHECK: autibsp
+; CHECK: eor x16, x30, x30, lsl #1
+; CHECK: tbnz x16, #62, [[BAD:.L.*]]
+; CHECK: br x9
+; CHECK: [[BAD]]:
+; CHECK: brk #0xc471
+define void @test_tailcall_indirect_in_x9(ptr sret(i64) %ret, [8 x i64] %in, ptr %fptr) #0 {
+ %ptr = alloca i8, i32 16
+ call i32 @test_tailcall()
+ tail call void %fptr(ptr sret(i64) %ret, [8 x i64] %in)
+ ret void
+}
+
+; CHECK-LABEL: test_auth_tailcall_indirect:
+; CHECK: autibsp
+; CHECK: eor x16, x30, x30, lsl #1
+; CHECK: tbnz x16, #62, [[BAD:.L.*]]
+; CHECK: mov x16, #42
+; CHECK: braa x0, x16
+; CHECK: [[BAD]]:
+; CHECK: brk #0xc471
+define void @test_auth_tailcall_indirect(ptr %fptr) #0 {
+ call i32 @test_tailcall()
+ tail call void %fptr() [ "ptrauth"(i32 0, i64 42) ]
+ ret void
+}
+
+; CHECK-LABEL: test_auth_tailcall_indirect_in_x9:
+; CHECK: autibsp
+; CHECK: eor x16, x30, x30, lsl #1
+; CHECK: tbnz x16, #62, [[BAD:.L.*]]
+; CHECK: brabz x9
+; CHECK: [[BAD]]:
+; CHECK: brk #0xc471
+define void @test_auth_tailcall_indirect_in_x9(ptr sret(i64) %ret, [8 x i64] %in, ptr %fptr) #0 {
+ %ptr = alloca i8, i32 16
+ call i32 @test_tailcall()
+ tail call void %fptr(ptr sret(i64) %ret, [8 x i64] %in) [ "ptrauth"(i32 1, i64 0) ]
+ ret void
+}
+
+; CHECK-LABEL: test_auth_tailcall_indirect_bti:
+; CHECK: autibsp
+; CHECK: eor x17, x30, x30, lsl #1
+; CHECK: tbnz x17, #62, [[BAD:.L.*]]
+; CHECK: brabz x16
+; CHECK: [[BAD]]:
+; CHECK: brk #0xc471
+define void @test_auth_tailcall_indirect_bti(ptr sret(i64) %ret, [8 x i64] %in, ptr %fptr) #0 "branch-target-enforcement"="true" {
+ %ptr = alloca i8, i32 16
+ call i32 @test_tailcall()
+ tail call void %fptr(ptr sret(i64) %ret, [8 x i64] %in) [ "ptrauth"(i32 1, i64 0) ]
+ ret void
+}
+
+declare i32 @bar()
+
+attributes #0 = { nounwind "ptrauth-returns" "ptrauth-auth-traps" }
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-ret.ll b/llvm/test/CodeGen/AArch64/ptrauth-ret.ll
new file mode 100644
index 0000000000000..61f5f6d9d23b7
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/ptrauth-ret.ll
@@ -0,0 +1,225 @@
+; RUN: llc < %s -mtriple aarch64-linux-gnu -mattr=+pauth -verify-machineinstrs -disable-post-ra \
+; RUN: -global-isel=0 -o - %s | FileCheck %s
+; RUN: llc < %s -mtriple aarch64-linux-gnu -mattr=+pauth -verify-machineinstrs -disable-post-ra \
+; RUN: -global-isel=1 -global-isel-abort=1 -o - %s | FileCheck %s
+
+define i32 @test() #0 {
+; CHECK-LABEL: test:
+; CHECK: %bb.0:
+; CHECK-NEXT: str x19, [sp, #-16]!
+; CHECK-NEXT: mov w0, wzr
+; CHECK-NEXT: //APP
+; CHECK-NEXT: //NO_APP
+; CHECK-NEXT: ldr x19, [sp], #16
+; CHECK-NEXT: ret
+ call void asm sideeffect "", "~{x19}"()
+ ret i32 0
+}
+
+define i32 @test_alloca() #0 {
+; CHECK-LABEL: test_alloca:
+; CHECK: %bb.0:
+; CHECK-NEXT: sub sp, sp, #32
+; CHECK-NEXT: mov x8, sp
+; CHECK-NEXT: mov w0, wzr
+; CHECK-NEXT: //APP
+; CHECK-NEXT: //NO_APP
+; CHECK-NEXT: add sp, sp, #32
+; CHECK-NEXT: ret
+ %p = alloca i8, i32 32
+ call void asm sideeffect "", "r"(ptr %p)
+ ret i32 0
+}
+
+define i32 @test_realign_alloca() #0 {
+; CHECK-LABEL: test_realign_alloca:
+; CHECK: %bb.0:
+; CHECK-NEXT: pacibsp
+; CHECK-NEXT: stp x29, x30, [sp, #-16]!
+; CHECK-NEXT: mov x29, sp
+; CHECK-NEXT: sub x9, sp, #112
+; CHECK-NEXT: and sp, x9, #0xffffffffffffff80
+; CHECK-NEXT: mov x8, sp
+; CHECK-NEXT: mov w0, wzr
+; CHECK-NEXT: //APP
+; CHECK-NEXT: //NO_APP
+; CHECK-NEXT: mov sp, x29
+; CHECK-NEXT: ldp x29, x30, [sp], #16
+; CHECK-NEXT: retab
+ %p = alloca i8, i32 32, align 128
+ call void asm sideeffect "", "r"(ptr %p)
+ ret i32 0
+}
+
+define i32 @test_big_alloca() #0 {
+; CHECK-LABEL: test_big_alloca:
+; CHECK: %bb.0:
+; CHECK-NEXT: str x29, [sp, #-16]!
+; CHECK-NEXT: sub sp, sp, #1024
+; CHECK-NEXT: mov x8, sp
+; CHECK-NEXT: mov w0, wzr
+; CHECK-NEXT: //APP
+; CHECK-NEXT: //NO_APP
+; CHECK-NEXT: add sp, sp, #1024
+; CHECK-NEXT: ldr x29, [sp], #16
+; CHECK-NEXT: ret
+ %p = alloca i8, i32 1024
+ call void asm sideeffect "", "r"(ptr %p)
+ ret i32 0
+}
+
+define i32 @test_var_alloca(i32 %s) #0 {
+ %p = alloca i8, i32 %s
+ call void asm sideeffect "", "r"(ptr %p)
+ ret i32 0
+}
+
+define i32 @test_noframe_saved(ptr %p) #0 {
+; CHECK-LABEL: test_noframe_saved:
+; CHECK: %bb.0:
+
+
+; CHECK-NEXT: str x29, [sp, #-96]!
+; CHECK-NEXT: stp x28, x27, [sp, #16]
+; CHECK-NEXT: stp x26, x25, [sp, #32]
+; CHECK-NEXT: stp x24, x23, [sp, #48]
+; CHECK-NEXT: stp x22, x21, [sp, #64]
+; CHECK-NEXT: stp x20, x19, [sp, #80]
+; CHECK-NEXT: ldr w29, [x0]
+; CHECK-NEXT: //APP
+; CHECK-NEXT: //NO_APP
+; CHECK-NEXT: mov w0, w29
+; CHECK-NEXT: ldp x20, x19, [sp, #80]
+; CHECK-NEXT: ldp x22, x21, [sp, #64]
+; CHECK-NEXT: ldp x24, x23, [sp, #48]
+; CHECK-NEXT: ldp x26, x25, [sp, #32]
+; CHECK-NEXT: ldp x28, x27, [sp, #16]
+; CHECK-NEXT: ldr x29, [sp], #96
+; CHECK-NEXT: ret
+ %v = load i32, ptr %p
+ call void asm sideeffect "", "~{x0},~{x1},~{x2},~{x3},~{x4},~{x5},~{x6},~{x7},~{x8},~{x9},~{x10},~{x11},~{x12},~{x13},~{x14},~{x15},~{x16},~{x17},~{x18},~{x19},~{x20},~{x21},~{x22},~{x23},~{x24},~{x25},~{x26},~{x27},~{x28}"()
+ ret i32 %v
+}
+
+define void @test_noframe() #0 {
+; CHECK-LABEL: test_noframe:
+; CHECK: %bb.0:
+; CHECK-NEXT: ret
+ ret void
+}
+
+; FIXME: Inefficient lowering of @llvm.returnaddress
+define ptr @test_returnaddress_0() #0 {
+; CHECK-LABEL: test_returnaddress_0:
+; CHECK: %bb.0:
+; CHECK-NEXT: pacibsp
+; CHECK-NEXT: str x30, [sp, #-16]!
+; CHECK-NEXT: xpaci x30
+; CHECK-NEXT: mov x0, x30
+; CHECK-NEXT: ldr x30, [sp], #16
+; CHECK-NEXT: retab
+ %r = call ptr @llvm.returnaddress(i32 0)
+ ret ptr %r
+}
+
+define ptr @test_returnaddress_1() #0 {
+; CHECK-LABEL: test_returnaddress_1:
+; CHECK: %bb.0:
+; CHECK-NEXT: pacibsp
+; CHECK-NEXT: stp x29, x30, [sp, #-16]!
+; CHECK-NEXT: mov x29, sp
+; CHECK-NEXT: ldr x8, [x29]
+; CHECK-NEXT: ldr x0, [x8, #8]
+; CHECK-NEXT: xpaci x0
+; CHECK-NEXT: ldp x29, x30, [sp], #16
+; CHECK-NEXT: retab
+ %r = call ptr @llvm.returnaddress(i32 1)
+ ret ptr %r
+}
+
+define void @test_noframe_alloca() #0 {
+; CHECK-LABEL: test_noframe_alloca:
+; CHECK: %bb.0:
+; CHECK-NEXT: sub sp, sp, #16
+; CHECK-NEXT: add x8, sp, #12
+; CHECK-NEXT: //APP
+; CHECK-NEXT: //NO_APP
+; CHECK-NEXT: add sp, sp, #16
+; CHECK-NEXT: ret
+ %p = alloca i8, i32 1
+ call void asm sideeffect "", "r"(ptr %p)
+ ret void
+}
+
+define void @test_call() #0 {
+; CHECK-LABEL: test_call:
+; CHECK: %bb.0:
+; CHECK-NEXT: pacibsp
+; CHECK-NEXT: str x30, [sp, #-16]!
+; CHECK-NEXT: bl bar
+; CHECK-NEXT: ldr x30, [sp], #16
+; CHECK-NEXT: retab
+ call i32 @bar()
+ ret void
+}
+
+define void @test_call_alloca() #0 {
+; CHECK-LABEL: test_call_alloca:
+; CHECK: %bb.0:
+; CHECK-NEXT: pacibsp
+; CHECK-NEXT: str x30, [sp, #-16]
+; CHECK-NEXT: bl bar
+; CHECK-NEXT: ldr x30, [sp], #16
+; CHECK-NEXT: retab
+ alloca i8
+ call i32 @bar()
+ ret void
+}
+
+define void @test_call_shrinkwrapping(i1 %c) #0 {
+; CHECK-LABEL: test_call_shrinkwrapping:
+; CHECK: %bb.0:
+; CHECK-NEXT: tbz w0, #0, .LBB12_2
+; CHECK-NEXT: %bb.1:
+; CHECK-NEXT: pacibsp
+; CHECK-NEXT: str x30, [sp, #-16]!
+; CHECK-NEXT: bl bar
+; CHECK-NEXT: ldr x30, [sp], #16
+; CHECK-NEXT: autibsp
+; CHECK-NEXT: LBB12_2:
+; CHECK-NEXT: ret
+ br i1 %c, label %tbb, label %fbb
+tbb:
+ call i32 @bar()
+ br label %fbb
+fbb:
+ ret void
+}
+
+define i32 @test_tailcall() #0 {
+; CHECK-LABEL: test_tailcall:
+; CHECK: %bb.0:
+; CHECK-NEXT: pacibsp
+; CHECK-NEXT: str x30, [sp, #-16]!
+; CHECK-NEXT: bl bar
+; CHECK-NEXT: ldr x30, [sp], #16
+; CHECK-NEXT: autibsp
+; CHECK-NEXT: b bar
+ call i32 @bar()
+ %c = tail call i32 @bar()
+ ret i32 %c
+}
+
+define i32 @test_tailcall_noframe() #0 {
+; CHECK-LABEL: test_tailcall_noframe:
+; CHECK: %bb.0:
+; CHECK-NEXT: b bar
+ %c = tail call i32 @bar()
+ ret i32 %c
+}
+
+declare i32 @bar()
+
+declare ptr @llvm.returnaddress(i32)
+
+attributes #0 = { nounwind "ptrauth-returns" }
More information about the llvm-commits
mailing list