Skip to content

Commit 5452cbc

Browse files
ostannardostannard
authored
[AArch64] Indirect tail-calls cannot use x16 with pac-ret+pc (#81020)
When using -mbranch-protection=pac-ret+pc, x16 is used in the function epilogue to hold the address of the signing instruction. This is used by a HINT instruction which can only use x16, so we can't change this. This means that we can't use it to hold the function pointer for an indirect tail-call. There is existing code to force indirect tail-calls to use x16 or x17 when BTI is enabled, so there are now 4 combinations: bti pac-ret+pc Valid function pointer registers off off Any non callee-saved register on off x16 or x17 off on Any non callee-saved register except x16 on on x17
1 parent 3e33b6f commit 5452cbc

9 files changed

+138
-24
lines changed

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1602,7 +1602,9 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
16021602
// attributes (isCall, isReturn, etc.). We lower them to the real
16031603
// instruction here.
16041604
case AArch64::TCRETURNri:
1605-
case AArch64::TCRETURNriBTI:
1605+
case AArch64::TCRETURNrix16x17:
1606+
case AArch64::TCRETURNrix17:
1607+
case AArch64::TCRETURNrinotx16:
16061608
case AArch64::TCRETURNriALL: {
16071609
MCInst TmpInst;
16081610
TmpInst.setOpcode(AArch64::BR);

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25700,7 +25700,9 @@ AArch64TargetLowering::EmitKCFICheck(MachineBasicBlock &MBB,
2570025700
case AArch64::BLR:
2570125701
case AArch64::BLRNoIP:
2570225702
case AArch64::TCRETURNri:
25703-
case AArch64::TCRETURNriBTI:
25703+
case AArch64::TCRETURNrix16x17:
25704+
case AArch64::TCRETURNrix17:
25705+
case AArch64::TCRETURNrinotx16:
2570425706
break;
2570525707
default:
2570625708
llvm_unreachable("Unexpected CFI call opcode");

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2503,7 +2503,9 @@ bool AArch64InstrInfo::isTailCallReturnInst(const MachineInstr &MI) {
25032503
return false;
25042504
case AArch64::TCRETURNdi:
25052505
case AArch64::TCRETURNri:
2506-
case AArch64::TCRETURNriBTI:
2506+
case AArch64::TCRETURNrix16x17:
2507+
case AArch64::TCRETURNrix17:
2508+
case AArch64::TCRETURNrinotx16:
25072509
case AArch64::TCRETURNriALL:
25082510
return true;
25092511
}

‎llvm/lib/Target/AArch64/AArch64InstrInfo.td

Lines changed: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -928,8 +928,25 @@ let RecomputePerFunction = 1 in {
928928
// Avoid generating STRQro if it is slow, unless we're optimizing for code size.
929929
def UseSTRQro : Predicate<"!Subtarget->isSTRQroSlow() || shouldOptForSize(MF)">;
930930

931-
def UseBTI : Predicate<[{ MF->getInfo<AArch64FunctionInfo>()->branchTargetEnforcement() }]>;
932-
def NotUseBTI : Predicate<[{ !MF->getInfo<AArch64FunctionInfo>()->branchTargetEnforcement() }]>;
931+
// Register restrictions for indirect tail-calls:
932+
// - If branch target enforcement is enabled, indirect calls must use x16 or
933+
// x17, because these are the only registers which can target the BTI C
934+
// instruction.
935+
// - If PAuthLR is enabled, x16 is used in the epilogue to hold the address
936+
// of the signing instruction. This can't be changed because it is used by a
937+
// HINT instruction which only accepts x16. We can't load anything from the
938+
// stack after this because the authentication instruction checks that SP is
939+
// the same as it was at function entry, so we can't have anything on the
940+
// stack.
941+
942+
// BTI on, PAuthLR off: x16 or x17
943+
def TailCallX16X17 : Predicate<[{ MF->getInfo<AArch64FunctionInfo>()->branchTargetEnforcement() && !MF->getInfo<AArch64FunctionInfo>()->branchProtectionPAuthLR() }]>;
944+
// BTI on, PAuthLR on: x17 only
945+
def TailCallX17 : Predicate<[{ MF->getInfo<AArch64FunctionInfo>()->branchTargetEnforcement() && MF->getInfo<AArch64FunctionInfo>()->branchProtectionPAuthLR() }]>;
946+
// BTI off, PAuthLR on: Any non-callee-saved register except x16
947+
def TailCallNotX16 : Predicate<[{ !MF->getInfo<AArch64FunctionInfo>()->branchTargetEnforcement() && MF->getInfo<AArch64FunctionInfo>()->branchProtectionPAuthLR() }]>;
948+
// BTI off, PAuthLR off: Any non-callee-saved register
949+
def TailCallAny : Predicate<[{ !MF->getInfo<AArch64FunctionInfo>()->branchTargetEnforcement() && !MF->getInfo<AArch64FunctionInfo>()->branchProtectionPAuthLR() }]>;
933950

934951
def SLSBLRMitigation : Predicate<[{ MF->getSubtarget<AArch64Subtarget>().hardenSlsBlr() }]>;
935952
def NoSLSBLRMitigation : Predicate<[{ !MF->getSubtarget<AArch64Subtarget>().hardenSlsBlr() }]>;
@@ -9121,18 +9138,30 @@ let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Uses = [SP] in {
91219138
// some verifier checks for outlined functions.
91229139
def TCRETURNriALL : Pseudo<(outs), (ins GPR64:$dst, i32imm:$FPDiff), []>,
91239140
Sched<[WriteBrReg]>;
9124-
// Indirect tail-call limited to only use registers (x16 and x17) which are
9125-
// allowed to tail-call a "BTI c" instruction.
9126-
def TCRETURNriBTI : Pseudo<(outs), (ins rtcGPR64:$dst, i32imm:$FPDiff), []>,
9141+
9142+
// Indirect tail-calls with reduced register classes, needed for BTI and
9143+
// PAuthLR.
9144+
def TCRETURNrix16x17 : Pseudo<(outs), (ins tcGPRx16x17:$dst, i32imm:$FPDiff), []>,
9145+
Sched<[WriteBrReg]>;
9146+
def TCRETURNrix17 : Pseudo<(outs), (ins tcGPRx17:$dst, i32imm:$FPDiff), []>,
9147+
Sched<[WriteBrReg]>;
9148+
def TCRETURNrinotx16 : Pseudo<(outs), (ins tcGPRnotx16:$dst, i32imm:$FPDiff), []>,
91279149
Sched<[WriteBrReg]>;
91289150
}
91299151

91309152
def : Pat<(AArch64tcret tcGPR64:$dst, (i32 timm:$FPDiff)),
91319153
(TCRETURNri tcGPR64:$dst, imm:$FPDiff)>,
9132-
Requires<[NotUseBTI]>;
9133-
def : Pat<(AArch64tcret rtcGPR64:$dst, (i32 timm:$FPDiff)),
9134-
(TCRETURNriBTI rtcGPR64:$dst, imm:$FPDiff)>,
9135-
Requires<[UseBTI]>;
9154+
Requires<[TailCallAny]>;
9155+
def : Pat<(AArch64tcret tcGPRx16x17:$dst, (i32 timm:$FPDiff)),
9156+
(TCRETURNrix16x17 tcGPRx16x17:$dst, imm:$FPDiff)>,
9157+
Requires<[TailCallX16X17]>;
9158+
def : Pat<(AArch64tcret tcGPRx17:$dst, (i32 timm:$FPDiff)),
9159+
(TCRETURNrix17 tcGPRx17:$dst, imm:$FPDiff)>,
9160+
Requires<[TailCallX17]>;
9161+
def : Pat<(AArch64tcret tcGPRnotx16:$dst, (i32 timm:$FPDiff)),
9162+
(TCRETURNrinotx16 tcGPRnotx16:$dst, imm:$FPDiff)>,
9163+
Requires<[TailCallNotX16]>;
9164+
91369165
def : Pat<(AArch64tcret tglobaladdr:$dst, (i32 timm:$FPDiff)),
91379166
(TCRETURNdi texternalsym:$dst, imm:$FPDiff)>;
91389167
def : Pat<(AArch64tcret texternalsym:$dst, (i32 timm:$FPDiff)),

‎llvm/lib/Target/AArch64/AArch64RegisterInfo.td

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -217,11 +217,16 @@ def tcGPR64 : RegisterClass<"AArch64", [i64], 64, (sub GPR64common, X19, X20, X2
217217
X22, X23, X24, X25, X26,
218218
X27, X28, FP, LR)>;
219219

220-
// Restricted set of tail call registers, for use when branch target
221-
// enforcement is enabled. These are the only registers which can be used to
222-
// indirectly branch (not call) to the "BTI c" instruction at the start of a
223-
// BTI-protected function.
224-
def rtcGPR64 : RegisterClass<"AArch64", [i64], 64, (add X16, X17)>;
220+
// Restricted sets of tail call registers, for use when branch target
221+
// enforcement or PAuthLR are enabled.
222+
// For BTI, x16 and x17 are the only registers which can be used to indirectly
223+
// branch (not call) to the "BTI c" instruction at the start of a BTI-protected
224+
// function.
225+
// For PAuthLR, x16 must be used in the function epilogue for other purposes,
226+
// so cannot hold the function pointer.
227+
def tcGPRx17 : RegisterClass<"AArch64", [i64], 64, (add X17)>;
228+
def tcGPRx16x17 : RegisterClass<"AArch64", [i64], 64, (add X16, X17)>;
229+
def tcGPRnotx16 : RegisterClass<"AArch64", [i64], 64, (sub tcGPR64, X16)>;
225230

226231
// Register set that excludes registers that are reserved for procedure calls.
227232
// This is used for pseudo-instructions that are actually implemented using a

‎llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1012,16 +1012,23 @@ bool AArch64CallLowering::isEligibleForTailCallOptimization(
10121012

10131013
static unsigned getCallOpcode(const MachineFunction &CallerF, bool IsIndirect,
10141014
bool IsTailCall) {
1015+
const AArch64FunctionInfo *FuncInfo = CallerF.getInfo<AArch64FunctionInfo>();
1016+
10151017
if (!IsTailCall)
10161018
return IsIndirect ? getBLRCallOpcode(CallerF) : (unsigned)AArch64::BL;
10171019

10181020
if (!IsIndirect)
10191021
return AArch64::TCRETURNdi;
10201022

1021-
// When BTI is enabled, we need to use TCRETURNriBTI to make sure that we use
1022-
// x16 or x17.
1023-
if (CallerF.getInfo<AArch64FunctionInfo>()->branchTargetEnforcement())
1024-
return AArch64::TCRETURNriBTI;
1023+
// When BTI or PAuthLR are enabled, there are restrictions on using x16 and
1024+
// x17 to hold the function pointer.
1025+
if (FuncInfo->branchTargetEnforcement()) {
1026+
if (FuncInfo->branchProtectionPAuthLR())
1027+
return AArch64::TCRETURNrix17;
1028+
else
1029+
return AArch64::TCRETURNrix16x17;
1030+
} else if (FuncInfo->branchProtectionPAuthLR())
1031+
return AArch64::TCRETURNrinotx16;
10251032

10261033
return AArch64::TCRETURNri;
10271034
}

‎llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,9 @@ AArch64RegisterBankInfo::getRegBankFromRegClass(const TargetRegisterClass &RC,
273273
case AArch64::GPR64common_and_GPR64noipRegClassID:
274274
case AArch64::GPR64noip_and_tcGPR64RegClassID:
275275
case AArch64::tcGPR64RegClassID:
276-
case AArch64::rtcGPR64RegClassID:
276+
case AArch64::tcGPRx16x17RegClassID:
277+
case AArch64::tcGPRx17RegClassID:
278+
case AArch64::tcGPRnotx16RegClassID:
277279
case AArch64::WSeqPairsClassRegClassID:
278280
case AArch64::XSeqPairsClassRegClassID:
279281
case AArch64::MatrixIndexGPR32_8_11RegClassID:

‎llvm/test/CodeGen/AArch64/branch-target-enforcement-indirect-calls.ll

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,3 +26,68 @@ entry:
2626
; CHECK: br {{x16|x17}}
2727
ret void
2828
}
29+
define void @bti_enabled_force_x10(ptr %p) "branch-target-enforcement"="true" {
30+
entry:
31+
%p_x10 = tail call ptr asm "", "={x10},{x10},~{lr}"(ptr %p)
32+
tail call void %p_x10()
33+
; CHECK: br {{x16|x17}}
34+
ret void
35+
}
36+
37+
; sign-return-address places no further restrictions on the tail-call register.
38+
39+
define void @bti_enabled_pac_enabled(ptr %p) "branch-target-enforcement"="true" "sign-return-address"="all" {
40+
entry:
41+
tail call void %p()
42+
; CHECK: br {{x16|x17}}
43+
ret void
44+
}
45+
define void @bti_enabled_pac_enabled_force_x10(ptr %p) "branch-target-enforcement"="true" "sign-return-address"="all" {
46+
entry:
47+
%p_x10 = tail call ptr asm "", "={x10},{x10},~{lr}"(ptr %p)
48+
tail call void %p_x10()
49+
; CHECK: br {{x16|x17}}
50+
ret void
51+
}
52+
53+
; PAuthLR needs to use x16 to hold the address of the signing instruction. That
54+
; can't be changed because the hint instruction only uses that register, so the
55+
; only choice for the tail-call function pointer is x17.
56+
57+
define void @bti_enabled_pac_pc_enabled(ptr %p) "branch-target-enforcement"="true" "sign-return-address"="all" "branch-protection-pauth-lr"="true" {
58+
entry:
59+
tail call void %p()
60+
; CHECK: br x17
61+
ret void
62+
}
63+
define void @bti_enabled_pac_pc_enabled_force_x16(ptr %p) "branch-target-enforcement"="true" "sign-return-address"="all" "branch-protection-pauth-lr"="true" {
64+
entry:
65+
%p_x16 = tail call ptr asm "", "={x16},{x16},~{lr}"(ptr %p)
66+
tail call void %p_x16()
67+
; CHECK: br x17
68+
ret void
69+
}
70+
71+
; PAuthLR by itself prevents x16 from being used, but any other
72+
; non-callee-saved register can be used.
73+
74+
define void @pac_pc_enabled(ptr %p) "sign-return-address"="all" "branch-protection-pauth-lr"="true" {
75+
entry:
76+
tail call void %p()
77+
; CHECK: br {{(x[0-9]|x1[0-578])$}}
78+
ret void
79+
}
80+
define void @pac_pc_enabled_force_x16(ptr %p) "sign-return-address"="all" "branch-protection-pauth-lr"="true" {
81+
entry:
82+
%p_x16 = tail call ptr asm "", "={x16},{x16},~{lr}"(ptr %p)
83+
tail call void %p_x16()
84+
; CHECK: br {{(x[0-9]|x1[0-578])$}}
85+
ret void
86+
}
87+
define void @pac_pc_enabled_force_x17(ptr %p) "sign-return-address"="all" "branch-protection-pauth-lr"="true" {
88+
entry:
89+
%p_x17 = tail call ptr asm "", "={x17},{x17},~{lr}"(ptr %p)
90+
tail call void %p_x17()
91+
; CHECK: br x17
92+
ret void
93+
}

‎llvm/test/CodeGen/AArch64/kcfi-bti.ll

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,11 @@ define void @f3(ptr noundef %x) {
7373
; MIR-LABEL: name: f3
7474
; MIR: body:
7575

76-
; ISEL: TCRETURNriBTI %1, 0, csr_aarch64_aapcs, implicit $sp, cfi-type 12345678
76+
; ISEL: TCRETURNrix16x17 %1, 0, csr_aarch64_aapcs, implicit $sp, cfi-type 12345678
7777

7878
; KCFI: BUNDLE{{.*}} {
7979
; KCFI-NEXT: KCFI_CHECK $x16, 12345678, implicit-def $x9, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
80-
; KCFI-NEXT: TCRETURNriBTI internal killed $x16, 0, csr_aarch64_aapcs, implicit $sp
80+
; KCFI-NEXT: TCRETURNrix16x17 internal killed $x16, 0, csr_aarch64_aapcs, implicit $sp
8181
; KCFI-NEXT: }
8282

8383
tail call void %x() [ "kcfi"(i32 12345678) ]

0 commit comments

Comments
 (0)