Skip to content

Commit 08a382e

Browse files
atrosinenkoatrosinenko
committed
[AArch64][PAC] Combine signing with address materialization
In pauthtest ABI, it is common to store a pointer signed with address diversity to a heap-allocated object (namely, storing a signed pointer to VTable as part of new object construction). This patch tries to prevent introducing a signing oracle by combining pointer materialization and its (re)signing into a single pseudo instruction which is not expanded until AsmPrinter, if possible. One of the typical patterns is materializing an unsigned pointer with `MOVaddr` pseudo and then signing it with `PAC[ID][AB]` instruction, which can be moved far away from `MOVaddr` by one of the passes in the machine pipeline. As the storage address is not a `Constant` value, one cannot simply emit a `ptrauth` constant in the frontend, which would be selected into `MOVaddrPAC` pseudo. Another pattern is fetching a pointer to VTable from a signed GOT entry using `LOADgotAUTH` pseudo, authenticating and checking it, and then re-signing after adding an offset. This commit adds an instruction insertion hook for `PAC[ID][AB]` which detects the above patterns and replaces it either with `MOVaddrPAC` or `LOADgotPAC` instruction.
1 parent e10b8cb commit 08a382e

File tree

6 files changed

+131
-4
lines changed

6 files changed

+131
-4
lines changed

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

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3349,6 +3349,78 @@ AArch64TargetLowering::EmitGetSMESaveSize(MachineInstr &MI,
33493349
MI.getOperand(0).getReg())
33503350
.addReg(AArch64::XZR);
33513351
BB->remove_instr(&MI);
3352+
3353+
return BB;
3354+
}
3355+
3356+
MachineBasicBlock *
3357+
AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
3358+
MachineBasicBlock *BB) const {
3359+
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
3360+
MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
3361+
const DebugLoc &DL = MI.getDebugLoc();
3362+
3363+
MachineInstr *AddrInstr = nullptr;
3364+
int64_t Offset = 0;
3365+
// Try to find the address-setting instruction, accumulating the offset
3366+
// along the way. If any unknown pattern is found, keep everything as-is.
3367+
MachineOperand *CurOp = &MI.getOperand(1);
3368+
while (CurOp) {
3369+
MachineOperand *Def = MRI.getOneDef(CurOp->getReg());
3370+
if (!Def)
3371+
return BB;
3372+
MachineInstr *DefMI = Def->getParent();
3373+
assert(DefMI != nullptr);
3374+
3375+
switch (DefMI->getOpcode()) {
3376+
case AArch64::COPY:
3377+
CurOp = &DefMI->getOperand(1);
3378+
break;
3379+
case AArch64::ADDXri:
3380+
if (DefMI->getOperand(3).getImm() != 0) // shifts are not handled
3381+
return BB;
3382+
CurOp = &DefMI->getOperand(1);
3383+
Offset += DefMI->getOperand(2).getImm();
3384+
break;
3385+
case AArch64::MOVaddr:
3386+
case AArch64::LOADgotAUTH:
3387+
AddrInstr = DefMI;
3388+
CurOp = nullptr;
3389+
break;
3390+
default:
3391+
return BB;
3392+
}
3393+
}
3394+
3395+
unsigned NewOpcode = AddrInstr->getOpcode() == AArch64::LOADgotAUTH
3396+
? AArch64::LOADgotPAC
3397+
: AArch64::MOVaddrPAC;
3398+
MachineOperand &AddrOp = AddrInstr->getOperand(1);
3399+
unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
3400+
Offset += AddrOp.getOffset();
3401+
3402+
// MOVaddrPAC and LOADgotPAC pseudos are expanded so that they use X16/X17
3403+
// internally, thus their restrictions on the register class of $AddrDisc
3404+
// operand are stricter than those of real PAC* instructions.
3405+
// If the original instruction accepts a discriminator operand, make sure
3406+
// it is moved out of X16/X17.
3407+
Register DiscReg = AArch64::XZR;
3408+
if (!isPACWithZeroDisc(MI.getOpcode())) {
3409+
DiscReg = MRI.createVirtualRegister(&AArch64::GPR64noipRegClass);
3410+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), DiscReg)
3411+
.addReg(MI.getOperand(2).getReg());
3412+
}
3413+
3414+
BuildMI(*BB, MI, DL, TII->get(NewOpcode))
3415+
.addGlobalAddress(AddrOp.getGlobal(), Offset, TargetFlags)
3416+
.addImm(getKeyForPACOpcode(MI.getOpcode()))
3417+
.addReg(DiscReg)
3418+
.addImm(0);
3419+
3420+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
3421+
.addReg(AArch64::X16);
3422+
3423+
MI.removeFromParent();
33523424
return BB;
33533425
}
33543426

@@ -3450,6 +3522,16 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
34503522
return EmitZTInstr(MI, BB, AArch64::ZERO_T, /*Op0IsDef=*/true);
34513523
case AArch64::MOVT_TIZ_PSEUDO:
34523524
return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
3525+
3526+
case AArch64::PACDA:
3527+
case AArch64::PACDB:
3528+
case AArch64::PACIA:
3529+
case AArch64::PACIB:
3530+
case AArch64::PACDZA:
3531+
case AArch64::PACDZB:
3532+
case AArch64::PACIZA:
3533+
case AArch64::PACIZB:
3534+
return tryRewritingPAC(MI, BB);
34533535
}
34543536
}
34553537

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -679,6 +679,8 @@ class AArch64TargetLowering : public TargetLowering {
679679
MachineBasicBlock *BB) const;
680680
MachineBasicBlock *EmitGetSMESaveSize(MachineInstr &MI,
681681
MachineBasicBlock *BB) const;
682+
MachineBasicBlock *tryRewritingPAC(MachineInstr &MI,
683+
MachineBasicBlock *BB) const;
682684

683685
MachineBasicBlock *
684686
EmitInstrWithCustomInserter(MachineInstr &MI,

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -790,6 +790,39 @@ static inline unsigned getPACOpcodeForKey(AArch64PACKey::ID K, bool Zero) {
790790
llvm_unreachable("Unhandled AArch64PACKey::ID enum");
791791
}
792792

793+
static inline AArch64PACKey::ID getKeyForPACOpcode(unsigned Opcode) {
794+
switch (Opcode) {
795+
case AArch64::PACDA:
796+
case AArch64::PACDZA:
797+
return AArch64PACKey::DA;
798+
case AArch64::PACDB:
799+
case AArch64::PACDZB:
800+
return AArch64PACKey::DB;
801+
case AArch64::PACIA:
802+
case AArch64::PACIZA:
803+
return AArch64PACKey::IA;
804+
case AArch64::PACIB:
805+
case AArch64::PACIZB:
806+
return AArch64PACKey::IB;
807+
}
808+
llvm_unreachable("Unhandled PAC opcode");
809+
}
810+
811+
static inline bool isPACWithZeroDisc(unsigned Opcode) {
812+
switch (Opcode) {
813+
case AArch64::PACDA:
814+
case AArch64::PACDB:
815+
case AArch64::PACIA:
816+
case AArch64::PACIB:
817+
return false;
818+
case AArch64::PACDZA:
819+
case AArch64::PACDZB:
820+
case AArch64::PACIZA:
821+
case AArch64::PACIZB:
822+
return true;
823+
}
824+
llvm_unreachable("Unhandled PAC opcode");
825+
}
793826
// struct TSFlags {
794827
#define TSFLAG_ELEMENT_SIZE_TYPE(X) (X) // 3-bits
795828
#define TSFLAG_DESTRUCTIVE_INST_TYPE(X) ((X) << 3) // 4-bits

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1845,7 +1845,9 @@ let Predicates = [HasPAuth] in {
18451845
def DZB : SignAuthZero<prefix_z, 0b11, !strconcat(asm, "dzb"), op>;
18461846
}
18471847

1848+
let usesCustomInserter = true in
18481849
defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
1850+
18491851
defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
18501852

18511853
def XPACI : ClearAuth<0, "xpaci">;

‎llvm/test/CodeGen/AArch64/GlobalISel/ptrauth-constant-in-code.ll‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,9 @@ define void @store_signed_const_local(ptr %dest) {
9696
; ISEL: %0:gpr64common = COPY $x0
9797
; ISEL-NEXT: %10:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
9898
; ISEL-NEXT: %2:gpr64common = PAUTH_BLEND %0, 1234
99-
; ISEL-NEXT: %4:gpr64 = PACDA %10, %2
99+
; ISEL-NEXT: %15:gpr64noip = COPY %2
100+
; ISEL-NEXT: MOVaddrPAC @const_table_local + 8, 2, %15, 0, implicit-def $x16, implicit-def $x17
101+
; ISEL-NEXT: %4:gpr64 = COPY $x16
100102
; ISEL-NEXT: %14:gpr64 = COPY %4
101103
; ISEL-NEXT: STRXui %14, %0, 0 :: (store (p0) into %ir.dest)
102104
; ISEL-NEXT: RET_ReallyLR
@@ -115,7 +117,9 @@ define void @store_signed_const_got(ptr %dest) {
115117
; ISEL-ELF-NEXT: %7:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got
116118
; ISEL-ELF-NEXT: %6:gpr64common = ADDXri %7, 8, 0
117119
; ISEL-ELF-NEXT: %2:gpr64common = PAUTH_BLEND %0, 1234
118-
; ISEL-ELF-NEXT: %4:gpr64 = PACDA %6, %2
120+
; ISEL-ELF-NEXT: %12:gpr64noip = COPY %2
121+
; ISEL-ELF-NEXT: LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %12, 0, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
122+
; ISEL-ELF-NEXT: %4:gpr64 = COPY $x16
119123
; ISEL-ELF-NEXT: %10:gpr64 = COPY %4
120124
; ISEL-ELF-NEXT: STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
121125
; ISEL-ELF-NEXT: RET_ReallyLR

‎llvm/test/CodeGen/AArch64/ptrauth-constant-in-code.ll‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,9 @@ define void @store_signed_const_local(ptr %dest) {
8585
; ISEL: %0:gpr64common = COPY $x0
8686
; ISEL-NEXT: %1:gpr64common = PAUTH_BLEND %0, 1234
8787
; ISEL-NEXT: %2:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
88-
; ISEL-NEXT: %3:gpr64 = PACDA %2, killed %1
88+
; ISEL-NEXT: %4:gpr64noip = COPY %1
89+
; ISEL-NEXT: MOVaddrPAC @const_table_local + 8, 2, %4, 0, implicit-def $x16, implicit-def $x17
90+
; ISEL-NEXT: %3:gpr64 = COPY $x16
8991
; ISEL-NEXT: STRXui killed %3, %0, 0 :: (store (s64) into %ir.dest)
9092
; ISEL-NEXT: RET_ReallyLR
9193
%dest.i = ptrtoint ptr %dest to i64
@@ -103,7 +105,9 @@ define void @store_signed_const_got(ptr %dest) {
103105
; ISEL-ELF-NEXT: %1:gpr64common = PAUTH_BLEND %0, 1234
104106
; ISEL-ELF-NEXT: %2:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got, implicit-def dead $x16, implicit-def dead $x17, implicit-def dead $nzcv
105107
; ISEL-ELF-NEXT: %3:gpr64common = ADDXri killed %2, 8, 0
106-
; ISEL-ELF-NEXT: %4:gpr64 = PACDA %3, killed %1
108+
; ISEL-ELF-NEXT: %5:gpr64noip = COPY %1
109+
; ISEL-ELF-NEXT: LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %5, 0, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
110+
; ISEL-ELF-NEXT: %4:gpr64 = COPY $x16
107111
; ISEL-ELF-NEXT: STRXui killed %4, %0, 0 :: (store (s64) into %ir.dest)
108112
; ISEL-ELF-NEXT: RET_ReallyLR
109113
%dest.i = ptrtoint ptr %dest to i64

0 commit comments

Comments
 (0)