diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp index 7509afaeb5fef..91d97cb376303 100644 --- a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp +++ b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp @@ -12,6 +12,7 @@ #include "AArch64InstrInfo.h" #include "AArch64MachineFunctionInfo.h" #include "AArch64Subtarget.h" +#include "Utils/AArch64BaseInfo.h" #include "llvm/CodeGen/MachineBasicBlock.h" #include "llvm/CodeGen/MachineInstrBuilder.h" #include "llvm/CodeGen/MachineModuleInfo.h" @@ -35,7 +36,10 @@ class AArch64PointerAuth : public MachineFunctionPass { private: /// An immediate operand passed to BRK instruction, if it is ever emitted. - const unsigned BrkOperand = 0xc471; + static unsigned BrkOperandForKey(AArch64PACKey::ID KeyId) { + const unsigned BrkOperandBase = 0xc470; + return BrkOperandBase + KeyId; + } const AArch64Subtarget *Subtarget = nullptr; const AArch64InstrInfo *TII = nullptr; @@ -238,7 +242,7 @@ MachineBasicBlock &llvm::AArch64PAuth::checkAuthenticatedRegister( return MBB; case AuthCheckMethod::DummyLoad: BuildMI(MBB, MBBI, DL, TII->get(AArch64::LDRWui), getWRegFromXReg(TmpReg)) - .addReg(AArch64::LR) + .addReg(AuthenticatedReg) .addImm(0) .addMemOperand(createCheckMemOperand(MF, Subtarget)); return MBB; @@ -314,6 +318,10 @@ unsigned llvm::AArch64PAuth::getCheckerSizeInBytes(AuthCheckMethod Method) { bool AArch64PointerAuth::checkAuthenticatedLR( MachineBasicBlock::iterator TI) const { + const AArch64FunctionInfo *MFnI = TI->getMF()->getInfo(); + AArch64PACKey::ID KeyId = + MFnI->shouldSignWithBKey() ? AArch64PACKey::IB : AArch64PACKey::IA; + AuthCheckMethod Method = Subtarget->getAuthenticatedLRCheckMethod(); if (Method == AuthCheckMethod::None) @@ -354,7 +362,7 @@ bool AArch64PointerAuth::checkAuthenticatedLR( "More than a single register is used by TCRETURN"); checkAuthenticatedRegister(TI, Method, AArch64::LR, TmpReg, /*UseIKey=*/true, - BrkOperand); + BrkOperandForKey(KeyId)); return true; } diff --git a/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll b/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll index ec04e553cac6e..cf033cb8208cc 100644 --- a/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll +++ b/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll @@ -23,7 +23,7 @@ define i32 @tailcall_direct() "sign-return-address"="non-leaf" { ; ; COMMON-NEXT: b callee ; BRK-NEXT: .[[FAIL]]: -; BRK-NEXT: brk #0xc471 +; BRK-NEXT: brk #0xc470 tail call void asm sideeffect "", "~{lr}"() %call = tail call i32 @callee() ret i32 %call @@ -48,7 +48,7 @@ define i32 @tailcall_indirect(ptr %fptr) "sign-return-address"="non-leaf" { ; ; COMMON-NEXT: br x0 ; BRK-NEXT: .[[FAIL]]: -; BRK-NEXT: brk #0xc471 +; BRK-NEXT: brk #0xc470 tail call void asm sideeffect "", "~{lr}"() %call = tail call i32 %fptr() ret i32 %call @@ -89,7 +89,7 @@ define i32 @tailcall_direct_noframe_sign_all() "sign-return-address"="all" { ; ; COMMON-NEXT: b callee ; BRK-NEXT: .[[FAIL]]: -; BRK-NEXT: brk #0xc471 +; BRK-NEXT: brk #0xc470 %call = tail call i32 @callee() ret i32 %call } @@ -113,9 +113,20 @@ define i32 @tailcall_indirect_noframe_sign_all(ptr %fptr) "sign-return-address"= ; ; COMMON-NEXT: br x0 ; BRK-NEXT: .[[FAIL]]: -; BRK-NEXT: brk #0xc471 +; BRK-NEXT: brk #0xc470 %call = tail call i32 %fptr() ret i32 %call } +define i32 @tailcall_ib_key() "sign-return-address"="all" "sign-return-address-key"="b_key" { +; COMMON-LABEL: tailcall_ib_key: +; +; COMMON: b callee +; BRK-NEXT: .{{LBB.*}}: +; BRK-NEXT: brk #0xc471 + tail call void asm sideeffect "", "~{lr}"() + %call = tail call i32 @callee() + ret i32 %call +} + declare i32 @callee()