[AArch64][PAC] Lower direct authenticated calls to ptrauth constants. #97664
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
llvmbot
added
backend:AArch64
llvm:globalisel
llvm:SelectionDAG
|
@llvm/pr-subscribers-llvm-globalisel @llvm/pr-subscribers-llvm-selectiondag Author: Ahmed Bougacha (ahmedbougacha) ChangesThis tries to turn indirect ptrauth calls into direct calls, using 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. Full diff: https://github.com/llvm/llvm-project/pull/97664.diff 5 Files Affected:
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@GOTPAGE
+; DARWIN-NEXT: ldr x16, [x16, _f@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 }
+
+@f.ref.ib.42.addr = external global ptr
+@f.ref.ib.0.addr = external global ptr
+@f_struct.ref.ib.0.addr = external global ptr
+@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)
|
|
@ahmedbougacha Could you please explicitly add smth like "Depends on #97665" to the PR description? Without applying that changes, tests fail since they require ptrauth constants lowering on MachO. |
kovdan01
approved these changes
Jul 4, 2024
| ; 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 { |
There was a problem hiding this comment.
In ptrauth-call.ll, we have more tests: for different constant discrimination values, with/without address discrimination, with mismatching signing parameters (!isKnownCompatibleWith), etc. I suppose that here some of such tests are not needed, but it looks like that it's worth to at least have a test for mismatching params. Above we have a test when we use invoke with ptr %arg0, but it would be nice to have a test with a mismatching ptrauth constant as well.
I'm OK with merging this "as is" and adding new tests in a separate patch later.
There was a problem hiding this comment.
Sure, added a test to verify codegen on mismatch. I didn't add all the variations since that's covered for calls; we can add more if the two diverge.
kovdan01
reviewed
Jul 4, 2024
- add mismatch test, move direct tests to the end of the file - nits
ahmedbougacha
force-pushed
the
users/ahmedbougacha/ptrauth-direct-call
branch
from
adffb5b to
85fdea7
Compare
yuxuanchen1997
pushed a commit
that referenced
this pull request
Jul 25, 2024
…#97664) Summary: 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. Test Plan: Reviewers: Subscribers: Tasks: Tags: Differential Revision: https://phabricator.intern.facebook.com/D60251467
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This tries to turn indirect ptrauth calls into direct calls, using
ConstantPtrAuth::isKnownEquivalentto compare theConstantPtrAuthtarget 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.
Depends on #97665