[AArch64][PAC] Refine authenticated pointer check methods

[atrosinenko]

Align the values of the immediate operand of BRK instruction with those used by the existing arm64e implementation.

Make AuthCheckMethod::DummyLoad use the requested register instead of LR.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

Changes

Align the values of the immediate operand of BRK instruction with those used by the existing arm64e implementation.

Make AuthCheckMethod::DummyLoad use the requested register instead of LR.

---

Full diff: https://github.com/llvm/llvm-project/pull/74074.diff

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64PointerAuth.cpp (+11-3)
• (modified) llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll (+15-4)

diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 7576d2a899d1afb..8597700d5e6f080 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;
@@ -174,7 +178,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;
@@ -250,6 +254,10 @@ unsigned llvm::AArch64PAuth::getCheckerSizeInBytes(AuthCheckMethod Method) {
 
 bool AArch64PointerAuth::checkAuthenticatedLR(
     MachineBasicBlock::iterator TI) const {
+  const AArch64FunctionInfo *MFnI = TI->getMF()->getInfo<AArch64FunctionInfo>();
+  AArch64PACKey::ID KeyId =
+      MFnI->shouldSignWithBKey() ? AArch64PACKey::IB : AArch64PACKey::IA;
+
   AuthCheckMethod Method = Subtarget->getAuthenticatedLRCheckMethod();
 
   if (Method == AuthCheckMethod::None)
@@ -290,7 +298,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 ec04e553cac6e37..cf033cb8208cc90 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()

[atrosinenko]

This PR was factored out from #72502.

[atrosinenko]

Ping.

[atrosinenko]

Ping.

A few notes on the patch:

• instead of defining a private constant BrkOperandBase = 0xc470 to add to an LLVM's key id here and there (as it was done in [AArch64][PAC] Implement code generation for @llvm.ptrauth.auth #72502), an opaque BrkOperandForKey(key) helper function is introduced as it is kind of coincidence that the "reference" BRK operands are computed based on the same AArch64PACKey enum :)
• it is worth to further change AuthCheckMethod::DummyLoad to perform a 1-byte load instead of a 4-byte load if checking a "data pointer" (loading the last byte of the last mapped page in a range is not relevant for code but possible for data). This is not changed by this patch as @llvm.ptrauth.auth should be the first suitable user of checkAuthenticatedRegister function for data pointers, so no tests can be added here. This patch does not include adding the generic AuthCheckMethod::XPAC method for a similar reason

[atrosinenko]

Ping.

[atrosinenko]

Ping.

[atrosinenko]

Ping.

[kbeyls]

These changes make sense to me.
I'm assuming that aligning the values of the brk operand with what arm64e does helps reduce possible incompatibilities with any users/tools that might interpret these brk operand values?
Or is there another rationale to align the brk operand values with what arm64e does?
I didn't try to check independently if these are indeed the values arm64e uses.

[DavidSpickett]

The support in lldb was added for arm64e and I see the value in there:

llvm-project/lldb/source/Plugins/Disassembler/LLVMC/DisassemblerLLVMC.cpp

Line 1426 in f956e7f

bool DisassemblerLLVMC::MCDisasmInstance::IsAuthenticated(

https://github.com/llvm/llvm-project/blob/main/lldb/test/API/functionalities/ptrauth_diagnostics/brkC47x_code/brkC47x.c

[DavidSpickett]

Currently only StopInfoMachException.cpp (aka Darwin exception) uses that method but if the values match up it wouldn't be difficult to get it working for others.

[atrosinenko]

Thank you for the review!

I'm assuming that aligning the values of the brk operand with what arm64e does helps reduce possible incompatibilities with any users/tools that might interpret these brk operand values?

Yes, the reason was merely to not diverge unless needed.

[DavidSpickett]

Definitely agree, it's ok for Apple who only use lldb but once we get this on Linux there's the question of GDB as well so very good to stick with one set of values.

Are these values something that will be included in the ABI, or should be? Strange thing to be in an ABI I guess, but if we're agreeing on values sounds like it could be.

[atrosinenko]

@DavidSpickett Did I understand your comment right: it is OK to merge this PR as-is AND it may be worth clarifying ABI document afterwards?

[DavidSpickett]

Yes, that's what I meant.

Even if it's "consistent numbers were chosen for consistency's sake only", it's good to record that.

[atrosinenko]

@DavidSpickett Thank you for the clarification!