Skip to content

Commit c7f3bdb

Browse files
atrosinenkoatrosinenko
authored
[AArch64][PAC] Fix clobbering registers by BLRA and AUTH_TCRETURN (llvm#155373)
After `isX16X17Safer` predicate was introduced, it became possible for `emitPtrauthDiscriminator` function to return its `AddrDisc` argument when it is neither X16 nor X17 (which are declared as implicit-def'ed by `BLRA` and `AUTH_TCRETURN[_BTI]` pseudo instructions). This resulted in the above pseudos being able to accidentally clobber unexpected register. As a quick fix for miscompilation possibility, this patch virtually restores the old behavior for the affected pseudo instructions. It should be possible to improve the efficiency via subsequent patches by accounting for `killed` flags and register masks.
1 parent c3b2530 commit c7f3bdb

File tree

2 files changed

+120
-9
lines changed

2 files changed

+120
-9
lines changed

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

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2229,13 +2229,24 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
22292229
if (BrTarget == AddrDisc)
22302230
report_fatal_error("Branch target is signed with its own value");
22312231

2232-
// If we are printing BLRA pseudo instruction, then x16 and x17 are
2233-
// implicit-def'ed by the MI and AddrDisc is not used as any other input, so
2234-
// try to save one MOV by setting MayUseAddrAsScratch.
2232+
// If we are printing BLRA pseudo, try to save one MOV by making use of the
2233+
// fact that x16 and x17 are described as clobbered by the MI instruction and
2234+
// AddrDisc is not used as any other input.
2235+
//
2236+
// Back in the day, emitPtrauthDiscriminator was restricted to only returning
2237+
// either x16 or x17, meaning the returned register is always among the
2238+
// implicit-def'ed registers of BLRA pseudo. Now this property can be violated
2239+
// if isX16X17Safer predicate is false, thus manually check if AddrDisc is
2240+
// among x16 and x17 to prevent clobbering unexpected registers.
2241+
//
22352242
// Unlike BLRA, BRA pseudo is used to perform computed goto, and thus not
22362243
// declared as clobbering x16/x17.
2244+
//
2245+
// FIXME: Make use of `killed` flags and register masks instead.
2246+
bool AddrDiscIsImplicitDef =
2247+
IsCall && (AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17);
22372248
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17,
2238-
/*MayUseAddrAsScratch=*/IsCall);
2249+
AddrDiscIsImplicitDef);
22392250
bool IsZeroDisc = DiscReg == AArch64::XZR;
22402251

22412252
unsigned Opc;
@@ -2968,8 +2979,15 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
29682979
// See the comments in emitPtrauthBranch.
29692980
if (Callee == AddrDisc)
29702981
report_fatal_error("Call target is signed with its own value");
2982+
2983+
// After isX16X17Safer predicate was introduced, emitPtrauthDiscriminator is
2984+
// no longer restricted to only reusing AddrDisc when it is X16 or X17
2985+
// (which are implicit-def'ed by AUTH_TCRETURN pseudos), thus impose this
2986+
// restriction manually not to clobber an unexpected register.
2987+
bool AddrDiscIsImplicitDef =
2988+
AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17;
29712989
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, ScratchReg,
2972-
/*MayUseAddrAsScratch=*/true);
2990+
AddrDiscIsImplicitDef);
29732991

29742992
const bool IsZero = DiscReg == AArch64::XZR;
29752993
const unsigned Opcodes[2][2] = {{AArch64::BRAA, AArch64::BRAAZ},

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

Lines changed: 97 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -181,8 +181,9 @@ define void @test_tailcall_omit_mov_x16_x16(ptr %objptr) #0 {
181181
; ELF-NEXT: movk x8, #6503, lsl #48
182182
; ELF-NEXT: autda x1, x8
183183
; ELF-NEXT: ldr x2, [x1]
184-
; ELF-NEXT: movk x1, #54167, lsl #48
185-
; ELF-NEXT: braa x2, x1
184+
; ELF-NEXT: mov x16, x1
185+
; ELF-NEXT: movk x16, #54167, lsl #48
186+
; ELF-NEXT: braa x2, x16
186187
%vtable.signed = load ptr, ptr %objptr, align 8
187188
%objptr.int = ptrtoint ptr %objptr to i64
188189
%vtable.discr = tail call i64 @llvm.ptrauth.blend(i64 %objptr.int, i64 6503)
@@ -213,8 +214,9 @@ define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
213214
; ELF-NEXT: movk x9, #6503, lsl #48
214215
; ELF-NEXT: autda x8, x9
215216
; ELF-NEXT: ldr x9, [x8]
216-
; ELF-NEXT: movk x8, #34646, lsl #48
217-
; ELF-NEXT: blraa x9, x8
217+
; ELF-NEXT: mov x17, x8
218+
; ELF-NEXT: movk x17, #34646, lsl #48
219+
; ELF-NEXT: blraa x9, x17
218220
; ELF-NEXT: mov w0, #42
219221
; ELF-NEXT: ldr x30, [sp], #16
220222
; CHECK-NEXT: ret
@@ -230,6 +232,97 @@ define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
230232
ret i32 42
231233
}
232234

235+
; The second BLRA instruction should not reuse its AddrDisc operand as a scratch register (returned later).
236+
define i64 @test_call_discr_csr_live(ptr %fnptr, i64 %addr.discr) #0 {
237+
; ELF-LABEL: test_call_discr_csr_live:
238+
; ELF-NEXT: str x30, [sp, #-32]!
239+
; ELF-NEXT: stp x20, x19, [sp, #16]
240+
; ELF-DAG: mov x[[FNPTR:[0-9]+]], x0
241+
; ELF-DAG: mov x[[ADDR_DISC:[0-9]+]], x1
242+
; ELF-DAG: mov x17, x1
243+
; ELF-NEXT: movk x17, #6503, lsl #48
244+
; ELF-NEXT: blraa x0, x17
245+
; ELF-NEXT: mov x17, x[[ADDR_DISC]]
246+
; ELF-NEXT: movk x17, #6503, lsl #48
247+
; ELF-NEXT: blraa x[[FNPTR]], x17
248+
; ELF-NEXT: mov x0, x[[ADDR_DISC]]
249+
; ELF-NEXT: ldp x20, x19, [sp, #16]
250+
; ELF-NEXT: ldr x30, [sp], #32
251+
; ELF-NEXT: ret
252+
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
253+
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
254+
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
255+
ret i64 %addr.discr
256+
}
257+
258+
; The second BLRA instruction may reuse its AddrDisc operand as a scratch register.
259+
define i64 @test_call_discr_csr_killed(ptr %fnptr, i64 %addr.discr) #0 {
260+
; ELF-LABEL: test_call_discr_csr_killed:
261+
; ELF-NEXT: str x30, [sp, #-32]!
262+
; ELF-NEXT: stp x20, x19, [sp, #16]
263+
; ELF-DAG: mov x[[FNPTR:[0-9]+]], x0
264+
; ELF-DAG: mov x[[ADDR_DISC:[0-9]+]], x1
265+
; ELF-DAG: mov x17, x1
266+
; ELF-NEXT: movk x17, #6503, lsl #48
267+
; ELF-NEXT: blraa x0, x17
268+
; ELF-DAG: mov x17, x[[ADDR_DISC]]
269+
; ELF-NEXT: movk x17, #6503, lsl #48
270+
; ELF-NEXT: blraa x[[FNPTR]], x17
271+
; ELF-NEXT: ldp x20, x19, [sp, #16]
272+
; ELF-NEXT: mov w0, #42
273+
; ELF-NEXT: ldr x30, [sp], #32
274+
; ELF-NEXT: ret
275+
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
276+
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
277+
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
278+
ret i64 42
279+
}
280+
281+
; BLRA instruction should not reuse its AddrDisc operand as a scratch register (function argument).
282+
define i64 @test_call_discr_arg(ptr %fnptr, i64 %addr.discr) #0 {
283+
; ELF-LABEL: test_call_discr_arg:
284+
; ELF-NEXT: str x30, [sp, #-16]!
285+
; ELF-NEXT: mov x8, x0
286+
; ELF-NEXT: mov x0, xzr
287+
; ELF-NEXT: mov x17, x1
288+
; ELF-NEXT: movk x17, #6503, lsl #48
289+
; ELF-NEXT: blraa x8, x17
290+
; ELF-NEXT: mov w0, #42
291+
; ELF-NEXT: ldr x30, [sp], #16
292+
; ELF-NEXT: ret
293+
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
294+
tail call void %fnptr(ptr null, i64 %addr.discr) [ "ptrauth"(i32 0, i64 %discr) ]
295+
ret i64 42
296+
}
297+
298+
; BLRA instruction may reuse its AddrDisc operand as a scratch register.
299+
define i64 @test_call_discr_non_arg(ptr %fnptr, i64 %addr.discr) #0 {
300+
; ELF-LABEL: test_call_discr_non_arg:
301+
; ELF-NEXT: str x30, [sp, #-16]!
302+
; ELF-NEXT: mov x17, x1
303+
; ELF-NEXT: movk x17, #6503, lsl #48
304+
; ELF-NEXT: blraa x0, x17
305+
; ELF-NEXT: mov w0, #42
306+
; ELF-NEXT: ldr x30, [sp], #16
307+
; ELF-NEXT: ret
308+
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
309+
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
310+
ret i64 42
311+
}
312+
313+
; AUTH_TCRETURN instruction should not reuse its AddrDisc operand as a scratch register (function argument).
314+
define i64 @test_tailcall_discr_arg(ptr %fnptr, i64 %addr.discr) #0 {
315+
; ELF-LABEL: test_tailcall_discr_arg:
316+
; ELF-NEXT: mov x2, x0
317+
; ELF-NEXT: mov x0, xzr
318+
; ELF-NEXT: mov x16, x1
319+
; ELF-NEXT: movk x16, #6503, lsl #48
320+
; ELF-NEXT: braa x2, x16
321+
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
322+
%result = tail call i64 %fnptr(ptr null, i64 %addr.discr) [ "ptrauth"(i32 0, i64 %discr) ]
323+
ret i64 %result
324+
}
325+
233326
define i32 @test_call_ia_arg(ptr %arg0, i64 %arg1) #0 {
234327
; DARWIN-LABEL: test_call_ia_arg:
235328
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!

0 commit comments

Comments
 (0)