[llvm] [PAC][AArch64] Lower ptrauth constants in data (PR #94240)

Daniil Kovalev via llvm-commits llvm-commits at lists.llvm.org
Fri Jun 7 03:18:05 PDT 2024


https://github.com/kovdan01 updated https://github.com/llvm/llvm-project/pull/94240

>From eb7cc2fc12a8dd273fa78b0fc28958e5c6cee86d Mon Sep 17 00:00:00 2001
From: Daniil Kovalev <dkovalev at accesssoftek.com>
Date: Mon, 3 Jun 2024 08:41:00 +0300
Subject: [PATCH 1/2] [PAC][AArch64] Lower ptrauth constants in data

Lower global references to ptrauth constants into `@AUTH` `MCExpr`'s.
The logic is common for MachO and ELF - test both.

Co-authored-by: Ahmed Bougacha <ahmed at bougacha.org>
---
 llvm/include/llvm/CodeGen/AsmPrinter.h        |   6 +
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    |   3 +
 llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp |  48 +++++
 llvm/test/CodeGen/AArch64/ptrauth-reloc.ll    | 176 ++++++++++++++++++
 4 files changed, 233 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/ptrauth-reloc.ll

diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
index 81c3e4be95e9f..e918590e8193f 100644
--- a/llvm/include/llvm/CodeGen/AsmPrinter.h
+++ b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -38,6 +38,7 @@ class BasicBlock;
 class BlockAddress;
 class Constant;
 class ConstantArray;
+class ConstantPtrAuth;
 class DataLayout;
 class DIE;
 class DIEAbbrev;
@@ -585,6 +586,11 @@ class AsmPrinter : public MachineFunctionPass {
     emitGlobalConstant(DL, CV);
   }
 
+  /// Lower the specified ptrauth constant to an MCExpr.
+  virtual const MCExpr *lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
+    report_fatal_error("ptrauth constant lowering not implemented");
+  }
+
   /// Return true if the basic block has exactly one predecessor and the control
   /// transfer mechanism between the predecessor and this block is a
   /// fall-through.
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index c5755b9bdc8d0..a668b951e337e 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -3177,6 +3177,9 @@ const MCExpr *AsmPrinter::lowerConstant(const Constant *CV) {
   if (const ConstantInt *CI = dyn_cast<ConstantInt>(CV))
     return MCConstantExpr::create(CI->getZExtValue(), Ctx);
 
+  if (const ConstantPtrAuth *CPA = dyn_cast<ConstantPtrAuth>(CV))
+    return lowerConstantPtrAuth(*CPA);
+
   if (const GlobalValue *GV = dyn_cast<GlobalValue>(CV))
     return MCSymbolRefExpr::create(getSymbol(GV), Ctx);
 
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 7da540f8ef8e5..51f52bd2379eb 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -43,6 +43,8 @@
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInst.h"
@@ -90,6 +92,8 @@ class AArch64AsmPrinter : public AsmPrinter {
     return MCInstLowering.lowerOperand(MO, MCOp);
   }
 
+  const MCExpr *lowerConstantPtrAuth(const ConstantPtrAuth &CPA) override;
+
   void emitStartOfAsmFile(Module &M) override;
   void emitJumpTableInfo() override;
   std::tuple<const MCSymbol *, uint64_t, const MCSymbol *,
@@ -1575,6 +1579,50 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
   assert(STI->getInstrInfo()->getInstSizeInBytes(*MI) >= InstsEmitted * 4);
 }
 
+const MCExpr *
+AArch64AsmPrinter::lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
+  MCContext &Ctx = OutContext;
+
+  // Figure out the base symbol and the addend, if any.
+  APInt Offset(64, 0);
+  const Value *BaseGV = CPA.getPointer()->stripAndAccumulateConstantOffsets(
+      getDataLayout(), Offset, /*AllowNonInbounds=*/true);
+
+  auto *BaseGVB = dyn_cast<GlobalValue>(BaseGV);
+
+  // If we can't understand the referenced ConstantExpr, there's nothing
+  // else we can do: emit an error.
+  if (!BaseGVB) {
+    BaseGV->getContext().emitError(
+        "cannot resolve target base/addend of ptrauth constant");
+    return nullptr;
+  }
+
+  // If there is an addend, turn that into the appropriate MCExpr.
+  const MCExpr *Sym = MCSymbolRefExpr::create(getSymbol(BaseGVB), Ctx);
+  if (Offset.sgt(0))
+    Sym = MCBinaryExpr::createAdd(
+        Sym, MCConstantExpr::create(Offset.getSExtValue(), Ctx), Ctx);
+  else if (Offset.slt(0))
+    Sym = MCBinaryExpr::createSub(
+        Sym, MCConstantExpr::create((-Offset).getSExtValue(), Ctx), Ctx);
+
+  uint64_t KeyID = CPA.getKey()->getZExtValue();
+  // We later rely on valid KeyID value in AArch64PACKeyIDToString call from
+  // AArch64AuthMCExpr::printImpl, so fail fast.
+  if (KeyID > AArch64PACKey::LAST)
+    report_fatal_error("invalid AArch64 PAC Key ID '" + Twine(KeyID) + "'");
+
+  uint64_t Disc = CPA.getDiscriminator()->getZExtValue();
+  if (!isUInt<16>(Disc))
+    report_fatal_error("invalid AArch64 PAC Discriminator '" + Twine(Disc) +
+                       "'");
+
+  // Finally build the complete @AUTH expr.
+  return AArch64AuthMCExpr::create(Sym, Disc, AArch64PACKey::ID(KeyID),
+                                   CPA.hasAddressDiscriminator(), Ctx);
+}
+
 // Simple pseudo-instructions have their lowering (with expansion to real
 // instructions) auto-generated.
 #include "AArch64GenMCPseudoLowering.inc"
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll b/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll
new file mode 100644
index 0000000000000..8a8dc16f3dd9e
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll
@@ -0,0 +1,176 @@
+; RUN: rm -rf %t && split-file %s %t && cd %t
+
+;--- ok.ll
+
+; RUN: llc < ok.ll -mtriple arm64e-apple-darwin \
+; RUN:   | FileCheck %s --check-prefix=CHECK-MACHO
+; RUN: llc < ok.ll -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ELF
+
+; RUN: llc < ok.ll -mtriple arm64e-apple-darwin \
+; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-MACHO
+; RUN: llc < ok.ll -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ELF
+
+ at g = external global i32
+
+ at g_weak = extern_weak global i32
+
+ at g_strong_def = constant i32 42
+
+; CHECK-ELF-LABEL:     .globl g.ref.ia.0
+; CHECK-ELF-NEXT:      .p2align 4
+; CHECK-ELF-NEXT:    g.ref.ia.0:
+; CHECK-ELF-NEXT:      .xword 5
+; CHECK-ELF-NEXT:      .xword g at AUTH(ia,0)
+; CHECK-ELF-NEXT:      .xword 6
+
+; CHECK-MACHO-LABEL:   .section __DATA,__const
+; CHECK-MACHO-NEXT:    .globl _g.ref.ia.0
+; CHECK-MACHO-NEXT:    .p2align 4
+; CHECK-MACHO-NEXT:  _g.ref.ia.0:
+; CHECK-MACHO-NEXT:    .quad 5
+; CHECK-MACHO-NEXT:    .quad _g at AUTH(ia,0)
+; CHECK-MACHO-NEXT:    .quad 6
+
+ at g.ref.ia.0 = constant { i64, ptr, i64 } { i64 5, ptr ptrauth (ptr @g, i32 0), i64 6 }
+
+; CHECK-ELF-LABEL:     .globl g.ref.ia.42
+; CHECK-ELF-NEXT:      .p2align 3
+; CHECK-ELF-NEXT:    g.ref.ia.42:
+; CHECK-ELF-NEXT:      .xword g at AUTH(ia,42)
+
+; CHECK-MACHO-LABEL:   .globl _g.ref.ia.42
+; CHECK-MACHO-NEXT:    .p2align 3
+; CHECK-MACHO-NEXT:  _g.ref.ia.42:
+; CHECK-MACHO-NEXT:    .quad _g at AUTH(ia,42)
+
+ at g.ref.ia.42 = constant ptr ptrauth (ptr @g, i32 0, i64 42)
+
+; CHECK-ELF-LABEL:     .globl g.ref.ib.0
+; CHECK-ELF-NEXT:      .p2align 4
+; CHECK-ELF-NEXT:    g.ref.ib.0:
+; CHECK-ELF-NEXT:      .xword 5
+; CHECK-ELF-NEXT:      .xword g at AUTH(ib,0)
+; CHECK-ELF-NEXT:      .xword 6
+
+; CHECK-MACHO-LABEL:   .globl _g.ref.ib.0
+; CHECK-MACHO-NEXT:    .p2align 4
+; CHECK-MACHO-NEXT:  _g.ref.ib.0:
+; CHECK-MACHO-NEXT:    .quad 5
+; CHECK-MACHO-NEXT:    .quad _g at AUTH(ib,0)
+; CHECK-MACHO-NEXT:    .quad 6
+
+ at g.ref.ib.0 = constant { i64, ptr, i64 } { i64 5, ptr ptrauth (ptr @g, i32 1, i64 0), i64 6 }
+
+; CHECK-ELF-LABEL:     .globl g.ref.da.42.addr
+; CHECK-ELF-NEXT:      .p2align 3
+; CHECK-ELF-NEXT:    g.ref.da.42.addr:
+; CHECK-ELF-NEXT:      .xword g at AUTH(da,42,addr)
+
+; CHECK-MACHO-LABEL:   .globl _g.ref.da.42.addr
+; CHECK-MACHO-NEXT:    .p2align 3
+; CHECK-MACHO-NEXT:  _g.ref.da.42.addr:
+; CHECK-MACHO-NEXT:    .quad _g at AUTH(da,42,addr)
+
+ at g.ref.da.42.addr = constant ptr ptrauth (ptr @g, i32 2, i64 42, ptr @g.ref.da.42.addr)
+
+; CHECK-ELF-LABEL:     .globl g.offset.ref.da.0
+; CHECK-ELF-NEXT:      .p2align 3
+; CHECK-ELF-NEXT:    g.offset.ref.da.0:
+; CHECK-ELF-NEXT:      .xword (g+16)@AUTH(da,0)
+
+; CHECK-MACHO-LABEL:   .globl _g.offset.ref.da.0
+; CHECK-MACHO-NEXT:    .p2align 3
+; CHECK-MACHO-NEXT:  _g.offset.ref.da.0:
+; CHECK-MACHO-NEXT:    .quad (_g+16)@AUTH(da,0)
+
+ at g.offset.ref.da.0 = constant ptr ptrauth (i8* getelementptr (i8, ptr @g, i64 16), i32 2)
+
+; CHECK-ELF-LABEL:     .globl g.big_offset.ref.da.0
+; CHECK-ELF-NEXT:      .p2align 3
+; CHECK-ELF-NEXT:    g.big_offset.ref.da.0:
+; CHECK-ELF-NEXT:      .xword (g+2147549185)@AUTH(da,0)
+
+; CHECK-MACHO-LABEL:   .globl _g.big_offset.ref.da.0
+; CHECK-MACHO-NEXT:    .p2align 3
+; CHECK-MACHO-NEXT:  _g.big_offset.ref.da.0:
+; CHECK-MACHO-NEXT:    .quad (_g+2147549185)@AUTH(da,0)
+
+ at g.big_offset.ref.da.0 = constant ptr ptrauth (i8* getelementptr (i8, ptr @g, i64 add (i64 2147483648, i64 65537)), i32 2)
+
+; CHECK-ELF-LABEL:     .globl g.weird_ref.da.0
+; CHECK-ELF-NEXT:      .p2align 3
+; CHECK-ELF-NEXT:    g.weird_ref.da.0:
+; CHECK-ELF-NEXT:      .xword (g+16)@AUTH(da,0)
+
+; CHECK-MACHO-LABEL:   .globl _g.weird_ref.da.0
+; CHECK-MACHO-NEXT:    .p2align 3
+; CHECK-MACHO-NEXT:  _g.weird_ref.da.0:
+; CHECK-MACHO-NEXT:    .quad (_g+16)@AUTH(da,0)
+
+ at g.weird_ref.da.0 = constant i64 ptrtoint (ptr inttoptr (i64 ptrtoint (ptr ptrauth (i8* getelementptr (i8, ptr @g, i64 16), i32 2) to i64) to ptr) to i64)
+
+; CHECK-ELF-LABEL:     .globl g_weak.ref.ia.42
+; CHECK-ELF-NEXT:      .p2align 3
+; CHECK-ELF-NEXT:    g_weak.ref.ia.42:
+; CHECK-ELF-NEXT:      .xword g_weak at AUTH(ia,42)
+
+; CHECK-MACHO-LABEL:   .globl _g_weak.ref.ia.42
+; CHECK-MACHO-NEXT:    .p2align 3
+; CHECK-MACHO-NEXT:  _g_weak.ref.ia.42:
+; CHECK-MACHO-NEXT:    .quad _g_weak at AUTH(ia,42)
+
+ at g_weak.ref.ia.42 = constant ptr ptrauth (ptr @g_weak, i32 0, i64 42)
+
+; CHECK-ELF-LABEL:     .globl g_strong_def.ref.da.0
+; CHECK-ELF-NEXT:      .p2align 3
+; CHECK-ELF-NEXT:    g_strong_def.ref.da.0:
+; CHECK-ELF-NEXT:      .xword g_strong_def at AUTH(da,0)
+
+; CHECK-MACHO-LABEL:   .globl _g_strong_def.ref.da.0
+; CHECK-MACHO-NEXT:    .p2align 3
+; CHECK-MACHO-NEXT:  _g_strong_def.ref.da.0:
+; CHECK-MACHO-NEXT:    .quad _g_strong_def at AUTH(da,0)
+
+ at g_strong_def.ref.da.0 = constant ptr ptrauth (ptr @g_strong_def, i32 2)
+
+;--- err-key.ll
+
+; RUN: not --crash llc < err-key.ll -mtriple arm64e-apple-darwin 2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-KEY
+; RUN: not --crash llc < err-key.ll -mtriple aarch64-elf -mattr=+pauth 2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-KEY
+
+; RUN: not --crash llc < err-key.ll -mtriple arm64e-apple-darwin \
+; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1  2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-KEY
+; RUN: not --crash llc < err-key.ll -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1 2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-KEY
+
+; CHECK-ERR-KEY: LLVM ERROR: invalid AArch64 PAC Key ID '4'
+
+ at g = external global i32
+ at g.ref.4.0 = constant ptr ptrauth (ptr @g, i32 4, i64 0)
+
+;--- err-disc.ll
+
+; RUN: not --crash llc < err-disc.ll -mtriple arm64e-apple-darwin 2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-DISC
+; RUN: not --crash llc < err-disc.ll -mtriple aarch64-elf -mattr=+pauth 2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-DISC
+
+; RUN: not --crash llc < err-disc.ll -mtriple arm64e-apple-darwin \
+; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1  2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-DISC
+; RUN: not --crash llc < err-disc.ll -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1 2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-DISC
+
+; CHECK-ERR-DISC: LLVM ERROR: invalid AArch64 PAC Discriminator '65536'
+
+ at g = external global i32
+ at g.ref.ia.65536 = constant ptr ptrauth (ptr @g, i32 0, i64 65536)

>From 9633e04f21f0745dc76f905684168333c2c4ff4d Mon Sep 17 00:00:00 2001
From: Daniil Kovalev <dkovalev at accesssoftek.com>
Date: Fri, 7 Jun 2024 13:11:27 +0300
Subject: [PATCH 2/2] Address review comments

---
 llvm/include/llvm/CodeGen/AsmPrinter.h        |  1 -
 llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp | 10 +++++-----
 llvm/test/CodeGen/AArch64/ptrauth-reloc.ll    |  4 ++--
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
index e918590e8193f..011f8c6534b6a 100644
--- a/llvm/include/llvm/CodeGen/AsmPrinter.h
+++ b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -586,7 +586,6 @@ class AsmPrinter : public MachineFunctionPass {
     emitGlobalConstant(DL, CV);
   }
 
-  /// Lower the specified ptrauth constant to an MCExpr.
   virtual const MCExpr *lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
     report_fatal_error("ptrauth constant lowering not implemented");
   }
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 51f52bd2379eb..da11539eab348 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -43,8 +43,6 @@
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfoMetadata.h"
-#include "llvm/IR/DiagnosticInfo.h"
-#include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInst.h"
@@ -1611,12 +1609,14 @@ AArch64AsmPrinter::lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
   // We later rely on valid KeyID value in AArch64PACKeyIDToString call from
   // AArch64AuthMCExpr::printImpl, so fail fast.
   if (KeyID > AArch64PACKey::LAST)
-    report_fatal_error("invalid AArch64 PAC Key ID '" + Twine(KeyID) + "'");
+    report_fatal_error("AArch64 PAC Key ID '" + Twine(KeyID) +
+                       "' out of range [0, " +
+                       Twine((unsigned)AArch64PACKey::LAST) + "]");
 
   uint64_t Disc = CPA.getDiscriminator()->getZExtValue();
   if (!isUInt<16>(Disc))
-    report_fatal_error("invalid AArch64 PAC Discriminator '" + Twine(Disc) +
-                       "'");
+    report_fatal_error("AArch64 PAC Discriminator '" + Twine(Disc) +
+                       "' out of range [0, 0xFFFF]");
 
   // Finally build the complete @AUTH expr.
   return AArch64AuthMCExpr::create(Sym, Disc, AArch64PACKey::ID(KeyID),
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll b/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll
index 8a8dc16f3dd9e..b7304b957a001 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll
+++ b/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll
@@ -151,7 +151,7 @@
 ; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1 2>&1 \
 ; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-KEY
 
-; CHECK-ERR-KEY: LLVM ERROR: invalid AArch64 PAC Key ID '4'
+; CHECK-ERR-KEY: LLVM ERROR: AArch64 PAC Key ID '4' out of range [0, 3]
 
 @g = external global i32
 @g.ref.4.0 = constant ptr ptrauth (ptr @g, i32 4, i64 0)
@@ -170,7 +170,7 @@
 ; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1 2>&1 \
 ; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-DISC
 
-; CHECK-ERR-DISC: LLVM ERROR: invalid AArch64 PAC Discriminator '65536'
+; CHECK-ERR-DISC: LLVM ERROR: AArch64 PAC Discriminator '65536' out of range [0, 0xFFFF]
 
 @g = external global i32
 @g.ref.ia.65536 = constant ptr ptrauth (ptr @g, i32 0, i64 65536)



More information about the llvm-commits mailing list