Skip to content

Commit 6e04e1e

Browse files
atrosinenkoatrosinenko
authored
[AArch64][PAC] Introduce AArch64::PAC pseudo instruction (#146488)
Introduce a pseudo instruction carrying address and immediate modifiers as separate operands to be selected instead of a pair of `MOVKXi` and `PAC[ID][AB]` . The new pseudo instruction is expanded in AsmPrinter, so that `MOVKXi` is emitted immediately before `PAC[ID][AB]`. This way, an attacker cannot control the immediate modifier used to sign the value, even if address modifier can be substituted. To simplify the instruction selection, select `AArch64::PAC` pseudo using TableGen pattern and post-process its `$AddrDisc` operand by custom inserter hook - this eliminates duplication of the logic for DAGISel and GlobalISel. Furthermore, this improves cross-BB analysis in case of DAGISel.
1 parent 33f4582 commit 6e04e1e

File tree

6 files changed

+622
-1
lines changed

6 files changed

+622
-1
lines changed

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,9 @@ class AArch64AsmPrinter : public AsmPrinter {
176176
std::optional<AArch64PACKey::ID> PACKey,
177177
uint64_t PACDisc, Register PACAddrDisc);
178178

179+
// Emit the sequence for PAC.
180+
void emitPtrauthSign(const MachineInstr *MI);
181+
179182
// Emit the sequence to compute the discriminator.
180183
//
181184
// The returned register is either unmodified AddrDisc or ScratchReg.
@@ -2175,6 +2178,37 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
21752178
OutStreamer->emitLabel(EndSym);
21762179
}
21772180

2181+
void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) {
2182+
Register Val = MI->getOperand(1).getReg();
2183+
auto Key = (AArch64PACKey::ID)MI->getOperand(2).getImm();
2184+
uint64_t Disc = MI->getOperand(3).getImm();
2185+
Register AddrDisc = MI->getOperand(4).getReg();
2186+
bool AddrDiscKilled = MI->getOperand(4).isKill();
2187+
2188+
// As long as at least one of Val and AddrDisc is in GPR64noip, a scratch
2189+
// register is available.
2190+
Register ScratchReg = Val == AArch64::X16 ? AArch64::X17 : AArch64::X16;
2191+
assert(ScratchReg != AddrDisc &&
2192+
"Neither X16 nor X17 is available as a scratch register");
2193+
2194+
// Compute pac discriminator
2195+
assert(isUInt<16>(Disc));
2196+
Register DiscReg = emitPtrauthDiscriminator(
2197+
Disc, AddrDisc, ScratchReg, /*MayUseAddrAsScratch=*/AddrDiscKilled);
2198+
bool IsZeroDisc = DiscReg == AArch64::XZR;
2199+
unsigned Opc = getPACOpcodeForKey(Key, IsZeroDisc);
2200+
2201+
// paciza x16 ; if IsZeroDisc
2202+
// pacia x16, x17 ; if !IsZeroDisc
2203+
MCInst PACInst;
2204+
PACInst.setOpcode(Opc);
2205+
PACInst.addOperand(MCOperand::createReg(Val));
2206+
PACInst.addOperand(MCOperand::createReg(Val));
2207+
if (!IsZeroDisc)
2208+
PACInst.addOperand(MCOperand::createReg(DiscReg));
2209+
EmitToStreamer(*OutStreamer, PACInst);
2210+
}
2211+
21782212
void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
21792213
bool IsCall = MI->getOpcode() == AArch64::BLRA;
21802214
unsigned BrTarget = MI->getOperand(0).getReg();
@@ -2890,6 +2924,10 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
28902924
MI->getOperand(4).getImm(), MI->getOperand(5).getReg());
28912925
return;
28922926

2927+
case AArch64::PAC:
2928+
emitPtrauthSign(MI);
2929+
return;
2930+
28932931
case AArch64::LOADauthptrstatic:
28942932
LowerLOADauthptrstatic(*MI);
28952933
return;

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

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3101,6 +3101,83 @@ AArch64TargetLowering::EmitGetSMESaveSize(MachineInstr &MI,
31013101
return BB;
31023102
}
31033103

3104+
// Helper function to find the instruction that defined a virtual register.
3105+
// If unable to find such instruction, returns nullptr.
3106+
static const MachineInstr *stripVRegCopies(const MachineRegisterInfo &MRI,
3107+
Register Reg) {
3108+
while (Reg.isVirtual()) {
3109+
MachineInstr *DefMI = MRI.getVRegDef(Reg);
3110+
assert(DefMI && "Virtual register definition not found");
3111+
unsigned Opcode = DefMI->getOpcode();
3112+
3113+
if (Opcode == AArch64::COPY) {
3114+
Reg = DefMI->getOperand(1).getReg();
3115+
// Vreg is defined by copying from physreg.
3116+
if (Reg.isPhysical())
3117+
return DefMI;
3118+
continue;
3119+
}
3120+
if (Opcode == AArch64::SUBREG_TO_REG) {
3121+
Reg = DefMI->getOperand(2).getReg();
3122+
continue;
3123+
}
3124+
3125+
return DefMI;
3126+
}
3127+
return nullptr;
3128+
}
3129+
3130+
void AArch64TargetLowering::fixupPtrauthDiscriminator(
3131+
MachineInstr &MI, MachineBasicBlock *BB, MachineOperand &IntDiscOp,
3132+
MachineOperand &AddrDiscOp, const TargetRegisterClass *AddrDiscRC) const {
3133+
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
3134+
MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
3135+
const DebugLoc &DL = MI.getDebugLoc();
3136+
3137+
Register AddrDisc = AddrDiscOp.getReg();
3138+
int64_t IntDisc = IntDiscOp.getImm();
3139+
assert(IntDisc == 0 && "Blend components are already expanded");
3140+
3141+
const MachineInstr *DiscMI = stripVRegCopies(MRI, AddrDisc);
3142+
if (DiscMI) {
3143+
switch (DiscMI->getOpcode()) {
3144+
case AArch64::MOVKXi:
3145+
// blend(addr, imm) which is lowered as "MOVK addr, #imm, #48".
3146+
// #imm should be an immediate and not a global symbol, for example.
3147+
if (DiscMI->getOperand(2).isImm() &&
3148+
DiscMI->getOperand(3).getImm() == 48) {
3149+
AddrDisc = DiscMI->getOperand(1).getReg();
3150+
IntDisc = DiscMI->getOperand(2).getImm();
3151+
}
3152+
break;
3153+
case AArch64::MOVi32imm:
3154+
case AArch64::MOVi64imm:
3155+
// Small immediate integer constant passed via VReg.
3156+
if (DiscMI->getOperand(1).isImm() &&
3157+
isUInt<16>(DiscMI->getOperand(1).getImm())) {
3158+
AddrDisc = AArch64::NoRegister;
3159+
IntDisc = DiscMI->getOperand(1).getImm();
3160+
}
3161+
break;
3162+
}
3163+
}
3164+
3165+
// For uniformity, always use NoRegister, as XZR is not necessarily contained
3166+
// in the requested register class.
3167+
if (AddrDisc == AArch64::XZR)
3168+
AddrDisc = AArch64::NoRegister;
3169+
3170+
// Make sure AddrDisc operand respects the register class imposed by MI.
3171+
if (AddrDisc && MRI.getRegClass(AddrDisc) != AddrDiscRC) {
3172+
Register TmpReg = MRI.createVirtualRegister(AddrDiscRC);
3173+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), TmpReg).addReg(AddrDisc);
3174+
AddrDisc = TmpReg;
3175+
}
3176+
3177+
AddrDiscOp.setReg(AddrDisc);
3178+
IntDiscOp.setImm(IntDisc);
3179+
}
3180+
31043181
MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
31053182
MachineInstr &MI, MachineBasicBlock *BB) const {
31063183

@@ -3199,6 +3276,11 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
31993276
return EmitZTInstr(MI, BB, AArch64::ZERO_T, /*Op0IsDef=*/true);
32003277
case AArch64::MOVT_TIZ_PSEUDO:
32013278
return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
3279+
3280+
case AArch64::PAC:
3281+
fixupPtrauthDiscriminator(MI, BB, MI.getOperand(3), MI.getOperand(4),
3282+
&AArch64::GPR64noipRegClass);
3283+
return BB;
32023284
}
32033285
}
32043286

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,13 @@ class AArch64TargetLowering : public TargetLowering {
182182
MachineBasicBlock *EmitGetSMESaveSize(MachineInstr &MI,
183183
MachineBasicBlock *BB) const;
184184

185+
/// Replace (0, vreg) discriminator components with the operands of blend
186+
/// or with (immediate, NoRegister) when possible.
187+
void fixupPtrauthDiscriminator(MachineInstr &MI, MachineBasicBlock *BB,
188+
MachineOperand &IntDiscOp,
189+
MachineOperand &AddrDiscOp,
190+
const TargetRegisterClass *AddrDiscRC) const;
191+
185192
MachineBasicBlock *
186193
EmitInstrWithCustomInserter(MachineInstr &MI,
187194
MachineBasicBlock *MBB) const override;

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2033,7 +2033,7 @@ let Predicates = [HasPAuth] in {
20332033
def DZB : SignAuthZero<prefix_z, 0b11, !strconcat(asm, "dzb"), op>;
20342034
}
20352035

2036-
defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
2036+
defm PAC : SignAuth<0b000, 0b010, "pac", null_frag>;
20372037
defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
20382038

20392039
def XPACI : ClearAuth<0, "xpaci">;
@@ -2153,6 +2153,26 @@ let Predicates = [HasPAuth] in {
21532153
let Uses = [];
21542154
}
21552155

2156+
// PAC pseudo instruction. In AsmPrinter, it is expanded into an actual PAC*
2157+
// instruction immediately preceded by the discriminator computation.
2158+
// This enforces the expected immediate modifier is used for signing, even
2159+
// if an attacker is able to substitute AddrDisc.
2160+
def PAC : Pseudo<(outs GPR64:$SignedVal),
2161+
(ins GPR64:$Val, i32imm:$Key, i64imm:$Disc, GPR64noip:$AddrDisc),
2162+
[], "$SignedVal = $Val">, Sched<[WriteI, ReadI]> {
2163+
let isCodeGenOnly = 1;
2164+
let hasSideEffects = 0;
2165+
let mayStore = 0;
2166+
let mayLoad = 0;
2167+
let Size = 12;
2168+
let Defs = [X16, X17];
2169+
let usesCustomInserter = 1;
2170+
}
2171+
2172+
// A standalone pattern is used, so that literal 0 can be passed as $Disc.
2173+
def : Pat<(int_ptrauth_sign GPR64:$Val, timm:$Key, GPR64noip:$AddrDisc),
2174+
(PAC GPR64:$Val, $Key, 0, GPR64noip:$AddrDisc)>;
2175+
21562176
// AUT and re-PAC a value, using different keys/data.
21572177
// This directly manipulates x16/x17, which are the only registers that
21582178
// certain OSs guarantee are safe to use for sensitive operations.

0 commit comments

Comments
 (0)