Skip to content

Commit d286efe

Browse files
ahmedbougachaahmedbougacha
authored
[AArch64][PAC] Lower direct authenticated calls to ptrauth constants. (#97664)
This tries to turn indirect ptrauth calls into direct calls, using `ConstantPtrAuth::isKnownEquivalent` to compare the `ConstantPtrAuth` target with the ptrauth call bundle. This should be straightforward, other than the somewhat awkward GISel handling, which has a handshake between CallLowering and IRTranslator to elide the ptrauth when possible.
1 parent ba66d60 commit d286efe

File tree

5 files changed

+363
-11
lines changed

5 files changed

+363
-11
lines changed

‎llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,14 @@ bool CallLowering::lowerCall(MachineIRBuilder &MIRBuilder, const CallBase &CB,
149149
// Try looking through a bitcast from one function type to another.
150150
// Commonly happens with calls to objc_msgSend().
151151
const Value *CalleeV = CB.getCalledOperand()->stripPointerCasts();
152+
153+
// If IRTranslator chose to drop the ptrauth info, we can turn this into
154+
// a direct call.
155+
if (!PAI && CB.countOperandBundlesOfType(LLVMContext::OB_ptrauth)) {
156+
CalleeV = cast<ConstantPtrAuth>(CalleeV)->getPointer();
157+
assert(isa<Function>(CalleeV));
158+
}
159+
152160
if (const Function *F = dyn_cast<Function>(CalleeV)) {
153161
if (F->hasFnAttribute(Attribute::NonLazyBind)) {
154162
LLT Ty = getLLTForType(*F->getType(), DL);

‎llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2649,17 +2649,24 @@ bool IRTranslator::translateCallBase(const CallBase &CB,
26492649
}
26502650

26512651
std::optional<CallLowering::PtrAuthInfo> PAI;
2652-
if (CB.countOperandBundlesOfType(LLVMContext::OB_ptrauth)) {
2652+
if (auto Bundle = CB.getOperandBundle(LLVMContext::OB_ptrauth)) {
26532653
// Functions should never be ptrauth-called directly.
26542654
assert(!CB.getCalledFunction() && "invalid direct ptrauth call");
26552655

2656-
auto PAB = CB.getOperandBundle("ptrauth");
2657-
const Value *Key = PAB->Inputs[0];
2658-
const Value *Discriminator = PAB->Inputs[1];
2659-
2660-
Register DiscReg = getOrCreateVReg(*Discriminator);
2661-
PAI = CallLowering::PtrAuthInfo{cast<ConstantInt>(Key)->getZExtValue(),
2662-
DiscReg};
2656+
const Value *Key = Bundle->Inputs[0];
2657+
const Value *Discriminator = Bundle->Inputs[1];
2658+
2659+
// Look through ptrauth constants to try to eliminate the matching bundle
2660+
// and turn this into a direct call with no ptrauth.
2661+
// CallLowering will use the raw pointer if it doesn't find the PAI.
2662+
const auto *CalleeCPA = dyn_cast<ConstantPtrAuth>(CB.getCalledOperand());
2663+
if (!CalleeCPA || !isa<Function>(CalleeCPA->getPointer()) ||
2664+
!CalleeCPA->isKnownCompatibleWith(Key, Discriminator, *DL)) {
2665+
// If we can't make it direct, package the bundle into PAI.
2666+
Register DiscReg = getOrCreateVReg(*Discriminator);
2667+
PAI = CallLowering::PtrAuthInfo{cast<ConstantInt>(Key)->getZExtValue(),
2668+
DiscReg};
2669+
}
26632670
}
26642671

26652672
Register ConvergenceCtrlToken = 0;

‎llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9454,6 +9454,14 @@ void SelectionDAGBuilder::LowerCallSiteWithPtrAuthBundle(
94549454
assert(Discriminator->getType()->isIntegerTy(64) &&
94559455
"Invalid ptrauth discriminator");
94569456

9457+
// Look through ptrauth constants to find the raw callee.
9458+
// Do a direct unauthenticated call if we found it and everything matches.
9459+
if (const auto *CalleeCPA = dyn_cast<ConstantPtrAuth>(CalleeV))
9460+
if (CalleeCPA->isKnownCompatibleWith(Key, Discriminator,
9461+
DAG.getDataLayout()))
9462+
return LowerCallTo(CB, getValue(CalleeCPA->getPointer()), CB.isTailCall(),
9463+
CB.isMustTailCall(), EHPadBB);
9464+
94579465
// Functions should never be ptrauth-called directly.
94589466
assert(!isa<Function>(CalleeV) && "invalid direct ptrauth call");
94599467

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

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,4 +269,136 @@ define i32 @test_tailcall_ib_arg_ind(ptr %arg0, i64 %arg1) #0 {
269269
ret i32 %tmp1
270270
}
271271

272+
; Test direct calls
273+
274+
define i32 @test_direct_call() #0 {
275+
; DARWIN-LABEL: test_direct_call:
276+
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
277+
; DARWIN-NEXT: bl _f
278+
; DARWIN-NEXT: ldp x29, x30, [sp], #16
279+
; DARWIN-NEXT: ret
280+
;
281+
; ELF-LABEL: test_direct_call:
282+
; ELF-NEXT: str x30, [sp, #-16]!
283+
; ELF-NEXT: bl f
284+
; ELF-NEXT: ldr x30, [sp], #16
285+
; ELF-NEXT: ret
286+
%tmp0 = call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 0, i64 42) ]
287+
ret i32 %tmp0
288+
}
289+
290+
define i32 @test_direct_tailcall(ptr %arg0) #0 {
291+
; DARWIN-LABEL: test_direct_tailcall:
292+
; DARWIN: b _f
293+
;
294+
; ELF-LABEL: test_direct_tailcall:
295+
; ELF-NEXT: b f
296+
%tmp0 = tail call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 0, i64 42) ]
297+
ret i32 %tmp0
298+
}
299+
300+
define i32 @test_direct_call_mismatch() #0 {
301+
; DARWIN-LABEL: test_direct_call_mismatch:
302+
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
303+
; DARWIN-NEXT: adrp x16, _f@GOTPAGE
304+
; DARWIN-NEXT: ldr x16, [x16, _f@GOTPAGEOFF]
305+
; DARWIN-NEXT: mov x17, #42
306+
; DARWIN-NEXT: pacia x16, x17
307+
; DARWIN-NEXT: mov x8, x16
308+
; DARWIN-NEXT: mov x17, #42
309+
; DARWIN-NEXT: blrab x8, x17
310+
; DARWIN-NEXT: ldp x29, x30, [sp], #16
311+
; DARWIN-NEXT: ret
312+
;
313+
; ELF-LABEL: test_direct_call_mismatch:
314+
; ELF-NEXT: str x30, [sp, #-16]!
315+
; ELF-NEXT: adrp x16, :got:f
316+
; ELF-NEXT: ldr x16, [x16, :got_lo12:f]
317+
; ELF-NEXT: mov x17, #42
318+
; ELF-NEXT: pacia x16, x17
319+
; ELF-NEXT: mov x8, x16
320+
; ELF-NEXT: mov x17, #42
321+
; ELF-NEXT: blrab x8, x17
322+
; ELF-NEXT: ldr x30, [sp], #16
323+
; ELF-NEXT: ret
324+
%tmp0 = call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 1, i64 42) ]
325+
ret i32 %tmp0
326+
}
327+
328+
define i32 @test_direct_call_addr() #0 {
329+
; DARWIN-LABEL: test_direct_call_addr:
330+
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
331+
; DARWIN-NEXT: bl _f
332+
; DARWIN-NEXT: ldp x29, x30, [sp], #16
333+
; DARWIN-NEXT: ret
334+
;
335+
; ELF-LABEL: test_direct_call_addr:
336+
; ELF-NEXT: str x30, [sp, #-16]!
337+
; ELF-NEXT: bl f
338+
; ELF-NEXT: ldr x30, [sp], #16
339+
; ELF-NEXT: ret
340+
%tmp0 = call i32 ptrauth(ptr @f, i32 1, i64 0, ptr @f.ref.ib.0.addr)() [ "ptrauth"(i32 1, i64 ptrtoint (ptr @f.ref.ib.0.addr to i64)) ]
341+
ret i32 %tmp0
342+
}
343+
344+
define i32 @test_direct_call_addr_blend() #0 {
345+
; DARWIN-LABEL: test_direct_call_addr_blend:
346+
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
347+
; DARWIN-NEXT: bl _f
348+
; DARWIN-NEXT: ldp x29, x30, [sp], #16
349+
; DARWIN-NEXT: ret
350+
;
351+
; ELF-LABEL: test_direct_call_addr_blend:
352+
; ELF-NEXT: str x30, [sp, #-16]!
353+
; ELF-NEXT: bl f
354+
; ELF-NEXT: ldr x30, [sp], #16
355+
; ELF-NEXT: ret
356+
%tmp0 = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f.ref.ib.42.addr to i64), i64 42)
357+
%tmp1 = call i32 ptrauth(ptr @f, i32 1, i64 42, ptr @f.ref.ib.42.addr)() [ "ptrauth"(i32 1, i64 %tmp0) ]
358+
ret i32 %tmp1
359+
}
360+
361+
define i32 @test_direct_call_addr_gep_different_index_types() #0 {
362+
; DARWIN-LABEL: test_direct_call_addr_gep_different_index_types:
363+
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
364+
; DARWIN-NEXT: bl _f
365+
; DARWIN-NEXT: ldp x29, x30, [sp], #16
366+
; DARWIN-NEXT: ret
367+
;
368+
; ELF-LABEL: test_direct_call_addr_gep_different_index_types:
369+
; ELF-NEXT: str x30, [sp, #-16]!
370+
; ELF-NEXT: bl f
371+
; ELF-NEXT: ldr x30, [sp], #16
372+
; ELF-NEXT: ret
373+
%tmp0 = call i32 ptrauth(ptr @f, i32 1, i64 0, ptr getelementptr ({ ptr }, ptr @f_struct.ref.ib.0.addr, i64 0, i32 0))() [ "ptrauth"(i32 1, i64 ptrtoint (ptr getelementptr ({ ptr }, ptr @f_struct.ref.ib.0.addr, i32 0, i32 0) to i64)) ]
374+
ret i32 %tmp0
375+
}
376+
377+
define i32 @test_direct_call_addr_blend_gep_different_index_types() #0 {
378+
; DARWIN-LABEL: test_direct_call_addr_blend_gep_different_index_types:
379+
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
380+
; DARWIN-NEXT: bl _f
381+
; DARWIN-NEXT: ldp x29, x30, [sp], #16
382+
; DARWIN-NEXT: ret
383+
;
384+
; ELF-LABEL: test_direct_call_addr_blend_gep_different_index_types:
385+
; ELF-NEXT: str x30, [sp, #-16]!
386+
; ELF-NEXT: bl f
387+
; ELF-NEXT: ldr x30, [sp], #16
388+
; ELF-NEXT: ret
389+
%tmp0 = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr getelementptr ({ ptr }, ptr @f_struct.ref.ib.123.addr, i32 0, i32 0) to i64), i64 123)
390+
%tmp1 = call i32 ptrauth(ptr @f, i32 1, i64 123, ptr getelementptr ({ ptr }, ptr @f_struct.ref.ib.123.addr, i64 0, i32 0))() [ "ptrauth"(i32 1, i64 %tmp0) ]
391+
ret i32 %tmp1
392+
}
393+
394+
@f.ref.ib.42.addr = external global ptr
395+
@f.ref.ib.0.addr = external global ptr
396+
@f_struct.ref.ib.0.addr = external global ptr
397+
@f_struct.ref.ib.123.addr = external global ptr
398+
399+
declare void @f()
400+
401+
declare i64 @llvm.ptrauth.auth(i64, i32, i64)
402+
declare i64 @llvm.ptrauth.blend(i64, i64)
403+
272404
attributes #0 = { nounwind }

0 commit comments

Comments
 (0)