Skip to content

[AArch64][PAC] Rework the expansion of AUT/AUTPAC pseudos #169699

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: users/atrosinenko/pauth-factor-out-low-level-emission
Choose a base branch
from
+192 −75
Open

[AArch64][PAC] Rework the expansion of AUT/AUTPAC pseudos #169699

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 113 additions & 64 deletions llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,13 +179,21 @@ class AArch64AsmPrinter : public AsmPrinter {
// Check authenticated LR before tail calling.
void emitPtrauthTailCallHardening(const MachineInstr *TC);

struct PtrAuthSchema {
PtrAuthSchema(AArch64PACKey::ID Key, uint64_t Disc,
const MachineOperand &AddrDiscOp);

AArch64PACKey::ID Key;
uint64_t Disc;
Register AddrDisc;
bool AddrDiscIsKilled;
};

// Emit the sequence for AUT or AUTPAC.
void emitPtrauthAuthResign(Register AUTVal, AArch64PACKey::ID AUTKey,
uint64_t AUTDisc,
const MachineOperand *AUTAddrDisc,
Register Scratch,
std::optional<AArch64PACKey::ID> PACKey,
uint64_t PACDisc, Register PACAddrDisc);
void
emitPtrauthAuthResign(Register Pointer, Register Scratch,
PtrAuthSchema AuthSchema,
std::optional<PtrAuthSchema> SignSchema = std::nullopt);

// Emit the sequence for PAC.
void emitPtrauthSign(const MachineInstr *MI);
Expand Down Expand Up @@ -2175,23 +2183,9 @@ void AArch64AsmPrinter::emitPtrauthTailCallHardening(const MachineInstr *TC) {
LRCheckMethod);
}

void AArch64AsmPrinter::emitPtrauthAuthResign(
Register AUTVal, AArch64PACKey::ID AUTKey, uint64_t AUTDisc,
const MachineOperand *AUTAddrDisc, Register Scratch,
std::optional<AArch64PACKey::ID> PACKey, uint64_t PACDisc,
Register PACAddrDisc) {
const bool IsAUTPAC = PACKey.has_value();

// We expand AUT/AUTPAC into a sequence of the form
//
// ; authenticate x16
// ; check pointer in x16
// Lsuccess:
// ; sign x16 (if AUTPAC)
// Lend: ; if not trapping on failure
//
// with the checking sequence chosen depending on whether/how we should check
// the pointer and whether we should trap on failure.
static std::pair<bool, bool> getCheckAndTrapMode(const MachineFunction *MF,
bool IsResign) {
const AArch64Subtarget &STI = MF->getSubtarget<AArch64Subtarget>();

// By default, auth/resign sequences check for auth failures.
bool ShouldCheck = true;
Expand All @@ -2200,7 +2194,7 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(

// On an FPAC CPU, you get traps whether you want them or not: there's
// no point in emitting checks or traps.
if (STI->hasFPAC())
if (STI.hasFPAC())
ShouldCheck = ShouldTrap = false;

// However, command-line flags can override this, for experimentation.
Expand All @@ -2219,38 +2213,79 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
break;
}

// Compute aut discriminator
Register AUTDiscReg = emitPtrauthDiscriminator(
AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill());
emitAUT(AUTKey, AUTVal, AUTDiscReg);
// Checked-but-not-trapping mode ("poison") only applies to resigning,
// replace with "unchecked" for standalone AUT.
if (!IsResign && ShouldCheck && !ShouldTrap)
ShouldCheck = ShouldTrap = false;

// Unchecked or checked-but-non-trapping AUT is just an "AUT": we're done.
if (!IsAUTPAC && (!ShouldCheck || !ShouldTrap))
return;
return std::make_pair(ShouldCheck, ShouldTrap);
}

MCSymbol *EndSym = nullptr;
AArch64AsmPrinter::PtrAuthSchema::PtrAuthSchema(
AArch64PACKey::ID Key, uint64_t Disc, const MachineOperand &AddrDiscOp)
: Key(Key), Disc(Disc), AddrDisc(AddrDiscOp.getReg()),
AddrDiscIsKilled(AddrDiscOp.isKill()) {}

if (ShouldCheck) {
if (IsAUTPAC && !ShouldTrap)
EndSym = createTempSymbol("resign_end_");
// We expand AUTx16x17/AUTxMxN into a sequence of the form
//
// ; authenticate Pointer
// ; check that Pointer is valid (optional, traps on failure)
//
// We expand AUTPAC into a sequence of the form
//
// ; authenticate Pointer
// ; check that Pointer is valid (optional, traps on failure)
// ; sign Pointer
//
// or
//
// ; authenticate Pointer
// ; check that Pointer is valid (skips re-sign on failure)
// ; sign Pointer
// Lon_failure:
//
void AArch64AsmPrinter::emitPtrauthAuthResign(
Register Pointer, Register Scratch, PtrAuthSchema AuthSchema,
std::optional<PtrAuthSchema> SignSchema) {
const bool IsResign = SignSchema.has_value();

const auto [ShouldCheck, ShouldTrap] = getCheckAndTrapMode(MF, IsResign);
const bool ShouldSkipSignOnAuthFailure = ShouldCheck && !ShouldTrap;
assert((ShouldCheck || !ShouldTrap) && "ShouldTrap implies ShouldCheck");

// It is hardly meaningful to authenticate or sign a pointer using its own
// value, thus we only have to take care not to early-clobber
// AuthSchema.AddrDisc that is aliased with SignSchema->AddrDisc.
assert(Pointer != AuthSchema.AddrDisc);
assert(!SignSchema || Pointer != SignSchema->AddrDisc);
bool IsResignWithAliasedAddrDiscs =
IsResign && AuthSchema.AddrDisc == SignSchema->AddrDisc;
bool MayReuseAUTAddrDisc =
!IsResignWithAliasedAddrDiscs && AuthSchema.AddrDiscIsKilled;
Register AUTDiscReg = emitPtrauthDiscriminator(
AuthSchema.Disc, AuthSchema.AddrDisc, Scratch, MayReuseAUTAddrDisc);
emitAUT(AuthSchema.Key, Pointer, AUTDiscReg);

emitPtrauthCheckAuthenticatedValue(
AUTVal, Scratch, AUTKey, AArch64PAuth::AuthCheckMethod::XPAC, EndSym);
}
MCSymbol *OnFailure =
ShouldSkipSignOnAuthFailure ? createTempSymbol("resign_end_") : nullptr;

if (ShouldCheck)
emitPtrauthCheckAuthenticatedValue(Pointer, Scratch, AuthSchema.Key,
AArch64PAuth::AuthCheckMethod::XPAC,
OnFailure);

// We already emitted unchecked and checked-but-non-trapping AUTs.
// That left us with trapping AUTs, and AUTPACs.
// Trapping AUTs don't need PAC: we're done.
if (!IsAUTPAC)
if (!IsResign) {
assert(!OnFailure && "Poison mode only applies to resigning");
return;
}

// Compute pac discriminator
Register PACDiscReg = emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
emitPAC(*PACKey, AUTVal, PACDiscReg);
Register PACDiscReg =
emitPtrauthDiscriminator(SignSchema->Disc, SignSchema->AddrDisc, Scratch,
SignSchema->AddrDiscIsKilled);
emitPAC(SignSchema->Key, Pointer, PACDiscReg);

// Lend:
if (EndSym)
OutStreamer->emitLabel(EndSym);
if (OnFailure)
OutStreamer->emitLabel(OnFailure);
}

void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) {
Expand Down Expand Up @@ -2936,27 +2971,41 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
return;
}

case AArch64::AUTx16x17:
emitPtrauthAuthResign(AArch64::X16,
(AArch64PACKey::ID)MI->getOperand(0).getImm(),
MI->getOperand(1).getImm(), &MI->getOperand(2),
AArch64::X17, std::nullopt, 0, 0);
case AArch64::AUTx16x17: {
const Register Pointer = AArch64::X16;
const Register Scratch = AArch64::X17;

PtrAuthSchema AuthSchema((AArch64PACKey::ID)MI->getOperand(0).getImm(),
MI->getOperand(1).getImm(), MI->getOperand(2));

emitPtrauthAuthResign(Pointer, Scratch, AuthSchema);
return;
}

case AArch64::AUTxMxN: {
const Register Pointer = MI->getOperand(0).getReg();
const Register Scratch = MI->getOperand(1).getReg();

case AArch64::AUTxMxN:
emitPtrauthAuthResign(MI->getOperand(0).getReg(),
(AArch64PACKey::ID)MI->getOperand(3).getImm(),
MI->getOperand(4).getImm(), &MI->getOperand(5),
MI->getOperand(1).getReg(), std::nullopt, 0, 0);
PtrAuthSchema AuthSchema((AArch64PACKey::ID)MI->getOperand(3).getImm(),
MI->getOperand(4).getImm(), MI->getOperand(5));

emitPtrauthAuthResign(Pointer, Scratch, AuthSchema);
return;
}

case AArch64::AUTPAC:
emitPtrauthAuthResign(
AArch64::X16, (AArch64PACKey::ID)MI->getOperand(0).getImm(),
MI->getOperand(1).getImm(), &MI->getOperand(2), AArch64::X17,
(AArch64PACKey::ID)MI->getOperand(3).getImm(),
MI->getOperand(4).getImm(), MI->getOperand(5).getReg());
case AArch64::AUTPAC: {
const Register Pointer = AArch64::X16;
const Register Scratch = AArch64::X17;

PtrAuthSchema AuthSchema((AArch64PACKey::ID)MI->getOperand(0).getImm(),
MI->getOperand(1).getImm(), MI->getOperand(2));

PtrAuthSchema SignSchema((AArch64PACKey::ID)MI->getOperand(3).getImm(),
MI->getOperand(4).getImm(), MI->getOperand(5));

emitPtrauthAuthResign(Pointer, Scratch, AuthSchema, SignSchema);
return;
}

case AArch64::PAC:
emitPtrauthSign(MI);
Expand Down
13 changes: 12 additions & 1 deletion llvm/lib/Target/AArch64/AArch64InstrInfo.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2188,10 +2188,21 @@ let Predicates = [HasPAuth] in {
let Uses = [X16];
}

// AsmPrinter can clobber the $AddrDisc register as long as it is marked
// "killed", otherwise an assertion checks that $Scratch != $AddrDisc.
// The problem is that it is always correct to *omit* the "killed" flag,
// thus an instruction like this is can be passed to AsmPrinter:
//
// $x8, $x9 = AUTxMxN $x8, 0, 51756, $x9, implicit-def $nzcv
//
// While it is possible to check if "$Scratch == $AddrDisc OR $AddrDisc
// is killed" instead, it is easier and more straightforward to mark
// $Scratch as @earlyclobber. Furthermore, this would be mandatory for a
// (hypotetical at the time of writing this comment) AUTPACxMxN pseudo.
def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64common:$Scratch),
(ins GPR64:$Val, i32imm:$Key,
i64imm:$Disc, GPR64:$AddrDisc),
[], "$AuthVal = $Val">, Sched<[WriteI, ReadI]> {
[], "$AuthVal = $Val,@earlyclobber $Scratch">, Sched<[WriteI, ReadI]> {
let isCodeGenOnly = 1;
let hasSideEffects = 1;
let mayStore = 0;
Expand Down
77 changes: 67 additions & 10 deletions llvm/test/CodeGen/AArch64/ptrauth-intrinsic-auth-resign-with-blend.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
; RUN: llc < %s -mtriple arm64e-apple-darwin -global-isel=0 -verify-machineinstrs \
; RUN: -aarch64-ptrauth-auth-checks=none | FileCheck %s -DL="L" --check-prefixes=UNCHECKED,UNCHECKED-DARWIN
; RUN: llc < %s -mtriple arm64e-apple-darwin -global-isel -global-isel-abort=1 -verify-machineinstrs \
Expand Down Expand Up @@ -84,12 +83,14 @@ define i64 @test_resign_blend(i64 %arg, i64 %arg1, i64 %arg2) {
; UNCHECKED-NEXT: mov x16, x0
; UNCHECKED-ELF-NEXT: movk x1, #12345, lsl #48
; UNCHECKED-ELF-NEXT: autda x16, x1
; UNCHECKED-ELF-NEXT: movk x2, #56789, lsl #48
; UNCHECKED-ELF-NEXT: pacdb x16, x2
; UNCHECKED-DARWIN-NEXT: mov x17, x1
; UNCHECKED-DARWIN-NEXT: movk x17, #12345, lsl #48
; UNCHECKED-DARWIN-NEXT: autda x16, x17
; UNCHECKED-NEXT: mov x17, x2
; UNCHECKED-NEXT: movk x17, #56789, lsl #48
; UNCHECKED-NEXT: pacdb x16, x17
; UNCHECKED-DARWIN-NEXT: mov x17, x2
; UNCHECKED-DARWIN-NEXT: movk x17, #56789, lsl #48
; UNCHECKED-DARWIN-NEXT: pacdb x16, x17
; UNCHECKED-NEXT: mov x0, x16
; UNCHECKED-NEXT: ret
;
Expand All @@ -108,9 +109,11 @@ define i64 @test_resign_blend(i64 %arg, i64 %arg1, i64 %arg2) {
; CHECKED-NEXT: mov x16, x17
; CHECKED-NEXT: b [[L]]resign_end_0
; CHECKED-NEXT: Lauth_success_0:
; CHECKED-NEXT: mov x17, x2
; CHECKED-NEXT: movk x17, #56789, lsl #48
; CHECKED-NEXT: pacdb x16, x17
; CHECKED-ELF-NEXT: movk x2, #56789, lsl #48
; CHECKED-ELF-NEXT: pacdb x16, x2
; CHECKED-DARWIN-NEXT: mov x17, x2
; CHECKED-DARWIN-NEXT: movk x17, #56789, lsl #48
; CHECKED-DARWIN-NEXT: pacdb x16, x17
; CHECKED-NEXT: Lresign_end_0:
; CHECKED-NEXT: mov x0, x16
; CHECKED-NEXT: ret
Expand All @@ -129,9 +132,11 @@ define i64 @test_resign_blend(i64 %arg, i64 %arg1, i64 %arg2) {
; TRAP-NEXT: b.eq [[L]]auth_success_1
; TRAP-NEXT: brk #0xc472
; TRAP-NEXT: Lauth_success_1:
; TRAP-NEXT: mov x17, x2
; TRAP-NEXT: movk x17, #56789, lsl #48
; TRAP-NEXT: pacdb x16, x17
; TRAP-ELF-NEXT: movk x2, #56789, lsl #48
; TRAP-ELF-NEXT: pacdb x16, x2
; TRAP-DARWIN-NEXT: mov x17, x2
; TRAP-DARWIN-NEXT: movk x17, #56789, lsl #48
; TRAP-DARWIN-NEXT: pacdb x16, x17
; TRAP-NEXT: mov x0, x16
; TRAP-NEXT: ret
%tmp0 = call i64 @llvm.ptrauth.blend(i64 %arg1, i64 12345)
Expand Down Expand Up @@ -299,6 +304,58 @@ define i64 @test_auth_too_large_discriminator(i64 %arg, i64 %arg1) {
ret i64 %tmp1
}

; As long as we support raw, non-blended 64-bit discriminators (which might be
; useful for low-level code such as dynamic loaders), the "auth" part of resign
; must not clobber %arg, if its upper bits are later used by the "sign" part.
define i64 @test_resign_aliased_discs_raw_sign_disc(i64 %p, i64 %arg) {
; UNCHECKED-LABEL: test_resign_aliased_discs_raw_sign_disc:
; UNCHECKED: %bb.0:
; UNCHECKED-NEXT: mov x16, x0
; UNCHECKED-NEXT: mov x17, x1
; UNCHECKED-NEXT: movk x17, #12345, lsl #48
; UNCHECKED-NEXT: autda x16, x17
; UNCHECKED-NEXT: pacdb x16, x1
; UNCHECKED-NEXT: mov x0, x16
; UNCHECKED-NEXT: ret
;
; CHECKED-LABEL: test_resign_aliased_discs_raw_sign_disc:
; CHECKED: %bb.0:
; CHECKED-NEXT: mov x16, x0
; CHECKED-NEXT: mov x17, x1
; CHECKED-NEXT: movk x17, #12345, lsl #48
; CHECKED-NEXT: autda x16, x17
; CHECKED-NEXT: mov x17, x16
; CHECKED-NEXT: xpacd x17
; CHECKED-NEXT: cmp x16, x17
; CHECKED-NEXT: b.eq [[L]]auth_success_3
; CHECKED-NEXT: mov x16, x17
; CHECKED-NEXT: b [[L]]resign_end_3
; CHECKED-NEXT: Lauth_success_3:
; CHECKED-NEXT: pacdb x16, x1
; CHECKED-NEXT: Lresign_end_3:
; CHECKED-NEXT: mov x0, x16
; CHECKED-NEXT: ret
;
; TRAP-LABEL: test_resign_aliased_discs_raw_sign_disc:
; TRAP: %bb.0:
; TRAP-NEXT: mov x16, x0
; TRAP-NEXT: mov x17, x1
; TRAP-NEXT: movk x17, #12345, lsl #48
; TRAP-NEXT: autda x16, x17
; TRAP-NEXT: mov x17, x16
; TRAP-NEXT: xpacd x17
; TRAP-NEXT: cmp x16, x17
; TRAP-NEXT: b.eq [[L]]auth_success_5
; TRAP-NEXT: brk #0xc472
; TRAP-NEXT: Lauth_success_5:
; TRAP-NEXT: pacdb x16, x1
; TRAP-NEXT: mov x0, x16
; TRAP-NEXT: ret
%auth.disc = call i64 @llvm.ptrauth.blend(i64 %arg, i64 12345)
%res = call i64 @llvm.ptrauth.resign(i64 %p, i32 2, i64 %auth.disc, i32 3, i64 %arg)
ret i64 %res
}

declare i64 @llvm.ptrauth.auth(i64, i32, i64)
declare i64 @llvm.ptrauth.resign(i64, i32, i64, i32, i64)
declare i64 @llvm.ptrauth.blend(i64, i64)
Loading