[AArch64][PAC] Support init/fini array signing

Implement an option `-fptrauth-init-fini` allowing to sign pointers in
init/fini arrays with a constant discriminator distinguishing them from
other function pointers.

The option is automatically enabled when building with
`-mbranch-protection=pauthabi` and disabled by default otherwise.

Author: kovdan01

clang/include/clang/Basic/Features.def

@@ -119,6 +119,7 @@ FEATURE(ptrauth_member_function_pointer_type_discrimination, LangOpts.PointerAut
 FEATURE(ptrauth_function_pointer_type_discrimination, LangOpts.FunctionPointerTypeDiscrimination)
 FEATURE(ptrauth_signed_block_descriptors, LangOpts.PointerAuthBlockDescriptorPointers)
 FEATURE(ptrauth_objc_isa_masking, LangOpts.PointerAuthObjcIsaMasking)
+FEATURE(ptrauth_init_fini, LangOpts.PointerAuthInitFini)
 FEATURE(swiftasynccc,
   PP.getTargetInfo().checkCallingConvention(CC_SwiftAsync) ==
   clang::TargetInfo::CCCR_OK)

clang/include/clang/Basic/LangOptions.def

@@ -166,6 +166,7 @@ LANGOPT(PointerAuthReturns, 1, 0, "return pointer authentication")
 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")
 ENUM_LANGOPT(PointerAuthObjcIsaAuthentication, PointerAuthenticationMode, 2,
              PointerAuthenticationMode::None, "authentication mode for objc isa")
 LANGOPT(PointerAuthObjcIsaMasking, 1, 0, "pre- and post-authentication masking for Objective-C isa pointer")

clang/include/clang/Basic/PointerAuthOptions.h

@@ -41,6 +41,10 @@ constexpr uint16_t SuperPointerConstantDiscriminator = 0xB5AB;
 /// is ptrauth_string_discriminator("block_descriptor")
 constexpr uint16_t BlockDescriptorConstantDiscriminator = 0xC0BB;
 
+/// 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 int PointerAuthKeyNone = -1;
 
 class PointerAuthSchema {
@@ -64,6 +68,7 @@ class PointerAuthSchema {
     ObjCMethodListPointer = 7,
     ObjCIsaPointer = 8,
     BlockDescriptorPointers = 9,
+    InitFiniPointers = 10,
   };
 
   /// Hardware pointer-signing keys in ARM8.3.
@@ -277,6 +282,9 @@ struct PointerAuthOptions {
 
   /// The ABI for Objective-C superclass pointers.
   PointerAuthSchema ObjCSuperPointers;
+
+  /// The ABI for function addresses in .init_array and .fini_array
+  PointerAuthSchema InitFiniPointers;
 };
 
 }  // end namespace clang

clang/include/clang/Driver/Options.td

@@ -3895,6 +3895,8 @@ let Group = f_Group in {
       HelpText<"enable signing block descriptors">;
     def fptrauth_function_pointer_type_discrimination : Flag<["-"], "fptrauth-function-pointer-type-discrimination">,
       HelpText<"Enabling type discrimination on C function pointers">;
+    def fptrauth_init_fini : Flag<["-"], "fptrauth-init-fini">,
+      HelpText<"Enable signing of function pointers in init/fini arrays">;
     def fptrauth_objc_isa : Flag<["-"], "fptrauth-objc-isa">,
       HelpText<"Enable signing and authentication of Objective-C object's 'isa' field">;
     def fptrauth_objc_isa_mode : Joined<["-"], "fptrauth-objc-isa-mode=">,
@@ -3912,6 +3914,7 @@ let Group = f_Group in {
   def fno_ptrauth_objc_isa : Flag<["-"], "fno-ptrauth-objc-isa">;
   def fno_ptrauth_block_descriptor_pointers : Flag<["-"], "fno-ptrauth-block-descriptor-pointers">;
   def fno_ptrauth_function_pointer_type_discrimination : Flag<["-"], "fno-ptrauth-function-pointer-type-discrimination">;
+  def fno_ptrauth_init_fini : Flag<["-"], "fno-ptrauth-init-fini">;
   def fno_branch_target_identification : Flag<["-"], "fno-branch-target-identification">;
 }
 

clang/lib/CodeGen/CodeGenModule.cpp

@@ -1944,37 +1944,81 @@ 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());
-
-  // Get the type of a ctor entry, { i32, void ()*, i8* }.
-  llvm::StructType *CtorStructTy = llvm::StructType::get(
-      Int32Ty, CtorPFTy, VoidPtrTy);
-
-  // Construct the constructor and destructor arrays.
-  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(llvm::ConstantExpr::getBitCast(I.Initializer, CtorPFTy));
-    if (I.AssociatedData)
-      ctor.add(llvm::ConstantExpr::getBitCast(I.AssociatedData, VoidPtrTy));
-    else
-      ctor.addNullPointer(VoidPtrTy);
-    ctor.finishAndAddTo(ctors);
-  }
+  const PointerAuthSchema &InitFiniAuthSchema =
+      getCodeGenOpts().PointerAuth.InitFiniPointers;
+  if (InitFiniAuthSchema) {
+    // Ctor function type is void()*.
+    llvm::FunctionType *CtorFTy = llvm::FunctionType::get(VoidTy, false);
+    llvm::Type *CtorPFTy = llvm::PointerType::get(
+        CtorFTy, TheModule.getDataLayout().getProgramAddressSpace());
+
+    llvm::StructType *CtorPauthTy =
+        llvm::StructType::get(CtorPFTy, Int32Ty, Int64Ty, Int64Ty);
+    llvm::Type *CtorPauthPTy = llvm::PointerType::get(
+        CtorPauthTy, TheModule.getDataLayout().getProgramAddressSpace());
+
+           // Get the type of a ctor entry, { i32, ptr, i8* }.
+    llvm::StructType *CtorStructTy =
+        llvm::StructType::get(Int32Ty, CtorPauthPTy, VoidPtrTy);
+
+           // Construct the constructor and destructor arrays.
+    ConstantInitBuilder builder(*this);
+    auto ctors = builder.beginArray(CtorStructTy);
+    for (const auto &I : Fns) {
+      auto ctor = ctors.beginStruct(CtorStructTy);
+      ctor.addInt(Int32Ty, I.Priority);
+      assert(!InitFiniAuthSchema.isAddressDiscriminated());
+      llvm::Constant *signedCtorAddr = getConstantSignedPointer(
+          I.Initializer, InitFiniAuthSchema.getKey(), 0,
+          llvm::ConstantInt::get(
+              SizeTy, InitFiniAuthSchema.getConstantDiscrimination()));
+      ctor.add(llvm::ConstantExpr::getBitCast(signedCtorAddr, CtorPauthPTy));
+      if (I.AssociatedData)
+        ctor.add(llvm::ConstantExpr::getBitCast(I.AssociatedData, VoidPtrTy));
+      else
+        ctor.addNullPointer(VoidPtrTy);
+      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);
+  } else {
+    // Ctor function type is void()*.
+    llvm::FunctionType *CtorFTy = llvm::FunctionType::get(VoidTy, false);
+    llvm::Type *CtorPFTy = llvm::PointerType::get(
+        CtorFTy, TheModule.getDataLayout().getProgramAddressSpace());
+
+           // Get the type of a ctor entry, { i32, void ()*, i8* }.
+    llvm::StructType *CtorStructTy =
+        llvm::StructType::get(Int32Ty, CtorPFTy, VoidPtrTy);
+
+           // Construct the constructor and destructor arrays.
+    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(llvm::ConstantExpr::getBitCast(I.Initializer, CtorPFTy));
+      if (I.AssociatedData)
+        ctor.add(llvm::ConstantExpr::getBitCast(I.AssociatedData, VoidPtrTy));
+      else
+        ctor.addNullPointer(VoidPtrTy);
+      ctor.finishAndAddTo(ctors);
+    }
 
-  // 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);
+    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);
+  }
 
   Fns.clear();
 }

clang/lib/Driver/ToolChains/Clang.cpp

@@ -1626,12 +1626,6 @@ static void handlePAuthABIOption(const ArgList &DriverArgs,
                           options::OPT_fno_ptrauth_auth_traps))
     CC1Args.push_back("-fptrauth-auth-traps");
 
-#if 1
-  if (!DriverArgs.hasArg(
-          options::OPT_fptrauth_block_descriptor_pointers,
-          options::OPT_fno_ptrauth_block_descriptor_pointers))
-    CC1Args.push_back("-fptrauth-block-descriptor-pointers");
-
   if (!DriverArgs.hasArg(
       options::OPT_fptrauth_vtable_pointer_address_discrimination,
       options::OPT_fno_ptrauth_vtable_pointer_address_discrimination))
@@ -1642,6 +1636,13 @@ static void handlePAuthABIOption(const ArgList &DriverArgs,
       options::OPT_fno_ptrauth_vtable_pointer_type_discrimination))
     CC1Args.push_back("-fptrauth-vtable-pointer-type-discrimination");
 
+  if (!DriverArgs.hasArg(options::OPT_fptrauth_init_fini,
+                         options::OPT_fno_ptrauth_init_fini))
+    CC1Args.push_back("-fptrauth-init-fini");
+
+#if 0
+  // Due to implicit casts, code with function pointer type discrimination
+  // enabled might be broken.
   if (!DriverArgs.hasArg(
       options::OPT_fptrauth_function_pointer_type_discrimination,
       options::OPT_fno_ptrauth_function_pointer_type_discrimination))
@@ -7079,6 +7080,11 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
                    false))
     CmdArgs.push_back("-fptrauth-vtable-pointer-type-discrimination");
 
+  if (Args.hasFlag(options::OPT_fptrauth_init_fini,
+                   options::OPT_fno_ptrauth_init_fini,
+                   false))
+    CmdArgs.push_back("-fptrauth-init-fini");
+
   if (Args.hasFlag(options::OPT_fptrauth_objc_isa,
                    options::OPT_fno_ptrauth_objc_isa, false))
     CmdArgs.push_back("-fptrauth-objc-isa-mode=sign-and-auth");

clang/lib/Frontend/CompilerInvocation.cpp

@@ -1491,6 +1491,11 @@ bool CompilerInvocation::setDefaultPointerAuthOptions(
             Key::ObjCIsaPointer, true, explicitIsaAuthenticationMode,
             Discrimination::Constant, SuperPointerConstantDiscriminator);
       }
+      if (LangOpts.PointerAuthInitFini) {
+        Opts.InitFiniPointers = PointerAuthSchema(
+            Key::InitFiniPointers, false, Discrimination::Constant,
+            InitFiniPointerConstantDiscriminator);
+      }
     }
     Opts.ReturnAddresses = LangOpts.PointerAuthReturns;
     Opts.AuthTraps = LangOpts.PointerAuthAuthTraps;
@@ -1546,6 +1551,11 @@ bool CompilerInvocation::setDefaultPointerAuthOptions(
             Key::ASDA, true, explicitIsaAuthenticationMode,
             Discrimination::Constant, SuperPointerConstantDiscriminator);
       }
+      if (LangOpts.PointerAuthInitFini) {
+        Opts.InitFiniPointers =
+            PointerAuthSchema(Key::ASIA, false, Discrimination::Constant,
+                              InitFiniPointerConstantDiscriminator);
+      }
     }
     Opts.ReturnAddresses = LangOpts.PointerAuthReturns;
     Opts.AuthTraps = LangOpts.PointerAuthAuthTraps;
@@ -3443,6 +3453,8 @@ static void GeneratePointerAuthArgs(const LangOptions &Opts,
   }
   if (Opts.PointerAuthObjcIsaMasking)
     GenerateArg(Consumer, OPT_fptrauth_objc_isa_masking);
+  if (Opts.PointerAuthInitFini)
+    GenerateArg(Consumer, OPT_fptrauth_init_fini);
 }
 
 static void ParsePointerAuthArgs(LangOptions &Opts, ArgList &Args,
@@ -3455,6 +3467,7 @@ static void ParsePointerAuthArgs(LangOptions &Opts, ArgList &Args,
       Args.hasArg(OPT_fptrauth_vtable_pointer_address_discrimination);
   Opts.PointerAuthVTPtrTypeDiscrimination =
       Args.hasArg(OPT_fptrauth_vtable_pointer_type_discrimination);
+  Opts.PointerAuthInitFini = Args.hasArg(OPT_fptrauth_init_fini);
   Opts.SoftPointerAuth = Args.hasArg(OPT_fptrauth_soft);
   Opts.PointerAuthBlockDescriptorPointers =
       Args.hasArg(OPT_fptrauth_block_descriptor_pointers);

clang/lib/Headers/ptrauth.h

@@ -80,6 +80,9 @@ typedef enum {
   /* The key used to sign block descriptor pointers. */
   ptrauth_key_block_descriptor_pointer = ptrauth_key_asda,
 
+  /* 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;
@@ -398,6 +401,9 @@ typedef __UINTPTR_TYPE__ ptrauth_generic_signature_t;
   __ptrauth(ptrauth_key_cxx_vtable_pointer,0,0)
 #define __ptrauth_swift_heap_object_destructor \
   __ptrauth(ptrauth_key_function_pointer,1,0xbbbf)
+#define __ptrauth_init_fini_discriminator 0xd9d4
+#define __ptrauth_init_fini_pointer                                            \
+  __ptrauth(ptrauth_key_init_fini_pointer, 0, __ptrauth_init_fini_discriminator)
 
 /* Some situations in the C++ and Swift ABIs use declaration-specific
    or type-specific extra discriminators. */

clang/test/CodeGen/ptrauth-constructor-attribute.c

@@ -0,0 +1,29 @@
+// REQUIRES: aarch64-registered-target
+// RUN: %clang -target aarch64-elf -march=armv8.3-a+pauth -mbranch-protection=pauthabi                        -S -emit-llvm -o - -c %s | FileCheck --check-prefix=SIGNED %s
+// RUN: %clang -target aarch64-elf -march=armv8.3-a+pauth -mbranch-protection=pauthabi -fptrauth-init-fini    -S -emit-llvm -o - -c %s | FileCheck --check-prefix=SIGNED %s
+// RUN: %clang -target aarch64-elf -march=armv8.3-a+pauth -mbranch-protection=pauthabi -fno-ptrauth-init-fini -S -emit-llvm -o - -c %s | FileCheck --check-prefix=UNSIGNED %s
+
+// SIGNED: @foo.ptrauth = private constant { ptr, i32, i64, i64 } { ptr @foo, i32 0, i64 0, i64 55764 }, section "llvm.ptrauth", align 8
+// SIGNED: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @foo.ptrauth, ptr null }]
+// SIGNED: @bar.ptrauth = private constant { ptr, i32, i64, i64 } { ptr @bar, i32 0, i64 0, i64 55764 }, section "llvm.ptrauth", align 8
+// SIGNED: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @bar.ptrauth, ptr null }]
+
+// UNSIGNED-NOT: @foo.ptrauth
+// UNSIGNED:     @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @foo, ptr null }]
+// UNSIGNED-NOT: @bar.ptrauth
+// 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;
+}

clang/test/Preprocessor/ptrauth_feature.c

@@ -3,6 +3,10 @@
 // RUN: %clang_cc1 %s -E -triple=arm64-- -fptrauth-returns | FileCheck %s --check-prefixes=NOCALLS,NOINTRIN,RETS,NOQUAL,NOFUNC
 // RUN: %clang_cc1 %s -E -triple=arm64-- -fptrauth-intrinsics | FileCheck %s --check-prefixes=NOCALLS,INTRIN,NORETS,QUAL,NOFUNC
 // RUN: %clang_cc1 %s -E -triple=arm64e-apple-ios6.0 -fptrauth-intrinsics -fptrauth-function-pointer-type-discrimination | FileCheck %s --check-prefixes=NOCALLS,INTRIN,NORETS,QUAL,FUNC
+// RUN: %clang -E %s --target=aarch64-elf -mbranch-protection=pauthabi | FileCheck %s --check-prefixes=INITFINI,VPTR_ADDR_DISCR,VPTR_TYPE_DISCR
+// RUN: %clang -E %s --target=aarch64-elf -mbranch-protection=pauthabi -fno-ptrauth-init-fini | FileCheck %s --check-prefixes=NOINITFINI
+// RUN: %clang -E %s --target=aarch64-elf -mbranch-protection=pauthabi -fno-ptrauth-vtable-pointer-address-discrimination | FileCheck %s --check-prefixes=NOVPTR_ADDR_DISCR
+// RUN: %clang -E %s --target=aarch64-elf -mbranch-protection=pauthabi -fno-ptrauth-vtable-pointer-type-discrimination | FileCheck %s --check-prefixes=NOVPTR_TYPE_DISCR
 
 #if __has_feature(ptrauth_calls)
 // CALLS: has_ptrauth_calls
@@ -45,6 +49,30 @@ void has_ptrauth_member_function_pointer_type_discrimination() {}
 void no_ptrauth_member_function_pointer_type_discrimination() {}
 #endif
 
+#if __has_feature(ptrauth_vtable_pointer_address_discrimination)
+// VPTR_ADDR_DISCR: has_ptrauth_vtable_pointer_address_discrimination
+void has_ptrauth_vtable_pointer_address_discrimination() {}
+#else
+// NOVPTR_ADDR_DISCR: no_ptrauth_vtable_pointer_address_discrimination
+void no_ptrauth_vtable_pointer_address_discrimination() {}
+#endif
+
+#if __has_feature(ptrauth_vtable_pointer_type_discrimination)
+// VPTR_TYPE_DISCR: has_ptrauth_vtable_pointer_type_discrimination
+void has_ptrauth_vtable_pointer_type_discrimination() {}
+#else
+// NOVPTR_TYPE_DISCR: no_ptrauth_vtable_pointer_type_discrimination
+void no_ptrauth_vtable_pointer_type_discrimination() {}
+#endif
+
+#if __has_feature(ptrauth_init_fini)
+// INITFINI: has_ptrauth_init_fini
+void has_ptrauth_init_fini() {}
+#else
+// NOINITFINI: no_ptrauth_init_fini
+void no_ptrauth_init_fini() {}
+#endif
+
 #include <ptrauth.h>
 
 #if __has_feature(ptrauth_function_pointer_type_discrimination)