Skip to content

Commit 77bcab8

Browse files
ahmedbougachaahmedbougacha
authored
[InstCombine] Combine ptrauth intrin. callee into same-key bundle. (#94707)
Try to optimize a call to the result of a ptrauth intrinsic, potentially into the ptrauth call bundle: call(ptrauth.resign(p)), ["ptrauth"()] -> call p, ["ptrauth"()] call(ptrauth.sign(p)), ["ptrauth"()] -> call p as long as the key/discriminator are the same in sign and auth-bundle, and we don't change the key in the bundle (to a potentially-invalid key.) Generating a plain call to a raw unauthenticated pointer is generally undesirable, but if we ended up seeing a naked ptrauth.sign in the first place, we already have suspicious code. Unauthenticated calls are also easier to spot than naked signs, so let the indirect call shine. Note that there is an arguably unsafe extension to this, where we don't bother checking that the key in bundle and intrinsic are the same (and also allow folding away an auth into a bundle.) This can end up generating calls with a bundle that has an invalid key (which an informed frontend wouldn't have otherwise done), which can be problematic. The C that generates that is straightforward but arguably unreasonable. That wouldn't be an issue if we were to bite the bullet and make these fully AArch64-specific, allowing key knowledge to be embedded here.
1 parent e685e4a commit 77bcab8

File tree

3 files changed

+231
-0
lines changed

3 files changed

+231
-0
lines changed

‎llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4050,6 +4050,83 @@ static IntrinsicInst *findInitTrampoline(Value *Callee) {
40504050
return nullptr;
40514051
}
40524052

4053+
Instruction *InstCombinerImpl::foldPtrAuthIntrinsicCallee(CallBase &Call) {
4054+
const Value *Callee = Call.getCalledOperand();
4055+
const auto *IPC = dyn_cast<IntToPtrInst>(Callee);
4056+
if (!IPC || !IPC->isNoopCast(DL))
4057+
return nullptr;
4058+
4059+
const auto *II = dyn_cast<IntrinsicInst>(IPC->getOperand(0));
4060+
if (!II)
4061+
return nullptr;
4062+
4063+
Intrinsic::ID IIID = II->getIntrinsicID();
4064+
if (IIID != Intrinsic::ptrauth_resign && IIID != Intrinsic::ptrauth_sign)
4065+
return nullptr;
4066+
4067+
// Isolate the ptrauth bundle from the others.
4068+
std::optional<OperandBundleUse> PtrAuthBundleOrNone;
4069+
SmallVector<OperandBundleDef, 2> NewBundles;
4070+
for (unsigned BI = 0, BE = Call.getNumOperandBundles(); BI != BE; ++BI) {
4071+
OperandBundleUse Bundle = Call.getOperandBundleAt(BI);
4072+
if (Bundle.getTagID() == LLVMContext::OB_ptrauth)
4073+
PtrAuthBundleOrNone = Bundle;
4074+
else
4075+
NewBundles.emplace_back(Bundle);
4076+
}
4077+
4078+
if (!PtrAuthBundleOrNone)
4079+
return nullptr;
4080+
4081+
Value *NewCallee = nullptr;
4082+
switch (IIID) {
4083+
// call(ptrauth.resign(p)), ["ptrauth"()] -> call p, ["ptrauth"()]
4084+
// assuming the call bundle and the sign operands match.
4085+
case Intrinsic::ptrauth_resign: {
4086+
// Resign result key should match bundle.
4087+
if (II->getOperand(3) != PtrAuthBundleOrNone->Inputs[0])
4088+
return nullptr;
4089+
// Resign result discriminator should match bundle.
4090+
if (II->getOperand(4) != PtrAuthBundleOrNone->Inputs[1])
4091+
return nullptr;
4092+
4093+
// Resign input (auth) key should also match: we can't change the key on
4094+
// the new call we're generating, because we don't know what keys are valid.
4095+
if (II->getOperand(1) != PtrAuthBundleOrNone->Inputs[0])
4096+
return nullptr;
4097+
4098+
Value *NewBundleOps[] = {II->getOperand(1), II->getOperand(2)};
4099+
NewBundles.emplace_back("ptrauth", NewBundleOps);
4100+
NewCallee = II->getOperand(0);
4101+
break;
4102+
}
4103+
4104+
// call(ptrauth.sign(p)), ["ptrauth"()] -> call p
4105+
// assuming the call bundle and the sign operands match.
4106+
// Non-ptrauth indirect calls are undesirable, but so is ptrauth.sign.
4107+
case Intrinsic::ptrauth_sign: {
4108+
// Sign key should match bundle.
4109+
if (II->getOperand(1) != PtrAuthBundleOrNone->Inputs[0])
4110+
return nullptr;
4111+
// Sign discriminator should match bundle.
4112+
if (II->getOperand(2) != PtrAuthBundleOrNone->Inputs[1])
4113+
return nullptr;
4114+
NewCallee = II->getOperand(0);
4115+
break;
4116+
}
4117+
default:
4118+
llvm_unreachable("unexpected intrinsic ID");
4119+
}
4120+
4121+
if (!NewCallee)
4122+
return nullptr;
4123+
4124+
NewCallee = Builder.CreateBitOrPointerCast(NewCallee, Callee->getType());
4125+
CallBase *NewCall = CallBase::Create(&Call, NewBundles);
4126+
NewCall->setCalledOperand(NewCallee);
4127+
return NewCall;
4128+
}
4129+
40534130
Instruction *InstCombinerImpl::foldPtrAuthConstantCallee(CallBase &Call) {
40544131
auto *CPA = dyn_cast<ConstantPtrAuth>(Call.getCalledOperand());
40554132
if (!CPA)
@@ -4238,6 +4315,10 @@ Instruction *InstCombinerImpl::visitCallBase(CallBase &Call) {
42384315
if (IntrinsicInst *II = findInitTrampoline(Callee))
42394316
return transformCallThroughTrampoline(Call, *II);
42404317

4318+
// Combine calls involving pointer authentication intrinsics.
4319+
if (Instruction *NewCall = foldPtrAuthIntrinsicCallee(Call))
4320+
return NewCall;
4321+
42414322
// Combine calls to ptrauth constants.
42424323
if (Instruction *NewCall = foldPtrAuthConstantCallee(Call))
42434324
return NewCall;

‎llvm/lib/Transforms/InstCombine/InstCombineInternal.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,14 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
283283
Instruction *transformCallThroughTrampoline(CallBase &Call,
284284
IntrinsicInst &Tramp);
285285

286+
/// Try to optimize a call to the result of a ptrauth intrinsic, potentially
287+
/// into the ptrauth call bundle:
288+
/// - call(ptrauth.resign(p)), ["ptrauth"()] -> call p, ["ptrauth"()]
289+
/// - call(ptrauth.sign(p)), ["ptrauth"()] -> call p
290+
/// as long as the key/discriminator are the same in sign and auth-bundle,
291+
/// and we don't change the key in the bundle (to a potentially-invalid key.)
292+
Instruction *foldPtrAuthIntrinsicCallee(CallBase &Call);
293+
286294
/// Try to optimize a call to a ptrauth constant, into its ptrauth bundle:
287295
/// call(ptrauth(f)), ["ptrauth"()] -> call f
288296
/// as long as the key/discriminator are the same in constant and bundle.

‎llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
2+
; RUN: opt < %s -passes=instcombine -S | FileCheck %s
3+
4+
define i32 @test_ptrauth_call_sign(ptr %p) {
5+
; CHECK-LABEL: @test_ptrauth_call_sign(
6+
; CHECK-NEXT: [[V3:%.*]] = call i32 [[P:%.*]]()
7+
; CHECK-NEXT: ret i32 [[V3]]
8+
;
9+
%v0 = ptrtoint ptr %p to i64
10+
%v1 = call i64 @llvm.ptrauth.sign(i64 %v0, i32 2, i64 5678)
11+
%v2 = inttoptr i64 %v1 to ptr
12+
%v3 = call i32 %v2() [ "ptrauth"(i32 2, i64 5678) ]
13+
ret i32 %v3
14+
}
15+
16+
define i32 @test_ptrauth_call_sign_otherbundle(ptr %p) {
17+
; CHECK-LABEL: @test_ptrauth_call_sign_otherbundle(
18+
; CHECK-NEXT: [[V3:%.*]] = call i32 [[P:%.*]]() [ "somebundle"(ptr null), "otherbundle"(i64 0) ]
19+
; CHECK-NEXT: ret i32 [[V3]]
20+
;
21+
%v0 = ptrtoint ptr %p to i64
22+
%v1 = call i64 @llvm.ptrauth.sign(i64 %v0, i32 2, i64 5678)
23+
%v2 = inttoptr i64 %v1 to ptr
24+
%v3 = call i32 %v2() [ "somebundle"(ptr null), "ptrauth"(i32 2, i64 5678), "otherbundle"(i64 0) ]
25+
ret i32 %v3
26+
}
27+
28+
define i32 @test_ptrauth_call_resign(ptr %p) {
29+
; CHECK-LABEL: @test_ptrauth_call_resign(
30+
; CHECK-NEXT: [[V3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 1, i64 1234) ]
31+
; CHECK-NEXT: ret i32 [[V3]]
32+
;
33+
%v0 = ptrtoint ptr %p to i64
34+
%v1 = call i64 @llvm.ptrauth.resign(i64 %v0, i32 1, i64 1234, i32 1, i64 5678)
35+
%v2 = inttoptr i64 %v1 to ptr
36+
%v3 = call i32 %v2() [ "ptrauth"(i32 1, i64 5678) ]
37+
ret i32 %v3
38+
}
39+
40+
define i32 @test_ptrauth_call_resign_blend(ptr %pp) {
41+
; CHECK-LABEL: @test_ptrauth_call_resign_blend(
42+
; CHECK-NEXT: [[V01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
43+
; CHECK-NEXT: [[V6:%.*]] = call i32 [[V01]]() [ "ptrauth"(i32 1, i64 1234) ]
44+
; CHECK-NEXT: ret i32 [[V6]]
45+
;
46+
%v0 = load ptr, ptr %pp, align 8
47+
%v1 = ptrtoint ptr %pp to i64
48+
%v2 = ptrtoint ptr %v0 to i64
49+
%v3 = call i64 @llvm.ptrauth.blend(i64 %v1, i64 5678)
50+
%v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 1, i64 1234, i32 1, i64 %v3)
51+
%v5 = inttoptr i64 %v4 to ptr
52+
%v6 = call i32 %v5() [ "ptrauth"(i32 1, i64 %v3) ]
53+
ret i32 %v6
54+
}
55+
56+
define i32 @test_ptrauth_call_resign_blend_2(ptr %pp) {
57+
; CHECK-LABEL: @test_ptrauth_call_resign_blend_2(
58+
; CHECK-NEXT: [[V01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
59+
; CHECK-NEXT: [[V1:%.*]] = ptrtoint ptr [[PP]] to i64
60+
; CHECK-NEXT: [[V3:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[V1]], i64 5678)
61+
; CHECK-NEXT: [[V6:%.*]] = call i32 [[V01]]() [ "ptrauth"(i32 0, i64 [[V3]]) ]
62+
; CHECK-NEXT: ret i32 [[V6]]
63+
;
64+
%v0 = load ptr, ptr %pp, align 8
65+
%v1 = ptrtoint ptr %pp to i64
66+
%v2 = ptrtoint ptr %v0 to i64
67+
%v3 = call i64 @llvm.ptrauth.blend(i64 %v1, i64 5678)
68+
%v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 0, i64 %v3, i32 0, i64 1234)
69+
%v5 = inttoptr i64 %v4 to ptr
70+
%v6 = call i32 %v5() [ "ptrauth"(i32 0, i64 1234) ]
71+
ret i32 %v6
72+
}
73+
74+
define i32 @test_ptrauth_call_resign_mismatch_key(ptr %p) {
75+
; CHECK-LABEL: @test_ptrauth_call_resign_mismatch_key(
76+
; CHECK-NEXT: [[V0:%.*]] = ptrtoint ptr [[P:%.*]] to i64
77+
; CHECK-NEXT: [[V1:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[V0]], i32 1, i64 1234, i32 0, i64 5678)
78+
; CHECK-NEXT: [[V2:%.*]] = inttoptr i64 [[V1]] to ptr
79+
; CHECK-NEXT: [[V3:%.*]] = call i32 [[V2]]() [ "ptrauth"(i32 1, i64 5678) ]
80+
; CHECK-NEXT: ret i32 [[V3]]
81+
;
82+
%v0 = ptrtoint ptr %p to i64
83+
%v1 = call i64 @llvm.ptrauth.resign(i64 %v0, i32 1, i64 1234, i32 0, i64 5678)
84+
%v2 = inttoptr i64 %v1 to ptr
85+
%v3 = call i32 %v2() [ "ptrauth"(i32 1, i64 5678) ]
86+
ret i32 %v3
87+
}
88+
89+
define i32 @test_ptrauth_call_resign_mismatch_disc(ptr %p) {
90+
; CHECK-LABEL: @test_ptrauth_call_resign_mismatch_disc(
91+
; CHECK-NEXT: [[V0:%.*]] = ptrtoint ptr [[P:%.*]] to i64
92+
; CHECK-NEXT: [[V1:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[V0]], i32 1, i64 1234, i32 0, i64 9900)
93+
; CHECK-NEXT: [[V2:%.*]] = inttoptr i64 [[V1]] to ptr
94+
; CHECK-NEXT: [[V3:%.*]] = call i32 [[V2]]() [ "ptrauth"(i32 1, i64 5678) ]
95+
; CHECK-NEXT: ret i32 [[V3]]
96+
;
97+
%v0 = ptrtoint ptr %p to i64
98+
%v1 = call i64 @llvm.ptrauth.resign(i64 %v0, i32 1, i64 1234, i32 0, i64 9900)
99+
%v2 = inttoptr i64 %v1 to ptr
100+
%v3 = call i32 %v2() [ "ptrauth"(i32 1, i64 5678) ]
101+
ret i32 %v3
102+
}
103+
104+
define i32 @test_ptrauth_call_resign_mismatch_blend(ptr %pp) {
105+
; CHECK-LABEL: @test_ptrauth_call_resign_mismatch_blend(
106+
; CHECK-NEXT: [[V0:%.*]] = load ptr, ptr [[PP:%.*]], align 8
107+
; CHECK-NEXT: [[V1:%.*]] = ptrtoint ptr [[PP]] to i64
108+
; CHECK-NEXT: [[V2:%.*]] = ptrtoint ptr [[V0]] to i64
109+
; CHECK-NEXT: [[V6:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[V1]], i64 5678)
110+
; CHECK-NEXT: [[V4:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[V2]], i32 1, i64 1234, i32 1, i64 [[V6]])
111+
; CHECK-NEXT: [[V5:%.*]] = inttoptr i64 [[V4]] to ptr
112+
; CHECK-NEXT: [[V3:%.*]] = call i32 [[V5]]() [ "ptrauth"(i32 1, i64 [[V1]]) ]
113+
; CHECK-NEXT: ret i32 [[V3]]
114+
;
115+
%v0 = load ptr, ptr %pp, align 8
116+
%v1 = ptrtoint ptr %pp to i64
117+
%v2 = ptrtoint ptr %v0 to i64
118+
%v3 = call i64 @llvm.ptrauth.blend(i64 %v1, i64 5678)
119+
%v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 1, i64 1234, i32 1, i64 %v3)
120+
%v5 = inttoptr i64 %v4 to ptr
121+
%v6 = call i32 %v5() [ "ptrauth"(i32 1, i64 %v1) ]
122+
ret i32 %v6
123+
}
124+
125+
define i32 @test_ptrauth_call_resign_changing_call_key(ptr %p) {
126+
; CHECK-LABEL: @test_ptrauth_call_resign_changing_call_key(
127+
; CHECK-NEXT: [[V0:%.*]] = ptrtoint ptr [[P:%.*]] to i64
128+
; CHECK-NEXT: [[V1:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[V0]], i32 2, i64 1234, i32 1, i64 5678)
129+
; CHECK-NEXT: [[V2:%.*]] = inttoptr i64 [[V1]] to ptr
130+
; CHECK-NEXT: [[V3:%.*]] = call i32 [[V2]]() [ "ptrauth"(i32 1, i64 5678) ]
131+
; CHECK-NEXT: ret i32 [[V3]]
132+
;
133+
%v0 = ptrtoint ptr %p to i64
134+
%v1 = call i64 @llvm.ptrauth.resign(i64 %v0, i32 2, i64 1234, i32 1, i64 5678)
135+
%v2 = inttoptr i64 %v1 to ptr
136+
%v3 = call i32 %v2() [ "ptrauth"(i32 1, i64 5678) ]
137+
ret i32 %v3
138+
}
139+
140+
declare i64 @llvm.ptrauth.sign(i64, i32, i64)
141+
declare i64 @llvm.ptrauth.resign(i64, i32, i64, i32, i64)
142+
declare i64 @llvm.ptrauth.blend(i64, i64)

0 commit comments

Comments
 (0)