[clang] [compiler-rt] [llvm] [PAC][AArch64] Support init/fini array signing (PR #96478)

Daniil Kovalev via cfe-commits cfe-commits at lists.llvm.org
Fri Jun 28 06:03:09 PDT 2024


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

>From 5399237a71c0ccd872821034d83ea2c3a04bed3f Mon Sep 17 00:00:00 2001
From: Daniil Kovalev <dkovalev at accesssoftek.com>
Date: Fri, 21 Jun 2024 12:32:51 +0300
Subject: [PATCH 1/2] [PAC][AArch64] Support init/fini array signing

If both `-fptrauth-init-fini` and `-fptrauth-calls` are passed, sign
function pointers in `llvm.global_ctors` and `llvm.global_dtors` with
constant discriminator 0xD9D4 (`ptrauth_string_discriminator("init_fini")`).
Additionally, if `-fptrauth-init-fini-address-discrimination` is passed,
address discrimination is used for signing (otherwise, just constant
discriminator is used).
---
 clang/include/clang/Basic/Features.def        |  1 +
 clang/include/clang/Basic/LangOptions.def     |  1 +
 .../include/clang/Basic/PointerAuthOptions.h  |  7 ++
 clang/include/clang/Driver/Options.td         |  1 +
 clang/lib/CodeGen/CodeGenModule.cpp           | 63 +++++++++------
 clang/lib/Driver/ToolChains/Clang.cpp         |  3 +
 clang/lib/Frontend/CompilerInvocation.cpp     |  9 +++
 clang/lib/Headers/ptrauth.h                   |  8 ++
 clang/test/CodeGen/aarch64-elf-pauthabi.c     | 12 ++-
 clang/test/CodeGen/ptrauth-init-fini.c        | 39 ++++++++++
 clang/test/Driver/aarch64-ptrauth.c           |  6 +-
 clang/test/Preprocessor/ptrauth_feature.c     | 52 +++++++++----
 compiler-rt/lib/builtins/crtbegin.c           | 16 ++++
 llvm/include/llvm/BinaryFormat/ELF.h          |  3 +-
 .../AArch64/note-gnu-property-elf-pauthabi.ll |  2 +-
 .../test/CodeGen/AArch64/ptrauth-init-fini.ll | 77 +++++++++++++++++++
 .../ELF/AArch64/aarch64-feature-pauth.s       | 18 ++---
 llvm/tools/llvm-readobj/ELFDumper.cpp         |  4 +-
 18 files changed, 268 insertions(+), 54 deletions(-)
 create mode 100644 clang/test/CodeGen/ptrauth-init-fini.c
 create mode 100644 llvm/test/CodeGen/AArch64/ptrauth-init-fini.ll

diff --git a/clang/include/clang/Basic/Features.def b/clang/include/clang/Basic/Features.def
index 53f410d3cb4bd..5dca40b261655 100644
--- a/clang/include/clang/Basic/Features.def
+++ b/clang/include/clang/Basic/Features.def
@@ -110,6 +110,7 @@ FEATURE(ptrauth_vtable_pointer_address_discrimination, LangOpts.PointerAuthVTPtr
 FEATURE(ptrauth_vtable_pointer_type_discrimination, LangOpts.PointerAuthVTPtrTypeDiscrimination)
 FEATURE(ptrauth_member_function_pointer_type_discrimination, LangOpts.PointerAuthCalls)
 FEATURE(ptrauth_init_fini, LangOpts.PointerAuthInitFini)
+FEATURE(ptrauth_init_fini_address_discrimination, LangOpts.PointerAuthInitFiniAddressDiscrimination)
 EXTENSION(swiftcc,
   PP.getTargetInfo().checkCallingConvention(CC_Swift) ==
   clang::TargetInfo::CCCR_OK)
diff --git a/clang/include/clang/Basic/LangOptions.def b/clang/include/clang/Basic/LangOptions.def
index 6dd6b5614f44c..2de8547888831 100644
--- a/clang/include/clang/Basic/LangOptions.def
+++ b/clang/include/clang/Basic/LangOptions.def
@@ -168,6 +168,7 @@ LANGOPT(PointerAuthAuthTraps, 1, 0, "pointer authentication failure traps")
 LANGOPT(PointerAuthVTPtrAddressDiscrimination, 1, 0, "incorporate address discrimination in authenticated vtable pointers")
 LANGOPT(PointerAuthVTPtrTypeDiscrimination, 1, 0, "incorporate type discrimination in authenticated vtable pointers")
 LANGOPT(PointerAuthInitFini, 1, 0, "sign function pointers in init/fini arrays")
+LANGOPT(PointerAuthInitFiniAddressDiscrimination, 1, 0, "incorporate address discrimination in authenticated function pointers in init/fini arrays")
 
 LANGOPT(DoubleSquareBracketAttributes, 1, 0, "'[[]]' attributes extension for all language standard modes")
 LANGOPT(ExperimentalLateParseAttributes, 1, 0, "experimental late parsing of attributes")
diff --git a/clang/include/clang/Basic/PointerAuthOptions.h b/clang/include/clang/Basic/PointerAuthOptions.h
index aaad4a2b2b5ae..9e2b64111e461 100644
--- a/clang/include/clang/Basic/PointerAuthOptions.h
+++ b/clang/include/clang/Basic/PointerAuthOptions.h
@@ -23,6 +23,10 @@
 
 namespace clang {
 
+/// Constant discriminator to be used with function pointers in .init_array and
+/// .fini_array. The value is ptrauth_string_discriminator("init_fini")
+constexpr uint16_t InitFiniPointerConstantDiscriminator = 0xD9D4;
+
 constexpr unsigned PointerAuthKeyNone = -1;
 
 class PointerAuthSchema {
@@ -150,6 +154,9 @@ class PointerAuthSchema {
 struct PointerAuthOptions {
   /// The ABI for C function pointers.
   PointerAuthSchema FunctionPointers;
+
+  /// The ABI for function addresses in .init_array and .fini_array
+  PointerAuthSchema InitFiniPointers;
 };
 
 } // end namespace clang
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
index c529cc9506667..5b8d4139d975b 100644
--- a/clang/include/clang/Driver/Options.td
+++ b/clang/include/clang/Driver/Options.td
@@ -4228,6 +4228,7 @@ defm ptrauth_vtable_pointer_address_discrimination :
 defm ptrauth_vtable_pointer_type_discrimination :
   OptInCC1FFlag<"ptrauth-vtable-pointer-type-discrimination", "Enable type discrimination of vtable pointers">;
 defm ptrauth_init_fini : OptInCC1FFlag<"ptrauth-init-fini", "Enable signing of function pointers in init/fini arrays">;
+defm ptrauth_init_fini_address_discrimination : OptInCC1FFlag<"ptrauth-init-fini-address-discrimination", "Enable address discrimination of function pointers in init/fini arrays">;
 }
 
 def fenable_matrix : Flag<["-"], "fenable-matrix">, Group<f_Group>,
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 76534475e88f7..53854c3c4b85b 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -1210,10 +1210,13 @@ void CodeGenModule::Release() {
           (LangOpts.PointerAuthVTPtrTypeDiscrimination
            << AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_VPTRTYPEDISCR) |
           (LangOpts.PointerAuthInitFini
-           << AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINI);
-      static_assert(AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINI ==
-                        AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_LAST,
-                    "Update when new enum items are defined");
+           << AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINI) |
+          (LangOpts.PointerAuthInitFiniAddressDiscrimination
+           << AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINIADDRDISC);
+      static_assert(
+          AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINIADDRDISC ==
+              AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_LAST,
+          "Update when new enum items are defined");
       if (PAuthABIVersion != 0) {
         getModule().addModuleFlag(llvm::Module::Error,
                                   "aarch64-elf-pauthabi-platform",
@@ -2051,37 +2054,49 @@ void CodeGenModule::AddGlobalDtor(llvm::Function *Dtor, int Priority,
 void CodeGenModule::EmitCtorList(CtorList &Fns, const char *GlobalName) {
   if (Fns.empty()) return;
 
-  // Ctor function type is void()*.
-  llvm::FunctionType* CtorFTy = llvm::FunctionType::get(VoidTy, false);
-  llvm::Type *CtorPFTy = llvm::PointerType::get(CtorFTy,
-      TheModule.getDataLayout().getProgramAddressSpace());
+  const PointerAuthSchema &InitFiniAuthSchema =
+      getCodeGenOpts().PointerAuth.InitFiniPointers;
 
-  // Get the type of a ctor entry, { i32, void ()*, i8* }.
-  llvm::StructType *CtorStructTy = llvm::StructType::get(
-      Int32Ty, CtorPFTy, VoidPtrTy);
+  // Ctor function type is ptr.
+  llvm::PointerType *PtrTy = llvm::PointerType::get(
+      getLLVMContext(), TheModule.getDataLayout().getProgramAddressSpace());
+
+  // Get the type of a ctor entry, { i32, ptr, ptr }.
+  llvm::StructType *CtorStructTy = llvm::StructType::get(Int32Ty, PtrTy, PtrTy);
 
   // Construct the constructor and destructor arrays.
-  ConstantInitBuilder builder(*this);
-  auto ctors = builder.beginArray(CtorStructTy);
+  ConstantInitBuilder Builder(*this);
+  auto Ctors = Builder.beginArray(CtorStructTy);
   for (const auto &I : Fns) {
-    auto ctor = ctors.beginStruct(CtorStructTy);
-    ctor.addInt(Int32Ty, I.Priority);
-    ctor.add(I.Initializer);
+    auto Ctor = Ctors.beginStruct(CtorStructTy);
+    Ctor.addInt(Int32Ty, I.Priority);
+    if (InitFiniAuthSchema) {
+      llvm::Constant *StorageAddress =
+          (InitFiniAuthSchema.isAddressDiscriminated()
+               ? StorageAddress = Ctor.getAddrOfCurrentPosition(PtrTy)
+               : nullptr);
+      llvm::Constant *SignedCtorPtr = getConstantSignedPointer(
+          I.Initializer, InitFiniAuthSchema.getKey(), StorageAddress,
+          llvm::ConstantInt::get(
+              SizeTy, InitFiniAuthSchema.getConstantDiscrimination()));
+      Ctor.add(SignedCtorPtr);
+    } else {
+      Ctor.add(I.Initializer);
+    }
     if (I.AssociatedData)
-      ctor.add(I.AssociatedData);
+      Ctor.add(I.AssociatedData);
     else
-      ctor.addNullPointer(VoidPtrTy);
-    ctor.finishAndAddTo(ctors);
+      Ctor.addNullPointer(PtrTy);
+    Ctor.finishAndAddTo(Ctors);
   }
 
-  auto list =
-    ctors.finishAndCreateGlobal(GlobalName, getPointerAlign(),
-                                /*constant*/ false,
-                                llvm::GlobalValue::AppendingLinkage);
+  auto List = Ctors.finishAndCreateGlobal(GlobalName, getPointerAlign(),
+                                          /*constant*/ false,
+                                          llvm::GlobalValue::AppendingLinkage);
 
   // The LTO linker doesn't seem to like it when we set an alignment
   // on appending variables.  Take it off as a workaround.
-  list->setAlignment(std::nullopt);
+  List->setAlignment(std::nullopt);
 
   Fns.clear();
 }
diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
index 2ce9e2f4bcfcd..e281f44b16641 100644
--- a/clang/lib/Driver/ToolChains/Clang.cpp
+++ b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -1789,6 +1789,9 @@ void Clang::AddAArch64TargetArgs(const ArgList &Args,
       options::OPT_fno_ptrauth_vtable_pointer_type_discrimination);
   Args.addOptInFlag(CmdArgs, options::OPT_fptrauth_init_fini,
                     options::OPT_fno_ptrauth_init_fini);
+  Args.addOptInFlag(CmdArgs,
+                    options::OPT_fptrauth_init_fini_address_discrimination,
+                    options::OPT_fno_ptrauth_init_fini_address_discrimination);
 }
 
 void Clang::AddLoongArchTargetArgs(const ArgList &Args,
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index a6d9f42ace9cc..9d01edfd253d5 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -1468,6 +1468,11 @@ void CompilerInvocation::setDefaultPointerAuthOptions(
     // If you change anything here, be sure to update <ptrauth.h>.
     Opts.FunctionPointers =
         PointerAuthSchema(Key::ASIA, false, Discrimination::None);
+    if (LangOpts.PointerAuthInitFini) {
+      Opts.InitFiniPointers = PointerAuthSchema(
+          Key::ASIA, LangOpts.PointerAuthInitFiniAddressDiscrimination,
+          Discrimination::Constant, InitFiniPointerConstantDiscriminator);
+    }
   }
 }
 
@@ -3387,6 +3392,8 @@ static void GeneratePointerAuthArgs(const LangOptions &Opts,
     GenerateArg(Consumer, OPT_fptrauth_vtable_pointer_type_discrimination);
   if (Opts.PointerAuthInitFini)
     GenerateArg(Consumer, OPT_fptrauth_init_fini);
+  if (Opts.PointerAuthInitFiniAddressDiscrimination)
+    GenerateArg(Consumer, OPT_fptrauth_init_fini_address_discrimination);
 }
 
 static void ParsePointerAuthArgs(LangOptions &Opts, ArgList &Args,
@@ -3400,6 +3407,8 @@ static void ParsePointerAuthArgs(LangOptions &Opts, ArgList &Args,
   Opts.PointerAuthVTPtrTypeDiscrimination =
       Args.hasArg(OPT_fptrauth_vtable_pointer_type_discrimination);
   Opts.PointerAuthInitFini = Args.hasArg(OPT_fptrauth_init_fini);
+  Opts.PointerAuthInitFiniAddressDiscrimination =
+      Args.hasArg(OPT_fptrauth_init_fini_address_discrimination);
 }
 
 /// Check if input file kind and language standard are compatible.
diff --git a/clang/lib/Headers/ptrauth.h b/clang/lib/Headers/ptrauth.h
index 1a4bd02933ea2..25d89b38adb44 100644
--- a/clang/lib/Headers/ptrauth.h
+++ b/clang/lib/Headers/ptrauth.h
@@ -32,6 +32,9 @@ typedef enum {
      The extra data is always 0. */
   ptrauth_key_function_pointer = ptrauth_key_process_independent_code,
 
+  /* The key used to sign pointers in ELF .init_array/.fini_array. */
+  ptrauth_key_init_fini_pointer = ptrauth_key_asia,
+
   /* Other pointers signed under the ABI use private ABI rules. */
 
 } ptrauth_key;
@@ -205,6 +208,11 @@ typedef __UINTPTR_TYPE__ ptrauth_generic_signature_t;
 #define ptrauth_sign_generic_data(__value, __data)                             \
   __builtin_ptrauth_sign_generic_data(__value, __data)
 
+#define __ptrauth_init_fini_discriminator 0xd9d4
+#define __ptrauth_init_fini_pointer(address_discrimination)                    \
+  __ptrauth(ptrauth_key_init_fini_pointer, address_discrimination,             \
+            __ptrauth_init_fini_discriminator)
+
 #else
 
 #define ptrauth_strip(__value, __key)                                          \
diff --git a/clang/test/CodeGen/aarch64-elf-pauthabi.c b/clang/test/CodeGen/aarch64-elf-pauthabi.c
index aa83ee3e0d7b0..d60dcee361888 100644
--- a/clang/test/CodeGen/aarch64-elf-pauthabi.c
+++ b/clang/test/CodeGen/aarch64-elf-pauthabi.c
@@ -5,7 +5,8 @@
 // RUN:   -fptrauth-auth-traps \
 // RUN:   -fptrauth-vtable-pointer-address-discrimination \
 // RUN:   -fptrauth-vtable-pointer-type-discrimination \
-// RUN:   -fptrauth-init-fini %s | \
+// RUN:   -fptrauth-init-fini %s \
+// RUN:   -fptrauth-init-fini-address-discrimination %s | \
 // RUN:   FileCheck %s --check-prefix=ALL
 
 // RUN: %clang_cc1 -triple aarch64-linux -emit-llvm -o - \
@@ -32,8 +33,12 @@
 // RUN:   -fptrauth-calls -fptrauth-init-fini %s | \
 // RUN:   FileCheck %s --check-prefix=INITFINI
 
+// RUN: %clang_cc1 -triple aarch64-linux -emit-llvm -o - \
+// RUN:   -fptrauth-calls -fptrauth-init-fini -fptrauth-init-fini-address-discrimination %s | \
+// RUN:   FileCheck %s --check-prefix=INITFINIADDR
+
 // ALL: !{i32 1, !"aarch64-elf-pauthabi-platform", i32 268435458}
-// ALL: !{i32 1, !"aarch64-elf-pauthabi-version", i32 127}
+// ALL: !{i32 1, !"aarch64-elf-pauthabi-version", i32 255}
 
 // INTRIN: !{i32 1, !"aarch64-elf-pauthabi-platform", i32 268435458}
 // INTRIN: !{i32 1, !"aarch64-elf-pauthabi-version", i32 1}
@@ -56,4 +61,7 @@
 // INITFINI: !{i32 1, !"aarch64-elf-pauthabi-platform", i32 268435458}
 // INITFINI: !{i32 1, !"aarch64-elf-pauthabi-version", i32 66}
 
+// INITFINIADDR: !{i32 1, !"aarch64-elf-pauthabi-platform", i32 268435458}
+// INITFINIADDR: !{i32 1, !"aarch64-elf-pauthabi-version", i32 194}
+
 void foo() {}
diff --git a/clang/test/CodeGen/ptrauth-init-fini.c b/clang/test/CodeGen/ptrauth-init-fini.c
new file mode 100644
index 0000000000000..8b1abd0681c90
--- /dev/null
+++ b/clang/test/CodeGen/ptrauth-init-fini.c
@@ -0,0 +1,39 @@
+// REQUIRES: aarch64-registered-target
+
+// RUN: %clang -target aarch64-elf -march=armv8.3-a+pauth -fptrauth-calls -fptrauth-init-fini    \
+// RUN:   -S -emit-llvm %s -o - | FileCheck --check-prefix=SIGNED %s
+
+// RUN: %clang -target aarch64-elf -march=armv8.3-a+pauth -fptrauth-calls -fptrauth-init-fini    \
+// RUN:   -fptrauth-init-fini-address-discrimination -S -emit-llvm %s -o - | FileCheck --check-prefix=ADDRDISC %s
+
+// RUN: %clang -target aarch64-elf -march=armv8.3-a+pauth -fptrauth-calls -fno-ptrauth-init-fini \
+// RUN:   -S -emit-llvm %s -o - | FileCheck --check-prefix=UNSIGNED %s
+
+// RUN: %clang -target aarch64-elf -march=armv8.3-a+pauth -fptrauth-calls -fptrauth-init-fini-address-discrimination \
+// RUN:   -S -emit-llvm %s -o - | FileCheck --check-prefix=UNSIGNED %s
+
+// RUN: %clang -target aarch64-elf -march=armv8.3-a+pauth                 -fptrauth-init-fini    \
+// RUN:   -S -emit-llvm %s -o - | FileCheck --check-prefix=UNSIGNED %s
+
+// SIGNED: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764), ptr null }]
+// SIGNED: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764), ptr null }]
+
+// ADDRDISC: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764, ptr getelementptr inbounds ([1 x { i32, ptr, ptr }], ptr @llvm.global_ctors, i32 0, i32 0, i32 1)), ptr null }]
+// ADDRDISC: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764, ptr getelementptr inbounds ([1 x { i32, ptr, ptr }], ptr @llvm.global_dtors, i32 0, i32 0, i32 1)), ptr null }]
+
+// UNSIGNED: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @foo, ptr null }]
+// UNSIGNED: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @bar, ptr null }]
+
+volatile int x = 0;
+
+__attribute__((constructor)) void foo(void) {
+  x = 42;
+}
+
+__attribute__((destructor)) void bar(void) {
+  x = 24;
+}
+
+int main() {
+  return x;
+}
diff --git a/clang/test/Driver/aarch64-ptrauth.c b/clang/test/Driver/aarch64-ptrauth.c
index fa0125f4b22a9..aac7bc32c562e 100644
--- a/clang/test/Driver/aarch64-ptrauth.c
+++ b/clang/test/Driver/aarch64-ptrauth.c
@@ -10,12 +10,13 @@
 // RUN:   -fno-ptrauth-vtable-pointer-address-discrimination -fptrauth-vtable-pointer-address-discrimination \
 // RUN:   -fno-ptrauth-vtable-pointer-type-discrimination -fptrauth-vtable-pointer-type-discrimination \
 // RUN:   -fno-ptrauth-init-fini -fptrauth-init-fini \
+// RUN:   -fno-ptrauth-init-fini-address-discrimination -fptrauth-init-fini-address-discrimination \
 // RUN:   %s 2>&1 | FileCheck %s --check-prefix=ALL
-// ALL: "-cc1"{{.*}} "-fptrauth-intrinsics" "-fptrauth-calls" "-fptrauth-returns" "-fptrauth-auth-traps" "-fptrauth-vtable-pointer-address-discrimination" "-fptrauth-vtable-pointer-type-discrimination" "-fptrauth-init-fini"
+// ALL: "-cc1"{{.*}} "-fptrauth-intrinsics" "-fptrauth-calls" "-fptrauth-returns" "-fptrauth-auth-traps" "-fptrauth-vtable-pointer-address-discrimination" "-fptrauth-vtable-pointer-type-discrimination" "-fptrauth-init-fini" "-fptrauth-init-fini-address-discrimination"
 
 // RUN: not %clang -### -c --target=x86_64 -fptrauth-intrinsics -fptrauth-calls -fptrauth-returns -fptrauth-auth-traps \
 // RUN:   -fptrauth-vtable-pointer-address-discrimination -fptrauth-vtable-pointer-type-discrimination \
-// RUN:   -fptrauth-init-fini %s 2>&1 | FileCheck %s --check-prefix=ERR
+// RUN:   -fptrauth-init-fini -fptrauth-init-fini-address-discrimination %s 2>&1 | FileCheck %s --check-prefix=ERR
 // ERR:      error: unsupported option '-fptrauth-intrinsics' for target '{{.*}}'
 // ERR-NEXT: error: unsupported option '-fptrauth-calls' for target '{{.*}}'
 // ERR-NEXT: error: unsupported option '-fptrauth-returns' for target '{{.*}}'
@@ -23,3 +24,4 @@
 // ERR-NEXT: error: unsupported option '-fptrauth-vtable-pointer-address-discrimination' for target '{{.*}}'
 // ERR-NEXT: error: unsupported option '-fptrauth-vtable-pointer-type-discrimination' for target '{{.*}}'
 // ERR-NEXT: error: unsupported option '-fptrauth-init-fini' for target '{{.*}}'
+// ERR-NEXT: error: unsupported option '-fptrauth-init-fini-address-discrimination' for target '{{.*}}'
diff --git a/clang/test/Preprocessor/ptrauth_feature.c b/clang/test/Preprocessor/ptrauth_feature.c
index 80e239110ffc7..f8a98f76d925b 100644
--- a/clang/test/Preprocessor/ptrauth_feature.c
+++ b/clang/test/Preprocessor/ptrauth_feature.c
@@ -4,56 +4,72 @@
 // RUN:   -fptrauth-returns \
 // RUN:   -fptrauth-vtable-pointer-address-discrimination \
 // RUN:   -fptrauth-vtable-pointer-type-discrimination \
-// RUN:   -fptrauth-init-fini | \
-// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI
+// RUN:   -fptrauth-init-fini \
+// RUN:   -fptrauth-init-fini-address-discrimination | \
+// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI,INITFINI_ADDR_DISCR
 
 // RUN: %clang_cc1 -E %s -triple=aarch64 \
 // RUN:   -fptrauth-calls \
 // RUN:   -fptrauth-returns \
 // RUN:   -fptrauth-vtable-pointer-address-discrimination \
 // RUN:   -fptrauth-vtable-pointer-type-discrimination \
-// RUN:   -fptrauth-init-fini | \
-// RUN:   FileCheck %s --check-prefixes=NOINTRIN,CALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI
+// RUN:   -fptrauth-init-fini \
+// RUN:   -fptrauth-init-fini-address-discrimination | \
+// RUN:   FileCheck %s --check-prefixes=NOINTRIN,CALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI,INITFINI_ADDR_DISCR
 
 // RUN: %clang_cc1 -E %s -triple=aarch64 \
 // RUN:   -fptrauth-intrinsics \
 // RUN:   -fptrauth-returns \
 // RUN:   -fptrauth-vtable-pointer-address-discrimination \
 // RUN:   -fptrauth-vtable-pointer-type-discrimination \
-// RUN:   -fptrauth-init-fini | \
-// RUN:   FileCheck %s --check-prefixes=INTRIN,NOCALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI
+// RUN:   -fptrauth-init-fini \
+// RUN:   -fptrauth-init-fini-address-discrimination | \
+// RUN:   FileCheck %s --check-prefixes=INTRIN,NOCALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI,INITFINI_ADDR_DISCR
 
 // RUN: %clang_cc1 -E %s -triple=aarch64 \
 // RUN:   -fptrauth-intrinsics \
 // RUN:   -fptrauth-calls \
 // RUN:   -fptrauth-vtable-pointer-address-discrimination \
 // RUN:   -fptrauth-vtable-pointer-type-discrimination \
-// RUN:   -fptrauth-init-fini | \
-// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,NORETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI
+// RUN:   -fptrauth-init-fini \
+// RUN:   -fptrauth-init-fini-address-discrimination | \
+// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,NORETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI,INITFINI_ADDR_DISCR
 
 // RUN: %clang_cc1 -E %s -triple=aarch64 \
 // RUN:   -fptrauth-intrinsics \
 // RUN:   -fptrauth-calls \
 // RUN:   -fptrauth-returns \
 // RUN:   -fptrauth-vtable-pointer-type-discrimination \
-// RUN:   -fptrauth-init-fini | \
-// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,NOVPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI
+// RUN:   -fptrauth-init-fini \
+// RUN:   -fptrauth-init-fini-address-discrimination | \
+// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,NOVPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI,INITFINI_ADDR_DISCR
 
 // RUN: %clang_cc1 -E %s -triple=aarch64 \
 // RUN:   -fptrauth-intrinsics \
 // RUN:   -fptrauth-calls \
 // RUN:   -fptrauth-returns \
 // RUN:   -fptrauth-vtable-pointer-address-discrimination \
-// RUN:   -fptrauth-init-fini | \
-// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,VPTR_ADDR_DISCR,NOVPTR_TYPE_DISCR,INITFINI
+// RUN:   -fptrauth-init-fini \
+// RUN:   -fptrauth-init-fini-address-discrimination | \
+// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,VPTR_ADDR_DISCR,NOVPTR_TYPE_DISCR,INITFINI,INITFINI_ADDR_DISCR
+
+// RUN: %clang_cc1 -E %s -triple=aarch64 \
+// RUN:   -fptrauth-intrinsics \
+// RUN:   -fptrauth-calls \
+// RUN:   -fptrauth-returns \
+// RUN:   -fptrauth-vtable-pointer-address-discrimination \
+// RUN:   -fptrauth-vtable-pointer-type-discrimination \
+// RUN:   -fptrauth-init-fini-address-discrimination | \
+// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,NOINITFINI,INITFINI_ADDR_DISCR
 
 // RUN: %clang_cc1 -E %s -triple=aarch64 \
 // RUN:   -fptrauth-intrinsics \
 // RUN:   -fptrauth-calls \
 // RUN:   -fptrauth-returns \
 // RUN:   -fptrauth-vtable-pointer-address-discrimination \
-// RUN:   -fptrauth-vtable-pointer-type-discrimination | \
-// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,NOINITFINI
+// RUN:   -fptrauth-vtable-pointer-type-discrimination \
+// RUN:   -fptrauth-init-fini | \
+// RUN:   FileCheck %s --check-prefixes=INTRIN,CALLS,RETS,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR,INITFINI,NOINITFINI_ADDR_DISCR
 
 #if __has_feature(ptrauth_intrinsics)
 // INTRIN: has_ptrauth_intrinsics
@@ -111,3 +127,11 @@ void has_ptrauth_init_fini() {}
 // NOINITFINI: no_ptrauth_init_fini
 void no_ptrauth_init_fini() {}
 #endif
+
+#if __has_feature(ptrauth_init_fini_address_discrimination)
+// INITFINI_ADDR_DISCR: has_ptrauth_init_fini_address_discrimination
+void has_ptrauth_init_fini_address_discrimination() {}
+#else
+// NOINITFINI_ADDR_DISCR: no_ptrauth_init_fini_address_discrimination
+void no_ptrauth_init_fini_address_discrimination() {}
+#endif
diff --git a/compiler-rt/lib/builtins/crtbegin.c b/compiler-rt/lib/builtins/crtbegin.c
index a0860ca12ea03..143af768de679 100644
--- a/compiler-rt/lib/builtins/crtbegin.c
+++ b/compiler-rt/lib/builtins/crtbegin.c
@@ -8,6 +8,10 @@
 
 #include <stddef.h>
 
+#if __has_feature(ptrauth_init_fini)
+#include <ptrauth.h>
+#endif
+
 __attribute__((visibility("hidden"))) void *__dso_handle = &__dso_handle;
 
 #ifdef EH_USE_FRAME_REGISTRY
@@ -46,8 +50,14 @@ static void __attribute__((used)) __do_init(void) {
 }
 
 #ifdef CRT_HAS_INITFINI_ARRAY
+#if __has_feature(ptrauth_init_fini)
+__attribute__((section(".init_array"),
+               used)) static void *__ptrauth_init_fini_pointer __init =
+    __do_init;
+#else
 __attribute__((section(".init_array"),
                used)) static void (*__init)(void) = __do_init;
+#endif
 #elif defined(__i386__) || defined(__x86_64__)
 __asm__(".pushsection .init,\"ax\", at progbits\n\t"
         "call __do_init\n\t"
@@ -103,8 +113,14 @@ static void __attribute__((used)) __do_fini(void) {
 }
 
 #ifdef CRT_HAS_INITFINI_ARRAY
+#if __has_feature(ptrauth_init_fini)
+__attribute__((section(".fini_array"),
+               used)) static void *__ptrauth_init_fini_pointer __fini =
+    __do_fini;
+#else
 __attribute__((section(".fini_array"),
                used)) static void (*__fini)(void) = __do_fini;
+#endif
 #elif defined(__i386__) || defined(__x86_64__)
 __asm__(".pushsection .fini,\"ax\", at progbits\n\t"
         "call __do_fini\n\t"
diff --git a/llvm/include/llvm/BinaryFormat/ELF.h b/llvm/include/llvm/BinaryFormat/ELF.h
index dfba180149916..a8227e3e1e2fa 100644
--- a/llvm/include/llvm/BinaryFormat/ELF.h
+++ b/llvm/include/llvm/BinaryFormat/ELF.h
@@ -1774,8 +1774,9 @@ enum : unsigned {
   AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_VPTRADDRDISCR = 4,
   AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_VPTRTYPEDISCR = 5,
   AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINI = 6,
+  AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINIADDRDISC = 7,
   AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_LAST =
-      AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINI,
+      AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINIADDRDISC,
 };
 
 // x86 processor feature bits.
diff --git a/llvm/test/CodeGen/AArch64/note-gnu-property-elf-pauthabi.ll b/llvm/test/CodeGen/AArch64/note-gnu-property-elf-pauthabi.ll
index 728cffeba02a2..4005ff461f5e2 100644
--- a/llvm/test/CodeGen/AArch64/note-gnu-property-elf-pauthabi.ll
+++ b/llvm/test/CodeGen/AArch64/note-gnu-property-elf-pauthabi.ll
@@ -27,7 +27,7 @@
 ; OBJ: Displaying notes found in: .note.gnu.property
 ; OBJ-NEXT:   Owner                 Data size	Description
 ; OBJ-NEXT:   GNU                   0x00000018	NT_GNU_PROPERTY_TYPE_0 (property note)
-; OBJ-NEXT:   AArch64 PAuth ABI core info: platform 0x10000002 (llvm_linux), version 0x55 (PointerAuthIntrinsics, !PointerAuthCalls, PointerAuthReturns, !PointerAuthAuthTraps, PointerAuthVTPtrAddressDiscrimination, !PointerAuthVTPtrTypeDiscrimination, PointerAuthInitFini)
+; OBJ-NEXT:   AArch64 PAuth ABI core info: platform 0x10000002 (llvm_linux), version 0x55 (PointerAuthIntrinsics, !PointerAuthCalls, PointerAuthReturns, !PointerAuthAuthTraps, PointerAuthVTPtrAddressDiscrimination, !PointerAuthVTPtrTypeDiscrimination, PointerAuthInitFini, !PointerAuthInitFiniAddressDiscrimination)
 
 ; ERR: either both or no 'aarch64-elf-pauthabi-platform' and 'aarch64-elf-pauthabi-version' module flags must be present
 
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-init-fini.ll b/llvm/test/CodeGen/AArch64/ptrauth-init-fini.ll
new file mode 100644
index 0000000000000..9ce413b51989b
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/ptrauth-init-fini.ll
@@ -0,0 +1,77 @@
+; RUN: rm -rf %t && split-file %s %t && cd %t
+
+;--- nodisc.ll
+
+; RUN: llc -mtriple aarch64-elf -mattr=+pauth -filetype=asm -o - nodisc.ll | \
+; RUN:   FileCheck %s --check-prefix=ASM
+; RUN: llc -mtriple aarch64-elf -mattr=+pauth -filetype=obj -o - nodisc.ll | \
+; RUN:   llvm-readelf -r -x .init_array -x .fini_array - | FileCheck %s --check-prefix=OBJ
+
+; ASM:      .section .init_array,"aw", at init_array
+; ASM-NEXT: .p2align 3, 0x0
+; ASM-NEXT: .xword   foo at AUTH(ia,55764)
+; ASM-NEXT: .section .fini_array,"aw", at fini_array
+; ASM-NEXT: .p2align 3, 0x0
+; ASM-NEXT: .xword   bar at AUTH(ia,55764)
+
+; OBJ:      Relocation section '.rela.init_array' at offset 0x[[#]] contains 1 entries:
+; OBJ-NEXT:     Offset             Info             Type               Symbol's Value  Symbol's Name + Addend
+; OBJ-NEXT: 0000000000000000  0000000700000244 R_AARCH64_AUTH_ABS64   0000000000000000 foo + 0
+; OBJ:      Relocation section '.rela.fini_array' at offset 0x[[#]] contains 1 entries:
+; OBJ-NEXT:     Offset             Info             Type               Symbol's Value  Symbol's Name + Addend
+; OBJ-NEXT: 0000000000000000  0000000800000244 R_AARCH64_AUTH_ABS64   0000000000000004 bar + 0
+; OBJ:      Hex dump of section '.init_array':
+; OBJ-NEXT: 0x00000000 00000000 d4d90000
+; OBJ:      Hex dump of section '.fini_array':
+; OBJ-NEXT: 0x00000000 00000000 d4d90000
+;;                              ^^^^ 0xD9D4: constant discriminator = 55764
+;;                                    ^^ 0x80: bits 61..60 key = IA; bit 63 addr disc = false
+
+ at llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764), ptr null }]
+ at llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764), ptr null }]
+
+define void @foo() {
+  ret void
+}
+
+define void @bar() {
+  ret void
+}
+
+;--- disc.ll
+
+; RUN: llc -mtriple aarch64-elf -mattr=+pauth -filetype=asm -o - disc.ll | \
+; RUN:   FileCheck %s --check-prefix=ASM-DISC
+; RUN: llc -mtriple aarch64-elf -mattr=+pauth -filetype=obj -o - disc.ll | \
+; RUN:   llvm-readelf -r -x .init_array -x .fini_array - | FileCheck %s --check-prefix=OBJ-DISC
+
+; ASM-DISC:      .section .init_array,"aw", at init_array
+; ASM-DISC-NEXT: .p2align 3, 0x0
+; ASM-DISC-NEXT: .xword   foo at AUTH(ia,55764,addr)
+; ASM-DISC-NEXT: .section .fini_array,"aw", at fini_array
+; ASM-DISC-NEXT: .p2align 3, 0x0
+; ASM-DISC-NEXT: .xword   bar at AUTH(ia,55764,addr)
+
+; OBJ-DISC:      Relocation section '.rela.init_array' at offset 0x[[#]] contains 1 entries:
+; OBJ-DISC-NEXT:     Offset             Info             Type               Symbol's Value  Symbol's Name + Addend
+; OBJ-DISC-NEXT: 0000000000000000  0000000700000244 R_AARCH64_AUTH_ABS64   0000000000000000 foo + 0
+; OBJ-DISC:      Relocation section '.rela.fini_array' at offset 0x[[#]] contains 1 entries:
+; OBJ-DISC-NEXT:     Offset             Info             Type               Symbol's Value  Symbol's Name + Addend
+; OBJ-DISC-NEXT: 0000000000000000  0000000800000244 R_AARCH64_AUTH_ABS64   0000000000000004 bar + 0
+; OBJ-DISC:      Hex dump of section '.init_array':
+; OBJ-DISC-NEXT: 0x00000000 00000000 d4d90080
+; OBJ-DISC:      Hex dump of section '.fini_array':
+; OBJ-DISC-NEXT: 0x00000000 00000000 d4d90080
+;;                                   ^^^^ 0xD9D4: constant discriminator = 55764
+;;                                         ^^ 0x80: bits 61..60 key = IA; bit 63 addr disc = true
+
+ at llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764, ptr getelementptr inbounds ([1 x { i32, ptr, ptr }], ptr @llvm.global_ctors, i32 0, i32 0, i32 1)), ptr null }]
+ at llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764, ptr getelementptr inbounds ([1 x { i32, ptr, ptr }], ptr @llvm.global_dtors, i32 0, i32 0, i32 1)), ptr null }]
+
+define void @foo() {
+  ret void
+}
+
+define void @bar() {
+  ret void
+}
diff --git a/llvm/test/tools/llvm-readobj/ELF/AArch64/aarch64-feature-pauth.s b/llvm/test/tools/llvm-readobj/ELF/AArch64/aarch64-feature-pauth.s
index 512531748cd25..67bf4ca6d2887 100644
--- a/llvm/test/tools/llvm-readobj/ELF/AArch64/aarch64-feature-pauth.s
+++ b/llvm/test/tools/llvm-readobj/ELF/AArch64/aarch64-feature-pauth.s
@@ -106,12 +106,12 @@ end:
 # RUN: llvm-mc -filetype=obj -triple aarch64-linux-gnu gnu-0x10000002-85.s -o gnu-0x10000002-85.o
 # RUN: llvm-readelf --notes gnu-0x10000002-85.o | \
 # RUN:   FileCheck --check-prefix=ELF -DPLATFORM="0x10000002 (llvm_linux)" \
-# RUN:   -DVERSION="0x55 (PointerAuthIntrinsics, !PointerAuthCalls, PointerAuthReturns, !PointerAuthAuthTraps, PointerAuthVTPtrAddressDiscrimination, !PointerAuthVTPtrTypeDiscrimination, PointerAuthInitFini)" %s
+# RUN:   -DVERSION="0x55 (PointerAuthIntrinsics, !PointerAuthCalls, PointerAuthReturns, !PointerAuthAuthTraps, PointerAuthVTPtrAddressDiscrimination, !PointerAuthVTPtrTypeDiscrimination, PointerAuthInitFini, !PointerAuthInitFiniAddressDiscrimination)" %s
 # RUN: llvm-readobj --notes gnu-0x10000002-85.o | \
 # RUN:   FileCheck --check-prefix=OBJ -DPLATFORM="0x10000002 (llvm_linux)" \
-# RUN:   -DVERSION="0x55 (PointerAuthIntrinsics, !PointerAuthCalls, PointerAuthReturns, !PointerAuthAuthTraps, PointerAuthVTPtrAddressDiscrimination, !PointerAuthVTPtrTypeDiscrimination, PointerAuthInitFini)" %s
+# RUN:   -DVERSION="0x55 (PointerAuthIntrinsics, !PointerAuthCalls, PointerAuthReturns, !PointerAuthAuthTraps, PointerAuthVTPtrAddressDiscrimination, !PointerAuthVTPtrTypeDiscrimination, PointerAuthInitFini, !PointerAuthInitFiniAddressDiscrimination)" %s
 
-#--- gnu-0x10000002-128.s
+#--- gnu-0x10000002-256.s
 .section ".note.gnu.property", "a"
   .long 4           // Name length is always 4 ("GNU")
   .long end - begin // Data length
@@ -123,15 +123,15 @@ begin:
   .long 0xc0000001  // Type: GNU_PROPERTY_AARCH64_FEATURE_PAUTH
   .long 16          // Data size
   .quad 0x10000002  // PAuth ABI platform
-  .quad 128         // PAuth ABI version
+  .quad 256         // PAuth ABI version
   .p2align 3        // Align to 8 byte for 64 bit
 end:
 
-# RUN: llvm-mc -filetype=obj -triple aarch64-linux-gnu gnu-0x10000002-128.s -o gnu-0x10000002-128.o
-# RUN: llvm-readelf --notes gnu-0x10000002-128.o | \
-# RUN:   FileCheck --check-prefix=ELF -DPLATFORM="0x10000002 (llvm_linux)" -DVERSION="0x80 (unknown)" %s
-# RUN: llvm-readobj --notes gnu-0x10000002-128.o | \
-# RUN:   FileCheck --check-prefix=OBJ -DPLATFORM="0x10000002 (llvm_linux)" -DVERSION="0x80 (unknown)" %s
+# RUN: llvm-mc -filetype=obj -triple aarch64-linux-gnu gnu-0x10000002-256.s -o gnu-0x10000002-256.o
+# RUN: llvm-readelf --notes gnu-0x10000002-256.o | \
+# RUN:   FileCheck --check-prefix=ELF -DPLATFORM="0x10000002 (llvm_linux)" -DVERSION="0x100 (unknown)" %s
+# RUN: llvm-readobj --notes gnu-0x10000002-256.o | \
+# RUN:   FileCheck --check-prefix=OBJ -DPLATFORM="0x10000002 (llvm_linux)" -DVERSION="0x100 (unknown)" %s
 
 #--- gnu-short.s
 .section ".note.gnu.property", "a"
diff --git a/llvm/tools/llvm-readobj/ELFDumper.cpp b/llvm/tools/llvm-readobj/ELFDumper.cpp
index c696934a959b2..e53db899c5f41 100644
--- a/llvm/tools/llvm-readobj/ELFDumper.cpp
+++ b/llvm/tools/llvm-readobj/ELFDumper.cpp
@@ -5216,8 +5216,10 @@ static bool printAArch64PAuthABICoreInfo(raw_ostream &OS, uint32_t DataSize,
     Flags[AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_VPTRTYPEDISCR] =
         "VTPtrTypeDiscrimination";
     Flags[AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINI] = "InitFini";
+    Flags[AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINIADDRDISC] =
+        "InitFiniAddressDiscrimination";
 
-    static_assert(AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINI ==
+    static_assert(AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_INITFINIADDRDISC ==
                       AARCH64_PAUTH_PLATFORM_LLVM_LINUX_VERSION_LAST,
                   "Update when new enum items are defined");
 

>From 5872beb3188b6e66c745e4a1c6eb91a54cecc123 Mon Sep 17 00:00:00 2001
From: Daniil Kovalev <dkovalev at accesssoftek.com>
Date: Fri, 28 Jun 2024 15:28:49 +0300
Subject: [PATCH 2/2] Implement special address discrimination for ctors/dtors

---
 clang/lib/CodeGen/CodeGenModule.cpp           |  6 +++-
 clang/test/CodeGen/ptrauth-init-fini.c        |  4 +--
 llvm/include/llvm/IR/Constants.h              | 13 ++++++++
 llvm/lib/IR/Constants.cpp                     | 12 +++++++
 llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp | 19 ++++++++++++
 .../test/CodeGen/AArch64/ptrauth-init-fini.ll | 31 +++++++++++++++++--
 6 files changed, 80 insertions(+), 5 deletions(-)

diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 5bd34b9edd9d6..4ff15c4bee7b1 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -2074,7 +2074,11 @@ void CodeGenModule::EmitCtorList(CtorList &Fns, const char *GlobalName) {
     if (InitFiniAuthSchema) {
       llvm::Constant *StorageAddress =
           (InitFiniAuthSchema.isAddressDiscriminated()
-               ? StorageAddress = Ctor.getAddrOfCurrentPosition(PtrTy)
+               ? llvm::ConstantExpr::getIntToPtr(
+                     llvm::ConstantInt::get(
+                         IntPtrTy,
+                         llvm::ConstantPtrAuth::AddrDiscriminator_CtorsDtors),
+                     PtrTy)
                : nullptr);
       llvm::Constant *SignedCtorPtr = getConstantSignedPointer(
           I.Initializer, InitFiniAuthSchema.getKey(), StorageAddress,
diff --git a/clang/test/CodeGen/ptrauth-init-fini.c b/clang/test/CodeGen/ptrauth-init-fini.c
index 8b1abd0681c90..d51ed1d98e81e 100644
--- a/clang/test/CodeGen/ptrauth-init-fini.c
+++ b/clang/test/CodeGen/ptrauth-init-fini.c
@@ -18,8 +18,8 @@
 // SIGNED: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764), ptr null }]
 // SIGNED: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764), ptr null }]
 
-// ADDRDISC: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764, ptr getelementptr inbounds ([1 x { i32, ptr, ptr }], ptr @llvm.global_ctors, i32 0, i32 0, i32 1)), ptr null }]
-// ADDRDISC: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764, ptr getelementptr inbounds ([1 x { i32, ptr, ptr }], ptr @llvm.global_dtors, i32 0, i32 0, i32 1)), ptr null }]
+// ADDRDISC: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764, ptr inttoptr (i64 1 to ptr)), ptr null }]
+// ADDRDISC: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764, ptr inttoptr (i64 1 to ptr)), ptr null }]
 
 // UNSIGNED: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @foo, ptr null }]
 // UNSIGNED: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @bar, ptr null }]
diff --git a/llvm/include/llvm/IR/Constants.h b/llvm/include/llvm/IR/Constants.h
index 4d13c3880c692..2788751e8b62a 100644
--- a/llvm/include/llvm/IR/Constants.h
+++ b/llvm/include/llvm/IR/Constants.h
@@ -1056,6 +1056,19 @@ class ConstantPtrAuth final : public Constant {
     return !getAddrDiscriminator()->isNullValue();
   }
 
+  /// A constant value for the address discriminator which has special
+  /// significance to ctors/dtors lowering. Regular address discrimination can't
+  /// be applied for them since uses of llvm.global_{c|d}tors are disallowed
+  /// (see Verifier::visitGlobalVariable) and we can't emit getelementptr
+  /// expressions referencing these special arrays.
+  enum { AddrDiscriminator_CtorsDtors = 1 };
+
+  /// Whether the address uses a special address discriminator.
+  /// These discriminators can't be used in real pointer-auth values; they
+  /// can only be used in "prototype" values that indicate how some real
+  /// schema is supposed to be produced.
+  bool hasSpecialAddressDiscriminator(uint64_t Value) const;
+
   /// Check whether an authentication operation with key \p Key and (possibly
   /// blended) discriminator \p Discriminator is known to be compatible with
   /// this ptrauth signed pointer.
diff --git a/llvm/lib/IR/Constants.cpp b/llvm/lib/IR/Constants.cpp
index bc91f904d7e87..9223574df43d3 100644
--- a/llvm/lib/IR/Constants.cpp
+++ b/llvm/lib/IR/Constants.cpp
@@ -2060,6 +2060,18 @@ Value *ConstantPtrAuth::handleOperandChangeImpl(Value *From, Value *ToV) {
       Values, this, From, To, NumUpdated, OperandNo);
 }
 
+bool ConstantPtrAuth::hasSpecialAddressDiscriminator(uint64_t Value) const {
+  const auto *CastV = dyn_cast<ConstantExpr>(getAddrDiscriminator());
+  if (!CastV || CastV->getOpcode() != Instruction::IntToPtr)
+    return false;
+
+  const auto *IntVal = dyn_cast<ConstantInt>(CastV->getOperand(0));
+  if (!IntVal)
+    return false;
+
+  return IntVal->getValue() == Value;
+}
+
 bool ConstantPtrAuth::isKnownCompatibleWith(const Value *Key,
                                             const Value *Discriminator,
                                             const DataLayout &DL) const {
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 00b8d7bd257e4..19f1d109e2375 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -101,6 +101,8 @@ class AArch64AsmPrinter : public AsmPrinter {
 
   void emitFunctionEntryLabel() override;
 
+  void emitXXStructor(const DataLayout &DL, const Constant *CV) override;
+
   void LowerJumpTableDest(MCStreamer &OutStreamer, const MachineInstr &MI);
 
   void LowerMOPS(MCStreamer &OutStreamer, const MachineInstr &MI);
@@ -1245,6 +1247,23 @@ void AArch64AsmPrinter::emitFunctionEntryLabel() {
   }
 }
 
+void AArch64AsmPrinter::emitXXStructor(const DataLayout &DL,
+                                       const Constant *CV) {
+  if (const auto *CPA = dyn_cast<ConstantPtrAuth>(CV))
+    if (CPA->hasAddressDiscriminator() &&
+        !CPA->hasSpecialAddressDiscriminator(
+            ConstantPtrAuth::AddrDiscriminator_CtorsDtors))
+      report_fatal_error(
+          "unexpected address discrimination value for ctors/dtors entry, only "
+          "'ptr inttoptr (i64 1 to ptr)' is allowed");
+  // If we have signed pointers in xxstructors list, they'll be lowered to @AUTH
+  // MCExpr's via AArch64AsmPrinter::lowerConstantPtrAuth. It does not look at
+  // actual address discrimination value and only checks
+  // hasAddressDiscriminator(), so it's OK to leave special address
+  // discrimination value here.
+  AsmPrinter::emitXXStructor(DL, CV);
+}
+
 /// Small jump tables contain an unsigned byte or half, representing the offset
 /// from the lowest-addressed possible destination to the desired basic
 /// block. Since all instructions are 4-byte aligned, this is further compressed
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-init-fini.ll b/llvm/test/CodeGen/AArch64/ptrauth-init-fini.ll
index 9ce413b51989b..186a31c63ba10 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-init-fini.ll
+++ b/llvm/test/CodeGen/AArch64/ptrauth-init-fini.ll
@@ -65,8 +65,8 @@ define void @bar() {
 ;;                                   ^^^^ 0xD9D4: constant discriminator = 55764
 ;;                                         ^^ 0x80: bits 61..60 key = IA; bit 63 addr disc = true
 
- at llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764, ptr getelementptr inbounds ([1 x { i32, ptr, ptr }], ptr @llvm.global_ctors, i32 0, i32 0, i32 1)), ptr null }]
- at llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764, ptr getelementptr inbounds ([1 x { i32, ptr, ptr }], ptr @llvm.global_dtors, i32 0, i32 0, i32 1)), ptr null }]
+ at llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764, ptr inttoptr (i64 1 to ptr)), ptr null }]
+ at llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764, ptr inttoptr (i64 1 to ptr)), ptr null }]
 
 define void @foo() {
   ret void
@@ -75,3 +75,30 @@ define void @foo() {
 define void @bar() {
   ret void
 }
+
+;--- err1.ll
+
+; RUN: not --crash llc -mtriple aarch64-elf -mattr=+pauth -filetype=asm -o - err1.ll 2>&1 | \
+; RUN:   FileCheck %s --check-prefix=ERR1
+
+; ERR1: LLVM ERROR: unexpected address discrimination value for ctors/dtors entry, only 'ptr inttoptr (i64 1 to ptr)' is allowed
+
+ at llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @foo, i32 0, i64 55764, ptr inttoptr (i64 2 to ptr)), ptr null }]
+
+define void @foo() {
+  ret void
+}
+
+;--- err2.ll
+
+; RUN: not --crash llc -mtriple aarch64-elf -mattr=+pauth -filetype=asm -o - err2.ll 2>&1 | \
+; RUN:   FileCheck %s --check-prefix=ERR2
+
+; ERR2: LLVM ERROR: unexpected address discrimination value for ctors/dtors entry, only 'ptr inttoptr (i64 1 to ptr)' is allowed
+
+ at g = external global ptr
+ at llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr ptrauth (ptr @bar, i32 0, i64 55764, ptr @g), ptr null }]
+
+define void @bar() {
+  ret void
+}



More information about the cfe-commits mailing list