Skip to content

Commit

Permalink
[PAC] Sign LR with B key for non-leaf functions with ptrauth-returns …
Browse files Browse the repository at this point in the history 
…attr (llvm#100552)

For pauthtest ABI, 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@bougacha.org>
Co-authored-by: Anatoly Trosinenko <atrosinenko@accesssoftek.com>
(cherry picked from commit 56fd247)
  • Loading branch information
@kovdan01 @llvmbot
kovdan01 authored and llvmbot committed Jul 26, 2024
1 parent 61b0a2f commit 58f851d
Show file tree
Hide file tree
Showing 7 changed files with 344 additions and 5 deletions.
7 changes: 6 additions & 1 deletion llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8339,7 +8339,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
Expand Down Expand Up @@ -8700,6 +8701,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"))
Expand Down
4 changes: 4 additions & 0 deletions llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
Expand Down
3 changes: 2 additions & 1 deletion llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
9 changes: 7 additions & 2 deletions llvm/lib/Target/AArch64/AArch64Subtarget.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
3 changes: 2 additions & 1 deletion llvm/lib/Target/AArch64/AArch64Subtarget.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
98 changes: 98 additions & 0 deletions llvm/test/CodeGen/AArch64/ptrauth-ret-trap.ll
View file
Original file line number Diff line number Diff line change
@@ -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" }
225 changes: 225 additions & 0 deletions llvm/test/CodeGen/AArch64/ptrauth-ret.ll
View file
Original file line number Diff line number Diff line change
@@ -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" }

0 comments on commit 58f851d

Please sign in to comment.