Skip to content

Commit 5c6ad3c

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 1e1844a commit 5c6ad3c

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

@@ -3440,6 +3512,16 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
34403512
return EmitZTInstr(MI, BB, AArch64::ZERO_T, /*Op0IsDef=*/true);
34413513
case AArch64::MOVT_TIZ_PSEUDO:
34423514
return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
3515+
3516+
case AArch64::PACDA:
3517+
case AArch64::PACDB:
3518+
case AArch64::PACIA:
3519+
case AArch64::PACIB:
3520+
case AArch64::PACDZA:
3521+
case AArch64::PACDZB:
3522+
case AArch64::PACIZA:
3523+
case AArch64::PACIZB:
3524+
return tryRewritingPAC(MI, BB);
34433525
}
34443526
}
34453527

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -691,6 +691,8 @@ class AArch64TargetLowering : public TargetLowering {
691691
MachineBasicBlock *BB) const;
692692
MachineBasicBlock *EmitGetSMESaveSize(MachineInstr &MI,
693693
MachineBasicBlock *BB) const;
694+
MachineBasicBlock *tryRewritingPAC(MachineInstr &MI,
695+
MachineBasicBlock *BB) const;
694696

695697
MachineBasicBlock *
696698
EmitInstrWithCustomInserter(MachineInstr &MI,

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

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

768+
static inline AArch64PACKey::ID getKeyForPACOpcode(unsigned Opcode) {
769+
switch (Opcode) {
770+
case AArch64::PACDA:
771+
case AArch64::PACDZA:
772+
return AArch64PACKey::DA;
773+
case AArch64::PACDB:
774+
case AArch64::PACDZB:
775+
return AArch64PACKey::DB;
776+
case AArch64::PACIA:
777+
case AArch64::PACIZA:
778+
return AArch64PACKey::IA;
779+
case AArch64::PACIB:
780+
case AArch64::PACIZB:
781+
return AArch64PACKey::IB;
782+
}
783+
llvm_unreachable("Unhandled PAC opcode");
784+
}
785+
786+
static inline bool isPACWithZeroDisc(unsigned Opcode) {
787+
switch (Opcode) {
788+
case AArch64::PACDA:
789+
case AArch64::PACDB:
790+
case AArch64::PACIA:
791+
case AArch64::PACIB:
792+
return false;
793+
case AArch64::PACDZA:
794+
case AArch64::PACDZB:
795+
case AArch64::PACIZA:
796+
case AArch64::PACIZB:
797+
return true;
798+
}
799+
llvm_unreachable("Unhandled PAC opcode");
800+
}
768801
// struct TSFlags {
769802
#define TSFLAG_ELEMENT_SIZE_TYPE(X) (X) // 3-bits
770803
#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
@@ -1821,7 +1821,9 @@ let Predicates = [HasPAuth] in {
18211821
def DZB : SignAuthZero<prefix_z, 0b11, !strconcat(asm, "dzb"), op>;
18221822
}
18231823

1824+
let usesCustomInserter = true in
18241825
defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
1826+
18251827
defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
18261828

18271829
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)