Skip to content

Commit 224c009

Browse files
committed
[AArch64][PAC] Lower auth/resign into checked sequence.
This introduces 3 hardening modes in the authentication step of auth/resign lowering: - unchecked, which uses the AUT instructions as-is - poison, which detects authentication failure (using an XPAC+CMP sequence), explicitly yielding the XPAC result rather than the AUT result, to avoid leaking - trap, which additionally traps on authentication failure, using BRK #0xC470 + key (IA C470, IB C471, DA C472, DB C473.) Not all modes are necessarily useful in all contexts, and there are more performant alternative lowerings in specific contexts (e.g., when I/D TBI enablement is a target ABI guarantee.) This is controlled by the `ptrauth-auth-traps` function attributes, and can be overridden using `-aarch64-ptrauth-auth-checks=`.
1 parent 30aa9fb commit 224c009

8 files changed

+1522
-0
lines changed

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

Lines changed: 272 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,16 @@
6767

6868
using namespace llvm;
6969

70+
enum PtrauthCheckMode { Default, Unchecked, Poison, Trap };
71+
static cl::opt<PtrauthCheckMode>
72+
PtrauthAuthChecks("aarch64-ptrauth-auth-checks", cl::Hidden,
73+
cl::values(
74+
clEnumValN(Unchecked, "none", "don't test for failure"),
75+
clEnumValN(Poison, "poison", "poison on failure"),
76+
clEnumValN(Trap, "trap", "trap on failure")),
77+
cl::desc("Check pointer authentication auth/resign failures"),
78+
cl::init(Default));
79+
7080
#define DEBUG_TYPE "asm-printer"
7181

7282
namespace {
@@ -124,6 +134,12 @@ class AArch64AsmPrinter : public AsmPrinter {
124134

125135
void emitSled(const MachineInstr &MI, SledKind Kind);
126136

137+
// Emit the sequence for AUT or AUTPAC.
138+
void emitPtrauthAuthResign(const MachineInstr *MI);
139+
// Emit the sequence to compute a discriminator into x17, or reuse AddrDisc.
140+
unsigned emitPtrauthDiscriminator(uint16_t Disc, unsigned AddrDisc,
141+
unsigned &InstsEmitted);
142+
127143
/// tblgen'erated driver function for lowering simple MI->MC
128144
/// pseudo instructions.
129145
bool emitPseudoExpansionLowering(MCStreamer &OutStreamer,
@@ -1417,6 +1433,257 @@ void AArch64AsmPrinter::emitFMov0(const MachineInstr &MI) {
14171433
}
14181434
}
14191435

1436+
unsigned AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
1437+
unsigned AddrDisc,
1438+
unsigned &InstsEmitted) {
1439+
// If there is no constant discriminator, there's no blend involved:
1440+
// just use the address discriminator register as-is (XZR or not).
1441+
if (!Disc)
1442+
return AddrDisc;
1443+
1444+
// If there's only a constant discriminator, MOV it into x17.
1445+
if (AddrDisc == AArch64::XZR) {
1446+
EmitToStreamer(*OutStreamer,
1447+
MCInstBuilder(AArch64::MOVZXi)
1448+
.addReg(AArch64::X17)
1449+
.addImm(Disc)
1450+
.addImm(/*shift=*/0));
1451+
++InstsEmitted;
1452+
return AArch64::X17;
1453+
}
1454+
1455+
// If there are both, emit a blend into x17.
1456+
EmitToStreamer(*OutStreamer,
1457+
MCInstBuilder(AArch64::ORRXrs)
1458+
.addReg(AArch64::X17)
1459+
.addReg(AArch64::XZR)
1460+
.addReg(AddrDisc)
1461+
.addImm(0));
1462+
++InstsEmitted;
1463+
EmitToStreamer(*OutStreamer,
1464+
MCInstBuilder(AArch64::MOVKXi)
1465+
.addReg(AArch64::X17)
1466+
.addReg(AArch64::X17)
1467+
.addImm(Disc)
1468+
.addImm(/*shift=*/48));
1469+
++InstsEmitted;
1470+
return AArch64::X17;
1471+
}
1472+
1473+
void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
1474+
unsigned InstsEmitted = 0;
1475+
const bool IsAUTPAC = MI->getOpcode() == AArch64::AUTPAC;
1476+
1477+
// We can expand AUT/AUTPAC into 3 possible sequences:
1478+
// - unchecked:
1479+
// autia x16, x0
1480+
// pacib x16, x1 ; if AUTPAC
1481+
//
1482+
// - checked and clearing:
1483+
// mov x17, x0
1484+
// movk x17, #disc, lsl #48
1485+
// autia x16, x17
1486+
// mov x17, x16
1487+
// xpaci x17
1488+
// cmp x16, x17
1489+
// b.eq Lsuccess
1490+
// mov x16, x17
1491+
// b Lend
1492+
// Lsuccess:
1493+
// mov x17, x1
1494+
// movk x17, #disc, lsl #48
1495+
// pacib x16, x17
1496+
// Lend:
1497+
// Where we only emit the AUT if we started with an AUT.
1498+
//
1499+
// - checked and trapping:
1500+
// mov x17, x0
1501+
// movk x17, #disc, lsl #48
1502+
// autia x16, x0
1503+
// mov x17, x16
1504+
// xpaci x17
1505+
// cmp x16, x17
1506+
// b.eq Lsuccess
1507+
// brk #<0xc470 + aut key>
1508+
// Lsuccess:
1509+
// mov x17, x1
1510+
// movk x17, #disc, lsl #48
1511+
// pacib x16, x17 ; if AUTPAC
1512+
// Where the b.eq skips over the trap if the PAC is valid.
1513+
//
1514+
// This sequence is expensive, but we need more information to be able to
1515+
// do better.
1516+
//
1517+
// We can't TBZ the poison bit because EnhancedPAC2 XORs the PAC bits
1518+
// on failure.
1519+
// We can't TST the PAC bits because we don't always know how the address
1520+
// space is setup for the target environment (and the bottom PAC bit is
1521+
// based on that).
1522+
// Either way, we also don't always know whether TBI is enabled or not for
1523+
// the specific target environment.
1524+
1525+
// By default, auth/resign sequences check for auth failures.
1526+
bool ShouldCheck = true;
1527+
// In the checked sequence, we only trap if explicitly requested.
1528+
bool ShouldTrap = MF->getFunction().hasFnAttribute("ptrauth-auth-traps");
1529+
1530+
// However, command-line flags can override this, for experimentation.
1531+
switch (PtrauthAuthChecks) {
1532+
case PtrauthCheckMode::Default: break;
1533+
case PtrauthCheckMode::Unchecked:
1534+
ShouldCheck = ShouldTrap = false;
1535+
break;
1536+
case PtrauthCheckMode::Poison:
1537+
ShouldCheck = true;
1538+
ShouldTrap = false;
1539+
break;
1540+
case PtrauthCheckMode::Trap:
1541+
ShouldCheck = ShouldTrap = true;
1542+
break;
1543+
}
1544+
1545+
auto AUTKey = (AArch64PACKey::ID)MI->getOperand(0).getImm();
1546+
uint64_t AUTDisc = MI->getOperand(1).getImm();
1547+
unsigned AUTAddrDisc = MI->getOperand(2).getReg();
1548+
1549+
unsigned XPACOpc = getXPACOpcodeForKey(AUTKey);
1550+
1551+
// Compute aut discriminator into x17
1552+
assert(isUInt<16>(AUTDisc));
1553+
unsigned AUTDiscReg = emitPtrauthDiscriminator(AUTDisc, AUTAddrDisc, InstsEmitted);
1554+
bool AUTZero = AUTDiscReg == AArch64::XZR;
1555+
unsigned AUTOpc = getAUTOpcodeForKey(AUTKey, AUTZero);
1556+
1557+
// autiza x16 ; if AUTZero
1558+
// autia x16, x17 ; if !AUTZero
1559+
MCInst AUTInst;
1560+
AUTInst.setOpcode(AUTOpc);
1561+
AUTInst.addOperand(MCOperand::createReg(AArch64::X16));
1562+
AUTInst.addOperand(MCOperand::createReg(AArch64::X16));
1563+
if (!AUTZero)
1564+
AUTInst.addOperand(MCOperand::createReg(AUTDiscReg));
1565+
EmitToStreamer(*OutStreamer, AUTInst);
1566+
++InstsEmitted;
1567+
1568+
// Unchecked or checked-but-non-trapping AUT is just an "AUT": we're done.
1569+
if (!IsAUTPAC && (!ShouldCheck || !ShouldTrap)) {
1570+
assert(STI->getInstrInfo()->getInstSizeInBytes(*MI) >= InstsEmitted * 4);
1571+
return;
1572+
}
1573+
1574+
MCSymbol *EndSym = nullptr;
1575+
1576+
// Checked sequences do an additional strip-and-compare.
1577+
if (ShouldCheck) {
1578+
MCSymbol *SuccessSym = createTempSymbol("auth_success_");
1579+
1580+
// XPAC has tied src/dst: use x17 as a temporary copy.
1581+
// mov x17, x16
1582+
EmitToStreamer(*OutStreamer,
1583+
MCInstBuilder(AArch64::ORRXrs)
1584+
.addReg(AArch64::X17)
1585+
.addReg(AArch64::XZR)
1586+
.addReg(AArch64::X16)
1587+
.addImm(0));
1588+
++InstsEmitted;
1589+
1590+
// xpaci x17
1591+
EmitToStreamer(*OutStreamer,
1592+
MCInstBuilder(XPACOpc)
1593+
.addReg(AArch64::X17)
1594+
.addReg(AArch64::X17));
1595+
++InstsEmitted;
1596+
1597+
// cmp x16, x17
1598+
EmitToStreamer(*OutStreamer,
1599+
MCInstBuilder(AArch64::SUBSXrs)
1600+
.addReg(AArch64::XZR)
1601+
.addReg(AArch64::X16)
1602+
.addReg(AArch64::X17)
1603+
.addImm(0));
1604+
++InstsEmitted;
1605+
1606+
// b.eq Lsuccess
1607+
EmitToStreamer(*OutStreamer,
1608+
MCInstBuilder(AArch64::Bcc)
1609+
.addImm(AArch64CC::EQ)
1610+
.addExpr(MCSymbolRefExpr::create(SuccessSym, OutContext)));
1611+
++InstsEmitted;
1612+
1613+
if (ShouldTrap) {
1614+
// Trapping sequences do a 'brk'.
1615+
// brk #<0xc470 + aut key>
1616+
EmitToStreamer(*OutStreamer,
1617+
MCInstBuilder(AArch64::BRK)
1618+
.addImm(0xc470 | AUTKey));
1619+
++InstsEmitted;
1620+
} else {
1621+
// Non-trapping checked sequences return the stripped result in x16,
1622+
// skipping over the PAC if there is one.
1623+
1624+
// FIXME: can we simply return the AUT result, already in x16? without..
1625+
// ..traps this is usable as an oracle anyway, based on high bits
1626+
// mov x17, x16
1627+
EmitToStreamer(*OutStreamer,
1628+
MCInstBuilder(AArch64::ORRXrs)
1629+
.addReg(AArch64::X16)
1630+
.addReg(AArch64::XZR)
1631+
.addReg(AArch64::X17)
1632+
.addImm(0));
1633+
++InstsEmitted;
1634+
1635+
if (IsAUTPAC) {
1636+
EndSym = createTempSymbol("resign_end_");
1637+
1638+
// b Lend
1639+
EmitToStreamer(*OutStreamer,
1640+
MCInstBuilder(AArch64::B)
1641+
.addExpr(MCSymbolRefExpr::create(EndSym, OutContext)));
1642+
++InstsEmitted;
1643+
}
1644+
}
1645+
1646+
// If the auth check succeeds, we can continue.
1647+
// Lsuccess:
1648+
OutStreamer->emitLabel(SuccessSym);
1649+
}
1650+
1651+
// We already emitted unchecked and checked-but-non-trapping AUTs.
1652+
// That left us with trapping AUTs, and AUTPACs.
1653+
// Trapping AUTs don't need PAC: we're done.
1654+
if (!IsAUTPAC) {
1655+
assert(STI->getInstrInfo()->getInstSizeInBytes(*MI) >= InstsEmitted * 4);
1656+
return;
1657+
}
1658+
1659+
auto PACKey = (AArch64PACKey::ID)MI->getOperand(3).getImm();
1660+
uint64_t PACDisc = MI->getOperand(4).getImm();
1661+
unsigned PACAddrDisc = MI->getOperand(5).getReg();
1662+
1663+
// Compute pac discriminator into x17
1664+
assert(isUInt<16>(PACDisc));
1665+
unsigned PACDiscReg =
1666+
emitPtrauthDiscriminator(PACDisc, PACAddrDisc, InstsEmitted);
1667+
bool PACZero = PACDiscReg == AArch64::XZR;
1668+
unsigned PACOpc = getPACOpcodeForKey(PACKey, PACZero);
1669+
1670+
// pacizb x16 ; if PACZero
1671+
// pacib x16, x17 ; if !PACZero
1672+
MCInst PACInst;
1673+
PACInst.setOpcode(PACOpc);
1674+
PACInst.addOperand(MCOperand::createReg(AArch64::X16));
1675+
PACInst.addOperand(MCOperand::createReg(AArch64::X16));
1676+
if (!PACZero)
1677+
PACInst.addOperand(MCOperand::createReg(PACDiscReg));
1678+
EmitToStreamer(*OutStreamer, PACInst);
1679+
++InstsEmitted;
1680+
1681+
assert(STI->getInstrInfo()->getInstSizeInBytes(*MI) >= InstsEmitted * 4);
1682+
// Lend:
1683+
if (EndSym)
1684+
OutStreamer->emitLabel(EndSym);
1685+
}
1686+
14201687
// Simple pseudo-instructions have their lowering (with expansion to real
14211688
// instructions) auto-generated.
14221689
#include "AArch64GenMCPseudoLowering.inc"
@@ -1552,6 +1819,11 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
15521819
return;
15531820
}
15541821

1822+
case AArch64::AUT:
1823+
case AArch64::AUTPAC:
1824+
emitPtrauthAuthResign(MI);
1825+
return;
1826+
15551827
// Tail calls use pseudo instructions so they have the proper code-gen
15561828
// attributes (isCall, isReturn, etc.). We lower them to the real
15571829
// instruction here.

0 commit comments

Comments
 (0)