Skip to content

Commit 56fd247

Browse files
kovdan01ahmedbougachaatrosinenko
authored
[PAC] Sign LR with B key for non-leaf functions with ptrauth-returns attr (#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 <[email protected]> Co-authored-by: Anatoly Trosinenko <[email protected]>
1 parent 81ce796 commit 56fd247

File tree

7 files changed

+344
-5
lines changed

7 files changed

+344
-5
lines changed

‎llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8340,7 +8340,8 @@ AArch64InstrInfo::getOutliningCandidateInfo(
83408340
NumBytesToCreateFrame += 8;
83418341

83428342
// PAuth is enabled - set extra tail call cost, if any.
8343-
auto LRCheckMethod = Subtarget.getAuthenticatedLRCheckMethod();
8343+
auto LRCheckMethod = Subtarget.getAuthenticatedLRCheckMethod(
8344+
*RepeatedSequenceLocs[0].getMF());
83448345
NumBytesToCheckLRInTCEpilogue =
83458346
AArch64PAuth::getCheckerSizeInBytes(LRCheckMethod);
83468347
// Checking the authenticated LR value may significantly impact
@@ -8701,6 +8702,10 @@ void AArch64InstrInfo::mergeOutliningCandidateAttributes(
87018702
// behaviour of one of them
87028703
const auto &CFn = Candidates.front().getMF()->getFunction();
87038704

8705+
if (CFn.hasFnAttribute("ptrauth-returns"))
8706+
F.addFnAttr(CFn.getFnAttribute("ptrauth-returns"));
8707+
if (CFn.hasFnAttribute("ptrauth-auth-traps"))
8708+
F.addFnAttr(CFn.getFnAttribute("ptrauth-auth-traps"));
87048709
// Since all candidates belong to the same module, just copy the
87058710
// function-level attributes of an arbitrary function.
87068711
if (CFn.hasFnAttribute("sign-return-address"))

‎llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ void AArch64FunctionInfo::initializeBaseYamlFields(
3838
}
3939

4040
static std::pair<bool, bool> GetSignReturnAddress(const Function &F) {
41+
if (F.hasFnAttribute("ptrauth-returns"))
42+
return {true, false}; // non-leaf
4143
// The function should be signed in the following situations:
4244
// - sign-return-address=all
4345
// - sign-return-address=non-leaf and the functions spills the LR
@@ -56,6 +58,8 @@ static std::pair<bool, bool> GetSignReturnAddress(const Function &F) {
5658
}
5759

5860
static bool ShouldSignWithBKey(const Function &F, const AArch64Subtarget &STI) {
61+
if (F.hasFnAttribute("ptrauth-returns"))
62+
return true;
5963
if (!F.hasFnAttribute("sign-return-address-key")) {
6064
if (STI.getTargetTriple().isOSWindows())
6165
return true;

‎llvm/lib/Target/AArch64/AArch64PointerAuth.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -341,7 +341,8 @@ bool AArch64PointerAuth::checkAuthenticatedLR(
341341
AArch64PACKey::ID KeyId =
342342
MFnI->shouldSignWithBKey() ? AArch64PACKey::IB : AArch64PACKey::IA;
343343

344-
AuthCheckMethod Method = Subtarget->getAuthenticatedLRCheckMethod();
344+
AuthCheckMethod Method =
345+
Subtarget->getAuthenticatedLRCheckMethod(*TI->getMF());
345346

346347
if (Method == AuthCheckMethod::None)
347348
return false;

‎llvm/lib/Target/AArch64/AArch64Subtarget.cpp

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -565,8 +565,13 @@ bool AArch64Subtarget::useAA() const { return UseAA; }
565565
// exception on its own. Later, if the callee spills the signed LR value and
566566
// neither FEAT_PAuth2 nor FEAT_EPAC are implemented, the valid PAC replaces
567567
// the higher bits of LR thus hiding the authentication failure.
568-
AArch64PAuth::AuthCheckMethod
569-
AArch64Subtarget::getAuthenticatedLRCheckMethod() const {
568+
AArch64PAuth::AuthCheckMethod AArch64Subtarget::getAuthenticatedLRCheckMethod(
569+
const MachineFunction &MF) const {
570+
// TODO: Check subtarget for the scheme. Present variant is a default for
571+
// pauthtest ABI.
572+
if (MF.getFunction().hasFnAttribute("ptrauth-returns") &&
573+
MF.getFunction().hasFnAttribute("ptrauth-auth-traps"))
574+
return AArch64PAuth::AuthCheckMethod::HighBitsNoTBI;
570575
if (AuthenticatedLRCheckMethod.getNumOccurrences())
571576
return AuthenticatedLRCheckMethod;
572577

‎llvm/lib/Target/AArch64/AArch64Subtarget.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -413,7 +413,8 @@ class AArch64Subtarget final : public AArch64GenSubtargetInfo {
413413
}
414414

415415
/// Choose a method of checking LR before performing a tail call.
416-
AArch64PAuth::AuthCheckMethod getAuthenticatedLRCheckMethod() const;
416+
AArch64PAuth::AuthCheckMethod
417+
getAuthenticatedLRCheckMethod(const MachineFunction &MF) const;
417418

418419
/// Compute the integer discriminator for a given BlockAddress constant, if
419420
/// blockaddress signing is enabled, or std::nullopt otherwise.

‎llvm/test/CodeGen/AArch64/ptrauth-ret-trap.ll

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
; RUN: llc -mtriple aarch64-linux-gnu -mattr=+pauth -asm-verbose=false -disable-post-ra -o - %s | FileCheck %s
2+
3+
; CHECK-LABEL: test_tailcall:
4+
; CHECK-NEXT: pacibsp
5+
; CHECK-NEXT: str x30, [sp, #-16]!
6+
; CHECK-NEXT: bl bar
7+
; CHECK-NEXT: ldr x30, [sp], #16
8+
; CHECK-NEXT: autibsp
9+
; CHECK-NEXT: eor x16, x30, x30, lsl #1
10+
; CHECK-NEXT: tbnz x16, #62, [[BAD:.L.*]]
11+
; CHECK-NEXT: b bar
12+
; CHECK-NEXT: [[BAD]]:
13+
; CHECK-NEXT: brk #0xc471
14+
define i32 @test_tailcall() #0 {
15+
call i32 @bar()
16+
%c = tail call i32 @bar()
17+
ret i32 %c
18+
}
19+
20+
; CHECK-LABEL: test_tailcall_noframe:
21+
; CHECK-NEXT: b bar
22+
define i32 @test_tailcall_noframe() #0 {
23+
%c = tail call i32 @bar()
24+
ret i32 %c
25+
}
26+
27+
; CHECK-LABEL: test_tailcall_indirect:
28+
; CHECK: autibsp
29+
; CHECK: eor x16, x30, x30, lsl #1
30+
; CHECK: tbnz x16, #62, [[BAD:.L.*]]
31+
; CHECK: br x0
32+
; CHECK: [[BAD]]:
33+
; CHECK: brk #0xc471
34+
define void @test_tailcall_indirect(ptr %fptr) #0 {
35+
call i32 @test_tailcall()
36+
tail call void %fptr()
37+
ret void
38+
}
39+
40+
; CHECK-LABEL: test_tailcall_indirect_in_x9:
41+
; CHECK: autibsp
42+
; CHECK: eor x16, x30, x30, lsl #1
43+
; CHECK: tbnz x16, #62, [[BAD:.L.*]]
44+
; CHECK: br x9
45+
; CHECK: [[BAD]]:
46+
; CHECK: brk #0xc471
47+
define void @test_tailcall_indirect_in_x9(ptr sret(i64) %ret, [8 x i64] %in, ptr %fptr) #0 {
48+
%ptr = alloca i8, i32 16
49+
call i32 @test_tailcall()
50+
tail call void %fptr(ptr sret(i64) %ret, [8 x i64] %in)
51+
ret void
52+
}
53+
54+
; CHECK-LABEL: test_auth_tailcall_indirect:
55+
; CHECK: autibsp
56+
; CHECK: eor x16, x30, x30, lsl #1
57+
; CHECK: tbnz x16, #62, [[BAD:.L.*]]
58+
; CHECK: mov x16, #42
59+
; CHECK: braa x0, x16
60+
; CHECK: [[BAD]]:
61+
; CHECK: brk #0xc471
62+
define void @test_auth_tailcall_indirect(ptr %fptr) #0 {
63+
call i32 @test_tailcall()
64+
tail call void %fptr() [ "ptrauth"(i32 0, i64 42) ]
65+
ret void
66+
}
67+
68+
; CHECK-LABEL: test_auth_tailcall_indirect_in_x9:
69+
; CHECK: autibsp
70+
; CHECK: eor x16, x30, x30, lsl #1
71+
; CHECK: tbnz x16, #62, [[BAD:.L.*]]
72+
; CHECK: brabz x9
73+
; CHECK: [[BAD]]:
74+
; CHECK: brk #0xc471
75+
define void @test_auth_tailcall_indirect_in_x9(ptr sret(i64) %ret, [8 x i64] %in, ptr %fptr) #0 {
76+
%ptr = alloca i8, i32 16
77+
call i32 @test_tailcall()
78+
tail call void %fptr(ptr sret(i64) %ret, [8 x i64] %in) [ "ptrauth"(i32 1, i64 0) ]
79+
ret void
80+
}
81+
82+
; CHECK-LABEL: test_auth_tailcall_indirect_bti:
83+
; CHECK: autibsp
84+
; CHECK: eor x17, x30, x30, lsl #1
85+
; CHECK: tbnz x17, #62, [[BAD:.L.*]]
86+
; CHECK: brabz x16
87+
; CHECK: [[BAD]]:
88+
; CHECK: brk #0xc471
89+
define void @test_auth_tailcall_indirect_bti(ptr sret(i64) %ret, [8 x i64] %in, ptr %fptr) #0 "branch-target-enforcement"="true" {
90+
%ptr = alloca i8, i32 16
91+
call i32 @test_tailcall()
92+
tail call void %fptr(ptr sret(i64) %ret, [8 x i64] %in) [ "ptrauth"(i32 1, i64 0) ]
93+
ret void
94+
}
95+
96+
declare i32 @bar()
97+
98+
attributes #0 = { nounwind "ptrauth-returns" "ptrauth-auth-traps" }

‎llvm/test/CodeGen/AArch64/ptrauth-ret.ll

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
1+
; RUN: llc < %s -mtriple aarch64-linux-gnu -mattr=+pauth -verify-machineinstrs -disable-post-ra \
2+
; RUN: -global-isel=0 -o - %s | FileCheck %s
3+
; RUN: llc < %s -mtriple aarch64-linux-gnu -mattr=+pauth -verify-machineinstrs -disable-post-ra \
4+
; RUN: -global-isel=1 -global-isel-abort=1 -o - %s | FileCheck %s
5+
6+
define i32 @test() #0 {
7+
; CHECK-LABEL: test:
8+
; CHECK: %bb.0:
9+
; CHECK-NEXT: str x19, [sp, #-16]!
10+
; CHECK-NEXT: mov w0, wzr
11+
; CHECK-NEXT: //APP
12+
; CHECK-NEXT: //NO_APP
13+
; CHECK-NEXT: ldr x19, [sp], #16
14+
; CHECK-NEXT: ret
15+
call void asm sideeffect "", "~{x19}"()
16+
ret i32 0
17+
}
18+
19+
define i32 @test_alloca() #0 {
20+
; CHECK-LABEL: test_alloca:
21+
; CHECK: %bb.0:
22+
; CHECK-NEXT: sub sp, sp, #32
23+
; CHECK-NEXT: mov x8, sp
24+
; CHECK-NEXT: mov w0, wzr
25+
; CHECK-NEXT: //APP
26+
; CHECK-NEXT: //NO_APP
27+
; CHECK-NEXT: add sp, sp, #32
28+
; CHECK-NEXT: ret
29+
%p = alloca i8, i32 32
30+
call void asm sideeffect "", "r"(ptr %p)
31+
ret i32 0
32+
}
33+
34+
define i32 @test_realign_alloca() #0 {
35+
; CHECK-LABEL: test_realign_alloca:
36+
; CHECK: %bb.0:
37+
; CHECK-NEXT: pacibsp
38+
; CHECK-NEXT: stp x29, x30, [sp, #-16]!
39+
; CHECK-NEXT: mov x29, sp
40+
; CHECK-NEXT: sub x9, sp, #112
41+
; CHECK-NEXT: and sp, x9, #0xffffffffffffff80
42+
; CHECK-NEXT: mov x8, sp
43+
; CHECK-NEXT: mov w0, wzr
44+
; CHECK-NEXT: //APP
45+
; CHECK-NEXT: //NO_APP
46+
; CHECK-NEXT: mov sp, x29
47+
; CHECK-NEXT: ldp x29, x30, [sp], #16
48+
; CHECK-NEXT: retab
49+
%p = alloca i8, i32 32, align 128
50+
call void asm sideeffect "", "r"(ptr %p)
51+
ret i32 0
52+
}
53+
54+
define i32 @test_big_alloca() #0 {
55+
; CHECK-LABEL: test_big_alloca:
56+
; CHECK: %bb.0:
57+
; CHECK-NEXT: str x29, [sp, #-16]!
58+
; CHECK-NEXT: sub sp, sp, #1024
59+
; CHECK-NEXT: mov x8, sp
60+
; CHECK-NEXT: mov w0, wzr
61+
; CHECK-NEXT: //APP
62+
; CHECK-NEXT: //NO_APP
63+
; CHECK-NEXT: add sp, sp, #1024
64+
; CHECK-NEXT: ldr x29, [sp], #16
65+
; CHECK-NEXT: ret
66+
%p = alloca i8, i32 1024
67+
call void asm sideeffect "", "r"(ptr %p)
68+
ret i32 0
69+
}
70+
71+
define i32 @test_var_alloca(i32 %s) #0 {
72+
%p = alloca i8, i32 %s
73+
call void asm sideeffect "", "r"(ptr %p)
74+
ret i32 0
75+
}
76+
77+
define i32 @test_noframe_saved(ptr %p) #0 {
78+
; CHECK-LABEL: test_noframe_saved:
79+
; CHECK: %bb.0:
80+
81+
82+
; CHECK-NEXT: str x29, [sp, #-96]!
83+
; CHECK-NEXT: stp x28, x27, [sp, #16]
84+
; CHECK-NEXT: stp x26, x25, [sp, #32]
85+
; CHECK-NEXT: stp x24, x23, [sp, #48]
86+
; CHECK-NEXT: stp x22, x21, [sp, #64]
87+
; CHECK-NEXT: stp x20, x19, [sp, #80]
88+
; CHECK-NEXT: ldr w29, [x0]
89+
; CHECK-NEXT: //APP
90+
; CHECK-NEXT: //NO_APP
91+
; CHECK-NEXT: mov w0, w29
92+
; CHECK-NEXT: ldp x20, x19, [sp, #80]
93+
; CHECK-NEXT: ldp x22, x21, [sp, #64]
94+
; CHECK-NEXT: ldp x24, x23, [sp, #48]
95+
; CHECK-NEXT: ldp x26, x25, [sp, #32]
96+
; CHECK-NEXT: ldp x28, x27, [sp, #16]
97+
; CHECK-NEXT: ldr x29, [sp], #96
98+
; CHECK-NEXT: ret
99+
%v = load i32, ptr %p
100+
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}"()
101+
ret i32 %v
102+
}
103+
104+
define void @test_noframe() #0 {
105+
; CHECK-LABEL: test_noframe:
106+
; CHECK: %bb.0:
107+
; CHECK-NEXT: ret
108+
ret void
109+
}
110+
111+
; FIXME: Inefficient lowering of @llvm.returnaddress
112+
define ptr @test_returnaddress_0() #0 {
113+
; CHECK-LABEL: test_returnaddress_0:
114+
; CHECK: %bb.0:
115+
; CHECK-NEXT: pacibsp
116+
; CHECK-NEXT: str x30, [sp, #-16]!
117+
; CHECK-NEXT: xpaci x30
118+
; CHECK-NEXT: mov x0, x30
119+
; CHECK-NEXT: ldr x30, [sp], #16
120+
; CHECK-NEXT: retab
121+
%r = call ptr @llvm.returnaddress(i32 0)
122+
ret ptr %r
123+
}
124+
125+
define ptr @test_returnaddress_1() #0 {
126+
; CHECK-LABEL: test_returnaddress_1:
127+
; CHECK: %bb.0:
128+
; CHECK-NEXT: pacibsp
129+
; CHECK-NEXT: stp x29, x30, [sp, #-16]!
130+
; CHECK-NEXT: mov x29, sp
131+
; CHECK-NEXT: ldr x8, [x29]
132+
; CHECK-NEXT: ldr x0, [x8, #8]
133+
; CHECK-NEXT: xpaci x0
134+
; CHECK-NEXT: ldp x29, x30, [sp], #16
135+
; CHECK-NEXT: retab
136+
%r = call ptr @llvm.returnaddress(i32 1)
137+
ret ptr %r
138+
}
139+
140+
define void @test_noframe_alloca() #0 {
141+
; CHECK-LABEL: test_noframe_alloca:
142+
; CHECK: %bb.0:
143+
; CHECK-NEXT: sub sp, sp, #16
144+
; CHECK-NEXT: add x8, sp, #12
145+
; CHECK-NEXT: //APP
146+
; CHECK-NEXT: //NO_APP
147+
; CHECK-NEXT: add sp, sp, #16
148+
; CHECK-NEXT: ret
149+
%p = alloca i8, i32 1
150+
call void asm sideeffect "", "r"(ptr %p)
151+
ret void
152+
}
153+
154+
define void @test_call() #0 {
155+
; CHECK-LABEL: test_call:
156+
; CHECK: %bb.0:
157+
; CHECK-NEXT: pacibsp
158+
; CHECK-NEXT: str x30, [sp, #-16]!
159+
; CHECK-NEXT: bl bar
160+
; CHECK-NEXT: ldr x30, [sp], #16
161+
; CHECK-NEXT: retab
162+
call i32 @bar()
163+
ret void
164+
}
165+
166+
define void @test_call_alloca() #0 {
167+
; CHECK-LABEL: test_call_alloca:
168+
; CHECK: %bb.0:
169+
; CHECK-NEXT: pacibsp
170+
; CHECK-NEXT: str x30, [sp, #-16]
171+
; CHECK-NEXT: bl bar
172+
; CHECK-NEXT: ldr x30, [sp], #16
173+
; CHECK-NEXT: retab
174+
alloca i8
175+
call i32 @bar()
176+
ret void
177+
}
178+
179+
define void @test_call_shrinkwrapping(i1 %c) #0 {
180+
; CHECK-LABEL: test_call_shrinkwrapping:
181+
; CHECK: %bb.0:
182+
; CHECK-NEXT: tbz w0, #0, .LBB12_2
183+
; CHECK-NEXT: %bb.1:
184+
; CHECK-NEXT: pacibsp
185+
; CHECK-NEXT: str x30, [sp, #-16]!
186+
; CHECK-NEXT: bl bar
187+
; CHECK-NEXT: ldr x30, [sp], #16
188+
; CHECK-NEXT: autibsp
189+
; CHECK-NEXT: LBB12_2:
190+
; CHECK-NEXT: ret
191+
br i1 %c, label %tbb, label %fbb
192+
tbb:
193+
call i32 @bar()
194+
br label %fbb
195+
fbb:
196+
ret void
197+
}
198+
199+
define i32 @test_tailcall() #0 {
200+
; CHECK-LABEL: test_tailcall:
201+
; CHECK: %bb.0:
202+
; CHECK-NEXT: pacibsp
203+
; CHECK-NEXT: str x30, [sp, #-16]!
204+
; CHECK-NEXT: bl bar
205+
; CHECK-NEXT: ldr x30, [sp], #16
206+
; CHECK-NEXT: autibsp
207+
; CHECK-NEXT: b bar
208+
call i32 @bar()
209+
%c = tail call i32 @bar()
210+
ret i32 %c
211+
}
212+
213+
define i32 @test_tailcall_noframe() #0 {
214+
; CHECK-LABEL: test_tailcall_noframe:
215+
; CHECK: %bb.0:
216+
; CHECK-NEXT: b bar
217+
%c = tail call i32 @bar()
218+
ret i32 %c
219+
}
220+
221+
declare i32 @bar()
222+
223+
declare ptr @llvm.returnaddress(i32)
224+
225+
attributes #0 = { nounwind "ptrauth-returns" }

0 commit comments

Comments
 (0)