[llvm] [AArch64][PAC] Lower direct authenticated calls to ptrauth constants. (PR #97664)
Ahmed Bougacha via llvm-commits
llvm-commits at lists.llvm.org
Wed Jul 3 20:12:11 PDT 2024
https://github.com/ahmedbougacha created https://github.com/llvm/llvm-project/pull/97664
None
>From adffb5b0ceefb7e720b2c3b77e98d2320195b064 Mon Sep 17 00:00:00 2001
From: Ahmed Bougacha <ahmed at bougacha.org>
Date: Tue, 5 Sep 2023 17:42:12 -0700
Subject: [PATCH] [AArch64][PAC] Lower direct authenticated calls to ptrauth
constants.
---
llvm/lib/CodeGen/GlobalISel/CallLowering.cpp | 8 ++
llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp | 23 +--
.../SelectionDAG/SelectionDAGBuilder.cpp | 8 ++
llvm/test/CodeGen/AArch64/ptrauth-call.ll | 134 ++++++++++++++++++
llvm/test/CodeGen/AArch64/ptrauth-invoke.ll | 85 +++++++++++
5 files changed, 250 insertions(+), 8 deletions(-)
diff --git a/llvm/lib/CodeGen/GlobalISel/CallLowering.cpp b/llvm/lib/CodeGen/GlobalISel/CallLowering.cpp
index ee94c0bfbf9d0..d16585b5650a7 100644
--- a/llvm/lib/CodeGen/GlobalISel/CallLowering.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CallLowering.cpp
@@ -149,6 +149,14 @@ bool CallLowering::lowerCall(MachineIRBuilder &MIRBuilder, const CallBase &CB,
// Try looking through a bitcast from one function type to another.
// Commonly happens with calls to objc_msgSend().
const Value *CalleeV = CB.getCalledOperand()->stripPointerCasts();
+
+ // If IRTranslator chose to drop the ptrauth info, we can turn this into
+ // a direct call.
+ if (!PAI && CB.countOperandBundlesOfType(LLVMContext::OB_ptrauth)) {
+ CalleeV = cast<ConstantPtrAuth>(CalleeV)->getPointer();
+ assert(isa<Function>(CalleeV));
+ }
+
if (const Function *F = dyn_cast<Function>(CalleeV)) {
if (F->hasFnAttribute(Attribute::NonLazyBind)) {
LLT Ty = getLLTForType(*F->getType(), DL);
diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index 7b96f4589f5c4..7c6fa1792db80 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -2649,17 +2649,24 @@ bool IRTranslator::translateCallBase(const CallBase &CB,
}
std::optional<CallLowering::PtrAuthInfo> PAI;
- if (CB.countOperandBundlesOfType(LLVMContext::OB_ptrauth)) {
+ if (auto Bundle = CB.getOperandBundle(LLVMContext::OB_ptrauth)) {
// Functions should never be ptrauth-called directly.
assert(!CB.getCalledFunction() && "invalid direct ptrauth call");
- auto PAB = CB.getOperandBundle("ptrauth");
- const Value *Key = PAB->Inputs[0];
- const Value *Discriminator = PAB->Inputs[1];
-
- Register DiscReg = getOrCreateVReg(*Discriminator);
- PAI = CallLowering::PtrAuthInfo{cast<ConstantInt>(Key)->getZExtValue(),
- DiscReg};
+ const Value *Key = Bundle->Inputs[0];
+ const Value *Discriminator = Bundle->Inputs[1];
+
+ // Look through ptrauth constants to try to eliminate the matching bundle
+ // and turn this into a direct call with no ptrauth.
+ // CallLowering will use the raw pointer if it doesn't find the PAI.
+ auto *CalleeCPA = dyn_cast<ConstantPtrAuth>(CB.getCalledOperand());
+ if (!CalleeCPA || !isa<Function>(CalleeCPA->getPointer()) ||
+ !CalleeCPA->isKnownCompatibleWith(Key, Discriminator, *DL)) {
+ // If we can't make it direct, package the bundle into PAI.
+ Register DiscReg = getOrCreateVReg(*Discriminator);
+ PAI = CallLowering::PtrAuthInfo{cast<ConstantInt>(Key)->getZExtValue(),
+ DiscReg};
+ }
}
Register ConvergenceCtrlToken = 0;
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 8db2708d41a69..d6654dfa1d05b 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -9374,6 +9374,14 @@ void SelectionDAGBuilder::LowerCallSiteWithPtrAuthBundle(
assert(Discriminator->getType()->isIntegerTy(64) &&
"Invalid ptrauth discriminator");
+ // Look through ptrauth constants to find the raw callee.
+ // Do a direct unauthenticated call if we found it and everything matches.
+ if (auto *CalleeCPA = dyn_cast<ConstantPtrAuth>(CalleeV))
+ if (CalleeCPA->isKnownCompatibleWith(Key, Discriminator,
+ DAG.getDataLayout()))
+ return LowerCallTo(CB, getValue(CalleeCPA->getPointer()), CB.isTailCall(),
+ CB.isMustTailCall(), EHPadBB);
+
// Functions should never be ptrauth-called directly.
assert(!isa<Function>(CalleeV) && "invalid direct ptrauth call");
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-call.ll b/llvm/test/CodeGen/AArch64/ptrauth-call.ll
index 72e158fdf9916..717ad2e0026b9 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-call.ll
+++ b/llvm/test/CodeGen/AArch64/ptrauth-call.ll
@@ -269,4 +269,138 @@ define i32 @test_tailcall_ib_arg_ind(ptr %arg0, i64 %arg1) #0 {
ret i32 %tmp1
}
+; Test direct calls
+
+define i32 @test_direct_call() #0 {
+; DARWIN-LABEL: test_direct_call:
+; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT: bl _f
+; DARWIN-NEXT: ldp x29, x30, [sp], #16
+; DARWIN-NEXT: ret
+;
+; ELF-LABEL: test_direct_call:
+; ELF-NEXT: str x30, [sp, #-16]!
+; ELF-NEXT: bl f
+; ELF-NEXT: ldr x30, [sp], #16
+; ELF-NEXT: ret
+ %tmp0 = call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 0, i64 42) ]
+ ret i32 %tmp0
+}
+
+define i32 @test_direct_tailcall(ptr %arg0) #0 {
+; DARWIN-LABEL: test_direct_tailcall:
+; DARWIN: b _f
+;
+; ELF-LABEL: test_direct_tailcall:
+; ELF-NEXT: b f
+ %tmp0 = tail call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 0, i64 42) ]
+ ret i32 %tmp0
+}
+
+define i32 @test_direct_call_mismatch() #0 {
+; DARWIN-LABEL: test_direct_call_mismatch:
+; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT: adrp x16, _f at GOTPAGE
+; DARWIN-NEXT: ldr x16, [x16, _f at GOTPAGEOFF]
+; DARWIN-NEXT: mov x17, #42
+; DARWIN-NEXT: pacia x16, x17
+; DARWIN-NEXT: mov x8, x16
+; DARWIN-NEXT: mov x17, #42
+; DARWIN-NEXT: blrab x8, x17
+; DARWIN-NEXT: ldp x29, x30, [sp], #16
+; DARWIN-NEXT: ret
+;
+; ELF-LABEL: test_direct_call_mismatch:
+; ELF-NEXT: str x30, [sp, #-16]!
+; ELF-NEXT: adrp x16, :got:f
+; ELF-NEXT: ldr x16, [x16, :got_lo12:f]
+; ELF-NEXT: mov x17, #42
+; ELF-NEXT: pacia x16, x17
+; ELF-NEXT: mov x8, x16
+; ELF-NEXT: mov x17, #42
+; ELF-NEXT: blrab x8, x17
+; ELF-NEXT: ldr x30, [sp], #16
+; ELF-NEXT: ret
+ %tmp0 = call i32 ptrauth(ptr @f, i32 0, i64 42)() [ "ptrauth"(i32 1, i64 42) ]
+ ret i32 %tmp0
+}
+
+define i32 @test_direct_call_addr() #0 {
+; DARWIN-LABEL: test_direct_call_addr:
+; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT: bl _f
+; DARWIN-NEXT: ldp x29, x30, [sp], #16
+; DARWIN-NEXT: ret
+;
+; ELF-LABEL: test_direct_call_addr:
+; ELF-NEXT: str x30, [sp, #-16]!
+; ELF-NEXT: bl f
+; ELF-NEXT: ldr x30, [sp], #16
+; ELF-NEXT: ret
+ %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)) ]
+ ret i32 %tmp0
+}
+
+define i32 @test_direct_call_addr_blend() #0 {
+; DARWIN-LABEL: test_direct_call_addr_blend:
+; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT: bl _f
+; DARWIN-NEXT: ldp x29, x30, [sp], #16
+; DARWIN-NEXT: ret
+;
+; ELF-LABEL: test_direct_call_addr_blend:
+; ELF-NEXT: str x30, [sp, #-16]!
+; ELF-NEXT: bl f
+; ELF-NEXT: ldr x30, [sp], #16
+; ELF-NEXT: ret
+ %tmp0 = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f.ref.ib.42.addr to i64), i64 42)
+ %tmp1 = call i32 ptrauth(ptr @f, i32 1, i64 42, ptr @f.ref.ib.42.addr)() [ "ptrauth"(i32 1, i64 %tmp0) ]
+ ret i32 %tmp1
+}
+
+define i32 @test_direct_call_addr_gep_different_index_types() #0 {
+; DARWIN-LABEL: test_direct_call_addr_gep_different_index_types:
+; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT: bl _f
+; DARWIN-NEXT: ldp x29, x30, [sp], #16
+; DARWIN-NEXT: ret
+;
+; ELF-LABEL: test_direct_call_addr_gep_different_index_types:
+; ELF-NEXT: str x30, [sp, #-16]!
+; ELF-NEXT: bl f
+; ELF-NEXT: ldr x30, [sp], #16
+; ELF-NEXT: ret
+ %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)) ]
+ ret i32 %tmp0
+}
+
+define i32 @test_direct_call_addr_blend_gep_different_index_types() #0 {
+; DARWIN-LABEL: test_direct_call_addr_blend_gep_different_index_types:
+; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
+; DARWIN-NEXT: bl _f
+; DARWIN-NEXT: ldp x29, x30, [sp], #16
+; DARWIN-NEXT: ret
+;
+; ELF-LABEL: test_direct_call_addr_blend_gep_different_index_types:
+; ELF-NEXT: str x30, [sp, #-16]!
+; ELF-NEXT: bl f
+; ELF-NEXT: ldr x30, [sp], #16
+; ELF-NEXT: ret
+ %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)
+ %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) ]
+ ret i32 %tmp1
+}
+
+attributes #0 = { nounwind }
+
+ at f.ref.ib.42.addr = external global ptr
+ at f.ref.ib.0.addr = external global ptr
+ at f_struct.ref.ib.0.addr = external global ptr
+ at f_struct.ref.ib.123.addr = external global ptr
+
+declare void @f()
+
+declare i64 @llvm.ptrauth.auth(i64, i32, i64)
+declare i64 @llvm.ptrauth.blend(i64, i64)
+
attributes #0 = { nounwind }
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-invoke.ll b/llvm/test/CodeGen/AArch64/ptrauth-invoke.ll
index fcd0ddb788336..dead82603935c 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-invoke.ll
+++ b/llvm/test/CodeGen/AArch64/ptrauth-invoke.ll
@@ -96,6 +96,90 @@ continuebb:
ret i32 %tmp0
}
+; DARWIN-LABEL: _test_invoke_ia_0_direct:
+; DARWIN-NEXT: [[FNBEGIN:L.*]]:
+; DARWIN-NEXT: .cfi_startproc
+; DARWIN-NEXT: .cfi_personality 155, ___gxx_personality_v0
+; DARWIN-NEXT: .cfi_lsda 16, [[EXCEPT:Lexception[0-9]+]]
+; DARWIN-NEXT: ; %bb.0:
+; DARWIN-NEXT: stp x20, x19, [sp, #-32]!
+; DARWIN-NEXT: stp x29, x30, [sp, #16]
+; DARWIN-NEXT: .cfi_def_cfa_offset 32
+; DARWIN-NEXT: .cfi_offset w30, -8
+; DARWIN-NEXT: .cfi_offset w29, -16
+; DARWIN-NEXT: .cfi_offset w19, -24
+; DARWIN-NEXT: .cfi_offset w20, -32
+; DARWIN-NEXT: [[PRECALL:L.*]]:
+; DARWIN-NEXT: bl _baz
+
+; DARWIN-SDAG-NEXT: [[POSTCALL:L.*]]:
+; DARWIN-SDAG-NEXT: ; %bb.1:
+; DARWIN-SDAG-NEXT: mov x19, x0
+
+; DARWIN-GISEL-NEXT: mov x19, x0
+; DARWIN-GISEL-NEXT: [[POSTCALL:L.*]]:
+
+; DARWIN-NEXT: [[CALLBB:L.*]]:
+; DARWIN-NEXT: bl _foo
+; DARWIN-NEXT: mov x0, x19
+; DARWIN-NEXT: ldp x29, x30, [sp, #16]
+; DARWIN-NEXT: ldp x20, x19, [sp], #32
+; DARWIN-NEXT: ret
+; DARWIN-NEXT: [[LPADBB:LBB[0-9_]+]]:
+; DARWIN-NEXT: [[LPAD:L.*]]:
+; DARWIN-NEXT: mov w19, #-1
+; DARWIN-NEXT: b [[CALLBB]]
+
+; ELF-LABEL: test_invoke_ia_0_direct:
+; ELF-NEXT: [[FNBEGIN:.L.*]]:
+; ELF-NEXT: .cfi_startproc
+; ELF-NEXT: .cfi_personality 156, DW.ref.__gxx_personality_v0
+; ELF-NEXT: .cfi_lsda 28, [[EXCEPT:.Lexception[0-9]+]]
+; ELF-NEXT: // %bb.0:
+; ELF-NEXT: stp x30, x19, [sp, #-16]!
+; ELF-NEXT: .cfi_def_cfa_offset 16
+; ELF-NEXT: .cfi_offset w19, -8
+; ELF-NEXT: .cfi_offset w30, -16
+; ELF-NEXT: [[PRECALL:.L.*]]:
+; ELF-NEXT: bl baz
+
+; ELF-SDAG-NEXT: [[POSTCALL:.L.*]]:
+; ELF-SDAG-NEXT: // %bb.1:
+; ELF-SDAG-NEXT: mov w19, w0
+
+; ELF-GISEL-NEXT: mov w19, w0
+; ELF-GISEL-NEXT: [[POSTCALL:.L.*]]:
+
+; ELF-NEXT: [[CALLBB:.L.*]]:
+; ELF-NEXT: bl foo
+; ELF-NEXT: mov w0, w19
+; ELF-NEXT: ldp x30, x19, [sp], #16
+; ELF-NEXT: ret
+; ELF-NEXT: [[LPADBB:.LBB[0-9_]+]]:
+; ELF-NEXT: [[LPAD:.L.*]]:
+; ELF-NEXT: mov w19, #-1
+; ELF-NEXT: b [[CALLBB]]
+
+; CHECK-LABEL: GCC_except_table{{.*}}:
+; CHECK-NEXT: [[EXCEPT]]:
+; CHECK: .uleb128 [[POSTCALL]]-[[PRECALL]] {{.*}} Call between [[PRECALL]] and [[POSTCALL]]
+; CHECK-NEXT: .uleb128 [[LPAD]]-[[FNBEGIN]] {{.*}} jumps to [[LPAD]]
+; CHECK-NEXT: .byte 0 {{.*}} On action: cleanup
+
+define i32 @test_invoke_ia_0_direct() #0 personality ptr @__gxx_personality_v0 {
+ %tmp0 = invoke i32 ptrauth (ptr @baz, i32 0)() [ "ptrauth"(i32 0, i64 0) ] to label %continuebb
+ unwind label %unwindbb
+
+unwindbb:
+ %tmp1 = landingpad { ptr, i32 } cleanup
+ call void @foo()
+ ret i32 -1
+
+continuebb:
+ call void @foo()
+ ret i32 %tmp0
+}
+
@_ZTIPKc = external constant ptr
@hello_str = private unnamed_addr constant [6 x i8] c"hello\00", align 1
@@ -265,6 +349,7 @@ continuebb:
declare void @foo()
declare void @bar(ptr)
+declare i32 @baz()
declare i32 @__gxx_personality_v0(...)
declare ptr @__cxa_allocate_exception(i64)
More information about the llvm-commits
mailing list