Skip to content

Commit 6049cd6

Browse files
ahmedbougachaahmedbougacha
authored
[AArch64][PAC] Lower jump-tables using hardened pseudo. (llvm#97666)
This introduces an alternative hardened lowering for jump-table dispatch, controlled by the function attribute `"aarch64-jump-table-hardening"`. The implementation is centered around a pseudo, BR_JumpTable: > A hardened but more expensive version of jump-table dispatch. > This combines the target address computation (otherwise done using > the JumpTableDest pseudos above) with the branch itself (otherwise > done using a plain BR) in a single non-attackable sequence. > > We take the final entry index as an operand to allow isel freedom. > This does mean that the index can be attacker-controlled. To > address that, we also do limited checking of the offset, mainly > ensuring it still points within the jump-table array. When it > doesn't, this branches to the first entry. We might want it to > trap instead. > > This is intended for use in conjunction with ptrauth for other > code pointers, to avoid signing jump-table entries and turning > them into pointers. > > Entry index is passed in x16. Clobbers x16/x17/nzcv. Jump-table compression isn't supported yet.
1 parent 03e92bf commit 6049cd6

File tree

5 files changed

+347
-1
lines changed

5 files changed

+347
-1
lines changed

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

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,8 @@ class AArch64AsmPrinter : public AsmPrinter {
104104

105105
void LowerJumpTableDest(MCStreamer &OutStreamer, const MachineInstr &MI);
106106

107+
void LowerHardenedBRJumpTable(const MachineInstr &MI);
108+
107109
void LowerMOPS(MCStreamer &OutStreamer, const MachineInstr &MI);
108110

109111
void LowerSTACKMAP(MCStreamer &OutStreamer, StackMaps &SM,
@@ -1354,6 +1356,139 @@ void AArch64AsmPrinter::LowerJumpTableDest(llvm::MCStreamer &OutStreamer,
13541356
.addImm(Size == 4 ? 0 : 2));
13551357
}
13561358

1359+
void AArch64AsmPrinter::LowerHardenedBRJumpTable(const MachineInstr &MI) {
1360+
unsigned InstsEmitted = 0;
1361+
1362+
const MachineJumpTableInfo *MJTI = MF->getJumpTableInfo();
1363+
assert(MJTI && "Can't lower jump-table dispatch without JTI");
1364+
1365+
const std::vector<MachineJumpTableEntry> &JTs = MJTI->getJumpTables();
1366+
assert(!JTs.empty() && "Invalid JT index for jump-table dispatch");
1367+
1368+
// Emit:
1369+
// mov x17, #<size of table> ; depending on table size, with MOVKs
1370+
// cmp x16, x17 ; or #imm if table size fits in 12-bit
1371+
// csel x16, x16, xzr, ls ; check for index overflow
1372+
//
1373+
// adrp x17, Ltable@PAGE ; materialize table address
1374+
// add x17, Ltable@PAGEOFF
1375+
// ldrsw x16, [x17, x16, lsl #2] ; load table entry
1376+
//
1377+
// Lanchor:
1378+
// adr x17, Lanchor ; compute target address
1379+
// add x16, x17, x16
1380+
// br x16 ; branch to target
1381+
1382+
MachineOperand JTOp = MI.getOperand(0);
1383+
1384+
unsigned JTI = JTOp.getIndex();
1385+
assert(!AArch64FI->getJumpTableEntryPCRelSymbol(JTI) &&
1386+
"unsupported compressed jump table");
1387+
1388+
const uint64_t NumTableEntries = JTs[JTI].MBBs.size();
1389+
1390+
// cmp only supports a 12-bit immediate. If we need more, materialize the
1391+
// immediate, using x17 as a scratch register.
1392+
uint64_t MaxTableEntry = NumTableEntries - 1;
1393+
if (isUInt<12>(MaxTableEntry)) {
1394+
EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::SUBSXri)
1395+
.addReg(AArch64::XZR)
1396+
.addReg(AArch64::X16)
1397+
.addImm(MaxTableEntry)
1398+
.addImm(0));
1399+
++InstsEmitted;
1400+
} else {
1401+
EmitToStreamer(*OutStreamer,
1402+
MCInstBuilder(AArch64::MOVZXi)
1403+
.addReg(AArch64::X17)
1404+
.addImm(static_cast<uint16_t>(MaxTableEntry))
1405+
.addImm(0));
1406+
++InstsEmitted;
1407+
// It's sad that we have to manually materialize instructions, but we can't
1408+
// trivially reuse the main pseudo expansion logic.
1409+
// A MOVK sequence is easy enough to generate and handles the general case.
1410+
for (int Offset = 16; Offset < 64; Offset += 16) {
1411+
if ((MaxTableEntry >> Offset) == 0)
1412+
break;
1413+
EmitToStreamer(*OutStreamer,
1414+
MCInstBuilder(AArch64::MOVKXi)
1415+
.addReg(AArch64::X17)
1416+
.addReg(AArch64::X17)
1417+
.addImm(static_cast<uint16_t>(MaxTableEntry >> Offset))
1418+
.addImm(Offset));
1419+
++InstsEmitted;
1420+
}
1421+
EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::SUBSXrs)
1422+
.addReg(AArch64::XZR)
1423+
.addReg(AArch64::X16)
1424+
.addReg(AArch64::X17)
1425+
.addImm(0));
1426+
++InstsEmitted;
1427+
}
1428+
1429+
// This picks entry #0 on failure.
1430+
// We might want to trap instead.
1431+
EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::CSELXr)
1432+
.addReg(AArch64::X16)
1433+
.addReg(AArch64::X16)
1434+
.addReg(AArch64::XZR)
1435+
.addImm(AArch64CC::LS));
1436+
++InstsEmitted;
1437+
1438+
// Prepare the @PAGE/@PAGEOFF low/high operands.
1439+
MachineOperand JTMOHi(JTOp), JTMOLo(JTOp);
1440+
MCOperand JTMCHi, JTMCLo;
1441+
1442+
JTMOHi.setTargetFlags(AArch64II::MO_PAGE);
1443+
JTMOLo.setTargetFlags(AArch64II::MO_PAGEOFF | AArch64II::MO_NC);
1444+
1445+
MCInstLowering.lowerOperand(JTMOHi, JTMCHi);
1446+
MCInstLowering.lowerOperand(JTMOLo, JTMCLo);
1447+
1448+
EmitToStreamer(
1449+
*OutStreamer,
1450+
MCInstBuilder(AArch64::ADRP).addReg(AArch64::X17).addOperand(JTMCHi));
1451+
++InstsEmitted;
1452+
1453+
EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::ADDXri)
1454+
.addReg(AArch64::X17)
1455+
.addReg(AArch64::X17)
1456+
.addOperand(JTMCLo)
1457+
.addImm(0));
1458+
++InstsEmitted;
1459+
1460+
EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::LDRSWroX)
1461+
.addReg(AArch64::X16)
1462+
.addReg(AArch64::X17)
1463+
.addReg(AArch64::X16)
1464+
.addImm(0)
1465+
.addImm(1));
1466+
++InstsEmitted;
1467+
1468+
MCSymbol *AdrLabel = MF->getContext().createTempSymbol();
1469+
const auto *AdrLabelE = MCSymbolRefExpr::create(AdrLabel, MF->getContext());
1470+
AArch64FI->setJumpTableEntryInfo(JTI, 4, AdrLabel);
1471+
1472+
OutStreamer->emitLabel(AdrLabel);
1473+
EmitToStreamer(
1474+
*OutStreamer,
1475+
MCInstBuilder(AArch64::ADR).addReg(AArch64::X17).addExpr(AdrLabelE));
1476+
++InstsEmitted;
1477+
1478+
EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::ADDXrs)
1479+
.addReg(AArch64::X16)
1480+
.addReg(AArch64::X17)
1481+
.addReg(AArch64::X16)
1482+
.addImm(0));
1483+
++InstsEmitted;
1484+
1485+
EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::BR).addReg(AArch64::X16));
1486+
++InstsEmitted;
1487+
1488+
(void)InstsEmitted;
1489+
assert(STI->getInstrInfo()->getInstSizeInBytes(MI) >= InstsEmitted * 4);
1490+
}
1491+
13571492
void AArch64AsmPrinter::LowerMOPS(llvm::MCStreamer &OutStreamer,
13581493
const llvm::MachineInstr &MI) {
13591494
unsigned Opcode = MI.getOpcode();
@@ -2232,6 +2367,10 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
22322367
LowerJumpTableDest(*OutStreamer, *MI);
22332368
return;
22342369

2370+
case AArch64::BR_JumpTable:
2371+
LowerHardenedBRJumpTable(*MI);
2372+
return;
2373+
22352374
case AArch64::FMOVH0:
22362375
case AArch64::FMOVS0:
22372376
case AArch64::FMOVD0:

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10742,6 +10742,30 @@ SDValue AArch64TargetLowering::LowerBR_JT(SDValue Op,
1074210742
auto *AFI = DAG.getMachineFunction().getInfo<AArch64FunctionInfo>();
1074310743
AFI->setJumpTableEntryInfo(JTI, 4, nullptr);
1074410744

10745+
// With aarch64-jump-table-hardening, we only expand the jump table dispatch
10746+
// sequence later, to guarantee the integrity of the intermediate values.
10747+
if (DAG.getMachineFunction().getFunction().hasFnAttribute(
10748+
"aarch64-jump-table-hardening")) {
10749+
CodeModel::Model CM = getTargetMachine().getCodeModel();
10750+
if (Subtarget->isTargetMachO()) {
10751+
if (CM != CodeModel::Small && CM != CodeModel::Large)
10752+
report_fatal_error("Unsupported code-model for hardened jump-table");
10753+
} else {
10754+
// Note that COFF support would likely also need JUMP_TABLE_DEBUG_INFO.
10755+
assert(Subtarget->isTargetELF() &&
10756+
"jump table hardening only supported on MachO/ELF");
10757+
if (CM != CodeModel::Small)
10758+
report_fatal_error("Unsupported code-model for hardened jump-table");
10759+
}
10760+
10761+
SDValue X16Copy = DAG.getCopyToReg(DAG.getEntryNode(), DL, AArch64::X16,
10762+
Entry, SDValue());
10763+
SDNode *B = DAG.getMachineNode(AArch64::BR_JumpTable, DL, MVT::Other,
10764+
DAG.getTargetJumpTable(JTI, MVT::i32),
10765+
X16Copy.getValue(0), X16Copy.getValue(1));
10766+
return SDValue(B, 0);
10767+
}
10768+
1074510769
SDNode *Dest =
1074610770
DAG.getMachineNode(AArch64::JumpTableDest32, DL, MVT::i64, MVT::i64, JT,
1074710771
Entry, DAG.getTargetJumpTable(JTI, MVT::i32));

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1143,6 +1143,33 @@ def JumpTableDest8 : Pseudo<(outs GPR64:$dst, GPR64sp:$scratch),
11431143
Sched<[]>;
11441144
}
11451145

1146+
// A hardened but more expensive version of jump-table dispatch.
1147+
// This combines the target address computation (otherwise done using the
1148+
// JumpTableDest pseudos above) with the branch itself (otherwise done using
1149+
// a plain BR) in a single non-attackable sequence.
1150+
//
1151+
// We take the final entry index as an operand to allow isel freedom. This does
1152+
// mean that the index can be attacker-controlled. To address that, we also do
1153+
// limited checking of the offset, mainly ensuring it still points within the
1154+
// jump-table array. When it doesn't, this branches to the first entry.
1155+
// We might want to trap instead.
1156+
//
1157+
// This is intended for use in conjunction with ptrauth for other code pointers,
1158+
// to avoid signing jump-table entries and turning them into pointers.
1159+
//
1160+
// Entry index is passed in x16. Clobbers x16/x17/nzcv.
1161+
let isNotDuplicable = 1 in
1162+
def BR_JumpTable : Pseudo<(outs), (ins i32imm:$jti), []>, Sched<[]> {
1163+
let isBranch = 1;
1164+
let isTerminator = 1;
1165+
let isIndirectBranch = 1;
1166+
let isBarrier = 1;
1167+
let isNotDuplicable = 1;
1168+
let Defs = [X16,X17,NZCV];
1169+
let Uses = [X16];
1170+
let Size = 44; // 28 fixed + 16 variable, for table size materialization
1171+
}
1172+
11461173
// Space-consuming pseudo to aid testing of placement and reachability
11471174
// algorithms. Immediate operand is the number of bytes this "instruction"
11481175
// occupies; register operands can be used to enforce dependency and constrain

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

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3602,10 +3602,33 @@ bool AArch64InstructionSelector::selectBrJT(MachineInstr &I,
36023602
unsigned JTI = I.getOperand(1).getIndex();
36033603
Register Index = I.getOperand(2).getReg();
36043604

3605+
MF->getInfo<AArch64FunctionInfo>()->setJumpTableEntryInfo(JTI, 4, nullptr);
3606+
3607+
// With aarch64-jump-table-hardening, we only expand the jump table dispatch
3608+
// sequence later, to guarantee the integrity of the intermediate values.
3609+
if (MF->getFunction().hasFnAttribute("aarch64-jump-table-hardening")) {
3610+
CodeModel::Model CM = TM.getCodeModel();
3611+
if (STI.isTargetMachO()) {
3612+
if (CM != CodeModel::Small && CM != CodeModel::Large)
3613+
report_fatal_error("Unsupported code-model for hardened jump-table");
3614+
} else {
3615+
// Note that COFF support would likely also need JUMP_TABLE_DEBUG_INFO.
3616+
assert(STI.isTargetELF() &&
3617+
"jump table hardening only supported on MachO/ELF");
3618+
if (CM != CodeModel::Small)
3619+
report_fatal_error("Unsupported code-model for hardened jump-table");
3620+
}
3621+
3622+
MIB.buildCopy({AArch64::X16}, I.getOperand(2).getReg());
3623+
MIB.buildInstr(AArch64::BR_JumpTable)
3624+
.addJumpTableIndex(I.getOperand(1).getIndex());
3625+
I.eraseFromParent();
3626+
return true;
3627+
}
3628+
36053629
Register TargetReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
36063630
Register ScratchReg = MRI.createVirtualRegister(&AArch64::GPR64spRegClass);
36073631

3608-
MF->getInfo<AArch64FunctionInfo>()->setJumpTableEntryInfo(JTI, 4, nullptr);
36093632
auto JumpTableInst = MIB.buildInstr(AArch64::JumpTableDest32,
36103633
{TargetReg, ScratchReg}, {JTAddr, Index})
36113634
.addJumpTableIndex(JTI);

‎llvm/test/CodeGen/AArch64/hardened-br-jump-table.ll‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
; RUN: rm -rf %t && split-file %s %t
2+
3+
;--- err1.ll
4+
5+
; RUN: not --crash llc %t/err1.ll -mtriple=aarch64-elf \
6+
; RUN: -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
7+
; RUN: -code-model=large \
8+
; RUN: -o - -verify-machineinstrs 2>&1 | FileCheck %s --check-prefix=ERR1
9+
10+
; RUN: not --crash llc %t/err1.ll -mtriple=aarch64-elf \
11+
; RUN: -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
12+
; RUN: -global-isel -global-isel-abort=1 \
13+
; RUN: -code-model=large \
14+
; RUN: -o - -verify-machineinstrs 2>&1 | FileCheck %s --check-prefix=ERR1
15+
16+
; ERR1: LLVM ERROR: Unsupported code-model for hardened jump-table
17+
define i32 @test_jumptable(i32 %in) "aarch64-jump-table-hardening" {
18+
19+
switch i32 %in, label %def [
20+
i32 0, label %lbl1
21+
i32 1, label %lbl2
22+
]
23+
24+
def:
25+
ret i32 0
26+
27+
lbl1:
28+
ret i32 1
29+
30+
lbl2:
31+
ret i32 2
32+
}
33+
34+
;--- test.ll
35+
36+
; RUN: llc %t/test.ll -mtriple=arm64-apple-darwin -aarch64-enable-collect-loh=0 \
37+
; RUN: -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
38+
; RUN: -o - -verify-machineinstrs | FileCheck %s --check-prefix=MACHO
39+
40+
; RUN: llc %t/test.ll -mtriple=arm64-apple-darwin -aarch64-enable-collect-loh=0 \
41+
; RUN: -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
42+
; RUN: -global-isel -global-isel-abort=1 \
43+
; RUN: -o - -verify-machineinstrs | FileCheck %s --check-prefix=MACHO
44+
45+
; RUN: llc %t/test.ll -mtriple=arm64-apple-darwin -aarch64-enable-collect-loh=0 \
46+
; RUN: -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
47+
; RUN: -code-model=large \
48+
; RUN: -o - -verify-machineinstrs | FileCheck %s --check-prefix=MACHO
49+
50+
; RUN: llc %t/test.ll -mtriple=arm64-apple-darwin -aarch64-enable-collect-loh=0 \
51+
; RUN: -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
52+
; RUN: -global-isel -global-isel-abort=1 \
53+
; RUN: -code-model=large \
54+
; RUN: -o - -verify-machineinstrs | FileCheck %s --check-prefix=MACHO
55+
56+
; RUN: llc %t/test.ll -mtriple=aarch64-elf \
57+
; RUN: -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
58+
; RUN: -o - -verify-machineinstrs | FileCheck %s --check-prefix=ELF
59+
60+
; RUN: llc %t/test.ll -mtriple=aarch64-elf \
61+
; RUN: -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
62+
; RUN: -global-isel -global-isel-abort=1 \
63+
; RUN: -o - -verify-machineinstrs | FileCheck %s --check-prefix=ELF
64+
65+
; MACHO-LABEL: test_jumptable:
66+
; MACHO: mov w16, w0
67+
; MACHO: cmp x16, #5
68+
; MACHO: csel x16, x16, xzr, ls
69+
; MACHO-NEXT: adrp x17, LJTI0_0@PAGE
70+
; MACHO-NEXT: add x17, x17, LJTI0_0@PAGEOFF
71+
; MACHO-NEXT: ldrsw x16, [x17, x16, lsl #2]
72+
; MACHO-NEXT: Ltmp0:
73+
; MACHO-NEXT: adr x17, Ltmp0
74+
; MACHO-NEXT: add x16, x17, x16
75+
; MACHO-NEXT: br x16
76+
77+
; ELF-LABEL: test_jumptable:
78+
; ELF: mov w16, w0
79+
; ELF: cmp x16, #5
80+
; ELF: csel x16, x16, xzr, ls
81+
; ELF-NEXT: adrp x17, .LJTI0_0
82+
; ELF-NEXT: add x17, x17, :lo12:.LJTI0_0
83+
; ELF-NEXT: ldrsw x16, [x17, x16, lsl #2]
84+
; ELF-NEXT: .Ltmp0:
85+
; ELF-NEXT: adr x17, .Ltmp0
86+
; ELF-NEXT: add x16, x17, x16
87+
; ELF-NEXT: br x16
88+
89+
define i32 @test_jumptable(i32 %in) "aarch64-jump-table-hardening" {
90+
91+
switch i32 %in, label %def [
92+
i32 0, label %lbl1
93+
i32 1, label %lbl2
94+
i32 2, label %lbl3
95+
i32 4, label %lbl4
96+
i32 5, label %lbl5
97+
]
98+
99+
def:
100+
ret i32 0
101+
102+
lbl1:
103+
ret i32 1
104+
105+
lbl2:
106+
ret i32 2
107+
108+
lbl3:
109+
ret i32 4
110+
111+
lbl4:
112+
ret i32 8
113+
114+
lbl5:
115+
ret i32 10
116+
117+
}
118+
119+
; MACHO-LABEL: LJTI0_0:
120+
; MACHO-NEXT: .long LBB{{[0-9_]+}}-Ltmp0
121+
; MACHO-NEXT: .long LBB{{[0-9_]+}}-Ltmp0
122+
; MACHO-NEXT: .long LBB{{[0-9_]+}}-Ltmp0
123+
; MACHO-NEXT: .long LBB{{[0-9_]+}}-Ltmp0
124+
; MACHO-NEXT: .long LBB{{[0-9_]+}}-Ltmp0
125+
; MACHO-NEXT: .long LBB{{[0-9_]+}}-Ltmp0
126+
127+
; ELF-LABEL: .LJTI0_0:
128+
; ELF-NEXT: .word .LBB{{[0-9_]+}}-.Ltmp0
129+
; ELF-NEXT: .word .LBB{{[0-9_]+}}-.Ltmp0
130+
; ELF-NEXT: .word .LBB{{[0-9_]+}}-.Ltmp0
131+
; ELF-NEXT: .word .LBB{{[0-9_]+}}-.Ltmp0
132+
; ELF-NEXT: .word .LBB{{[0-9_]+}}-.Ltmp0
133+
; ELF-NEXT: .word .LBB{{[0-9_]+}}-.Ltmp0

0 commit comments

Comments
 (0)