[AArch64][PAC] Handle signing of init/fini pointers in AsmPrinter

clang/include/clang/Basic/PointerAuthOptions.h

@@ -27,10 +27,6 @@ namespace clang {
 /// 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;
-
 /// Constant discriminator to be used with method list pointers. The value is
 /// ptrauth_string_discriminator("method_list_t")
 constexpr uint16_t MethodListPointerConstantDiscriminator = 0xC310;
@@ -224,9 +220,6 @@ struct PointerAuthOptions {
   /// The ABI for C++ member function pointers.
   PointerAuthSchema CXXMemberFunctionPointers;
 
-  /// The ABI for function addresses in .init_array and .fini_array
-  PointerAuthSchema InitFiniPointers;
-
   /// The ABI for block invocation function pointers.
   PointerAuthSchema BlockInvocationFunctionPointers;
 

clang/lib/CodeGen/CodeGenModule.cpp

@@ -1392,6 +1392,12 @@ void CodeGenModule::Release() {
 
     if (LangOpts.PointerAuthELFGOT)
       getModule().addModuleFlag(llvm::Module::Error, "ptrauth-elf-got", 1);
+    if (LangOpts.PointerAuthCalls && LangOpts.PointerAuthInitFini) {
+      getModule().addModuleFlag(llvm::Module::Error, "ptrauth-init-fini", 1);
+      if (LangOpts.PointerAuthInitFiniAddressDiscrimination)
+        getModule().addModuleFlag(
+            llvm::Module::Error, "ptrauth-init-fini-address-discrimination", 1);
+    }
 
     if (getTriple().isOSLinux()) {
       if (LangOpts.PointerAuthCalls)
@@ -2432,9 +2438,6 @@ void CodeGenModule::AddGlobalDtor(llvm::Function *Dtor, int Priority,
 void CodeGenModule::EmitCtorList(CtorList &Fns, const char *GlobalName) {
   if (Fns.empty()) return;
 
-  const PointerAuthSchema &InitFiniAuthSchema =
-      getCodeGenOpts().PointerAuth.InitFiniPointers;
-
   // Ctor function type is ptr.
   llvm::PointerType *PtrTy = llvm::PointerType::get(
       getLLVMContext(), TheModule.getDataLayout().getProgramAddressSpace());
@@ -2448,23 +2451,7 @@ void CodeGenModule::EmitCtorList(CtorList &Fns, const char *GlobalName) {
   for (const auto &I : Fns) {
     auto Ctor = Ctors.beginStruct(CtorStructTy);
     Ctor.addInt(Int32Ty, I.Priority);
-    if (InitFiniAuthSchema) {
-      llvm::Constant *StorageAddress =
-          (InitFiniAuthSchema.isAddressDiscriminated()
-               ? llvm::ConstantExpr::getIntToPtr(
-                     llvm::ConstantInt::get(
-                         IntPtrTy,
-                         llvm::ConstantPtrAuth::AddrDiscriminator_CtorsDtors),
-                     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);
-    }
+    Ctor.add(I.Initializer);
     if (I.AssociatedData)
       Ctor.add(I.AssociatedData);
     else

clang/lib/Frontend/CompilerInvocation.cpp

@@ -1475,12 +1475,6 @@ void CompilerInvocation::setDefaultPointerAuthOptions(
     Opts.CXXMemberFunctionPointers =
         PointerAuthSchema(Key::ASIA, false, Discrimination::Type);
 
-    if (LangOpts.PointerAuthInitFini) {
-      Opts.InitFiniPointers = PointerAuthSchema(
-          Key::ASIA, LangOpts.PointerAuthInitFiniAddressDiscrimination,
-          Discrimination::Constant, InitFiniPointerConstantDiscriminator);
-    }
-
     Opts.BlockInvocationFunctionPointers =
         PointerAuthSchema(Key::ASIA, true, Discrimination::None);
     Opts.BlockHelperFunctionPointers =

clang/test/CodeGen/ptrauth-init-fini.c

@@ -1,28 +1,36 @@
 // REQUIRES: aarch64-registered-target
 
 // RUN: %clang_cc1 -triple aarch64-elf -target-feature +pauth -fptrauth-calls -fptrauth-init-fini    \
-// RUN:   -emit-llvm %s -o - | FileCheck --check-prefix=SIGNED %s
+// RUN:   -emit-llvm %s -o - | FileCheck --check-prefix=COMMON,SIGNED %s
 
 // RUN: %clang_cc1 -triple aarch64-elf -target-feature +pauth -fptrauth-calls -fptrauth-init-fini    \
-// RUN:   -fptrauth-init-fini-address-discrimination -emit-llvm %s -o - | FileCheck --check-prefix=ADDRDISC %s
+// RUN:   -fptrauth-init-fini-address-discrimination -emit-llvm %s -o - | FileCheck --check-prefix=COMMON,ADDRDISC %s
 
 // RUN: %clang_cc1 -triple aarch64-elf -target-feature +pauth -fptrauth-calls \
-// RUN:   -emit-llvm %s -o - | FileCheck --check-prefix=UNSIGNED %s
+// RUN:   -emit-llvm %s -o - | FileCheck --check-prefix=COMMON,UNSIGNED %s
 
 // RUN: %clang_cc1 -triple aarch64-elf -target-feature +pauth -fptrauth-calls -fptrauth-init-fini-address-discrimination \
-// RUN:   -emit-llvm %s -o - | FileCheck --check-prefix=UNSIGNED %s
+// RUN:   -emit-llvm %s -o - | FileCheck --check-prefix=COMMON,UNSIGNED %s
 
 // RUN: %clang_cc1 -triple aarch64-elf -target-feature +pauth                 -fptrauth-init-fini    \
-// RUN:   -emit-llvm %s -o - | FileCheck --check-prefix=UNSIGNED %s
+// RUN:   -emit-llvm %s -o - | FileCheck --check-prefix=COMMON,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 }]
+// COMMON: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @foo, ptr null }]
+// COMMON: @llvm.global_dtors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @bar, 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 }]
+// The below checks assume no other module flags happens to be set.
 
-// 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 }]
+// UNSIGNED-NOT: !llvm.module.flags
+// UNSIGNED-NOT: !"ptrauth-init-fini"
+// UNSIGNED-NOT: !"ptrauth-init-fini-address-discrimination"
+
+// SIGNED: !llvm.module.flags = !{!0}
+// SIGNED: !0 = !{i32 1, !"ptrauth-init-fini", i32 1}
+// SIGNED-NOT: !"ptrauth-init-fini-address-discrimination"
+
+// ADDRDISC: !llvm.module.flags = !{!0, !1}
+// ADDRDISC: !0 = !{i32 1, !"ptrauth-init-fini", i32 1}
+// ADDRDISC: !1 = !{i32 1, !"ptrauth-init-fini-address-discrimination", i32 1}
 
 volatile int x = 0;
 

llvm/lib/IR/AutoUpgrade.cpp

@@ -6165,12 +6165,109 @@ void llvm::UpgradeARCRuntime(Module &M) {
     UpgradeToIntrinsic(I.first, I.second);
 }
 
+// Upgrade from wrapping each pointer stored in the @llvm.global_(ctors|dtors)
+// with ptrauth constant expression to storing plain pointers in these arrays
+// and requesting the particular signing schema globally via module flags.
+//
+// Only perform the upgrade if all elements of *both* arrays agree on a common
+// signing schema. The processing of the array is *not* stopped on the first
+// null function pointer.
+static bool upgradePtrauthInitFiniArrays(Module &M) {
+  // Either "not decided yet" or whether we should request address diversity
+  // in addition to the basic constant diversity.
+  // There is no value representing "decided not to sign", as this results
+  // in immediate return from upgradePtrauthInitFiniArrays.
+  std::optional<bool> UseAddressDisc;
+
+  // Do not attempt upgrading if the new module flags already exist.
+  if (const NamedMDNode *ModFlags = M.getModuleFlagsMetadata()) {
+    for (const MDNode *Flag : ModFlags->operands()) {
+      if (Flag->getNumOperands() != 3)
+        continue;
+      const MDString *ID = dyn_cast_or_null<MDString>(Flag->getOperand(1));
+      if (ID && (ID->getString() == "ptrauth-init-fini" ||
+                 ID->getString() == "ptrauth-init-fini-address-discriminator"))
+        return false;
+    }
+  }
+
+  auto UpgradeSinglePointer = [&UseAddressDisc](Constant *CV) -> Constant * {
+    const unsigned ExpectedConstDisc = 0xD9D4;
+    const unsigned ExpectedAddressMarker = 1;
+
+    auto *CPA = dyn_cast<ConstantPtrAuth>(CV);
+    if (!CPA || !CPA->getDiscriminator()->equalsInt(ExpectedConstDisc))
+      return nullptr; // Nothing to upgrade or unknown pattern found.
+
+    bool HasAddressDisc;
+    if (!CPA->hasAddressDiscriminator())
+      HasAddressDisc = false;
+    else if (CPA->hasSpecialAddressDiscriminator(ExpectedAddressMarker))
+      HasAddressDisc = true;
+    else
+      return nullptr; // Unknown pattern.
+
+    if (UseAddressDisc && *UseAddressDisc != HasAddressDisc)
+      return nullptr; // Disagreement with the decided mode.
+
+    UseAddressDisc = HasAddressDisc;
+    return CPA->getPointer();
+  };
+
+  SmallVector<std::pair<GlobalVariable *, Constant *>> PendingUpgrades;
+  for (const char *Name : {"llvm.global_ctors", "llvm.global_dtors"}) {
+    auto *GV = dyn_cast_if_present<GlobalVariable>(M.getNamedValue(Name));
+    if (!GV || !GV->hasInitializer())
+      continue; // Skip, but it is okay to upgrade the other variable.
+
+    auto *Init = dyn_cast<ConstantArray>(GV->getInitializer());
+    if (!Init)
+      return false;
+
+    std::vector<Constant *> NewStructors;
+    NewStructors.reserve(Init->getNumOperands());
+    for (Use &U : Init->operands()) {
+      auto *Structor = dyn_cast<ConstantStruct>(U.get());
+      if (!Structor || Structor->getNumOperands() != 3)
+        return false;
+
+      Constant *Prio = Structor->getOperand(0);
+      Constant *Func = UpgradeSinglePointer(Structor->getOperand(1));
+      Constant *Arg = Structor->getOperand(2);
+      if (!Func)
+        return false;
+
+      NewStructors.push_back(
+          ConstantStruct::get(Structor->getType(), {Prio, Func, Arg}));
+    }
+
+    Constant *NewInit = ConstantArray::get(Init->getType(), NewStructors);
+    PendingUpgrades.push_back({GV, NewInit});
+  }
+
+  if (PendingUpgrades.empty())
+    return false;
+  assert(UseAddressDisc.has_value());
+
+  for (auto [GV, NewInit] : PendingUpgrades)
+    GV->setInitializer(NewInit);
+  M.addModuleFlag(Module::Error, "ptrauth-init-fini", 1);
+  if (UseAddressDisc.value())
+    M.addModuleFlag(Module::Error, "ptrauth-init-fini-address-discriminator",
+                    1);
+
+  return true;
+}
+
 bool llvm::UpgradeModuleFlags(Module &M) {
+  bool Changed = false;
+  Changed |= upgradePtrauthInitFiniArrays(M);
+
   NamedMDNode *ModFlags = M.getModuleFlagsMetadata();
   if (!ModFlags)
-    return false;
+    return Changed;
 
-  bool HasObjCFlag = false, HasClassProperties = false, Changed = false;
+  bool HasObjCFlag = false, HasClassProperties = false;
   bool HasSwiftVersionFlag = false;
   uint8_t SwiftMajorVersion, SwiftMinorVersion;
   uint32_t SwiftABIVersion;

llvm/lib/IR/Verifier.cpp

@@ -910,6 +910,19 @@ void Verifier::visitGlobalVariable(const GlobalVariable &GV) {
       Check(ETy->isPointerTy(), "wrong type for intrinsic global variable",
             &GV);
     }
+
+    auto *Init = GV.hasInitializer()
+                     ? dyn_cast<ConstantArray>(GV.getInitializer())
+                     : nullptr;
+    if (Init) {
+      for (const Use &U : Init->operands()) {
+        auto *Structor = dyn_cast<ConstantStruct>(U);
+        if (!Structor || Structor->getNumOperands() != 3)
+          continue;
+        Check(!isa<ConstantPtrAuth>(Structor->getOperand(1)),
+              "signing of ctors/dtors should be requested via module flags");
+      }
+    }
   }
 
   if (GV.hasName() && (GV.getName() == "llvm.used" ||
@@ -1961,26 +1974,49 @@ void Verifier::visitModuleFlags() {
   // Scan each flag, and track the flags and requirements.
   DenseMap<const MDString*, const MDNode*> SeenIDs;
   SmallVector<const MDNode*, 16> Requirements;
-  uint64_t PAuthABIPlatform = -1;
-  uint64_t PAuthABIVersion = -1;
+  std::optional<uint64_t> PAuthABIPlatform;
+  std::optional<uint64_t> PAuthABIVersion;
+  std::optional<uint64_t> HasPtrauthInitFini;
+  std::optional<uint64_t> HasPtrauthInitFiniAddr;
+
   for (const MDNode *MDN : Flags->operands()) {
     visitModuleFlag(MDN, SeenIDs, Requirements);
     if (MDN->getNumOperands() != 3)
       continue;
+
     if (const auto *FlagName = dyn_cast_or_null<MDString>(MDN->getOperand(1))) {
-      if (FlagName->getString() == "aarch64-elf-pauthabi-platform") {
-        if (const auto *PAP =
+      auto GetFlagNamed = [&](StringRef Name) -> std::optional<uint64_t> {
+        if (FlagName->getString() != Name)
+          return std::nullopt;
+        if (const auto *FlagValue =
                 mdconst::dyn_extract_or_null<ConstantInt>(MDN->getOperand(2)))
-          PAuthABIPlatform = PAP->getZExtValue();
-      } else if (FlagName->getString() == "aarch64-elf-pauthabi-version") {
-        if (const auto *PAV =
-                mdconst::dyn_extract_or_null<ConstantInt>(MDN->getOperand(2)))
-          PAuthABIVersion = PAV->getZExtValue();
-      }
+          return FlagValue->getZExtValue();
+
+        CheckFailed(Name + ": module flag expects integer value");
+        return std::nullopt;
+      };
+
+      if (auto Value = GetFlagNamed("aarch64-elf-pauthabi-platform"))
+        PAuthABIPlatform = *Value;
+      else if (auto Value = GetFlagNamed("aarch64-elf-pauthabi-version"))
+        PAuthABIVersion = *Value;
+      else if (auto Value = GetFlagNamed("ptrauth-init-fini"))
+        HasPtrauthInitFini = *Value;
+      else if (auto Value =
+                   GetFlagNamed("ptrauth-init-fini-address-discrimination"))
+        HasPtrauthInitFiniAddr = *Value;
     }
   }
 
-  if ((PAuthABIPlatform == uint64_t(-1)) != (PAuthABIVersion == uint64_t(-1)))
+  Check(!HasPtrauthInitFini || HasPtrauthInitFini.value() == 1,
+        "ptrauth-init-fini must be set to 1 or unset");
+  Check(!HasPtrauthInitFiniAddr || HasPtrauthInitFiniAddr.value() == 1,
+        "ptrauth-init-fini-address-discrimination must be set to 1 or unset");
+  if (HasPtrauthInitFiniAddr)
+    Check(HasPtrauthInitFini, "ptrauth-init-fini-address-discrimination module "
+                              "flag requires ptrauth-init-fini");
+
+  if (PAuthABIPlatform.has_value() != PAuthABIVersion.has_value())
     CheckFailed("either both or no 'aarch64-elf-pauthabi-platform' and "
                 "'aarch64-elf-pauthabi-version' module flags must be present");
 

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -98,6 +98,8 @@ class AArch64AsmPrinter : public AsmPrinter {
   FaultMaps FM;
   const AArch64Subtarget *STI;
   bool ShouldEmitWeakSwiftAsyncExtendedFramePointerFlags = false;
+  bool PtrauthInitFini = false;
+  bool PtrauthInitFiniAddressDisc = false;
 #ifndef NDEBUG
   unsigned InstsEmitted;
 #endif
@@ -405,6 +407,11 @@ void AArch64AsmPrinter::emitStartOfAsmFile(Module &M) {
       EnableImportCallOptimization = true;
   }
 
+  if (M.getModuleFlag("ptrauth-init-fini"))
+    PtrauthInitFini = true;
+  if (M.getModuleFlag("ptrauth-init-fini-address-discrimination"))
+    PtrauthInitFiniAddressDisc = true;
+
   if (!TT.isOSBinFormatELF())
     return;
 
@@ -1464,18 +1471,31 @@ 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.
+  LLVMContext &C = CV->getContext();
+  assert(!isa<ConstantPtrAuth>(CV) &&
+         "ctors/dtors are to be signed by asm printer");
+
+  if (PtrauthInitFini) {
+    IntegerType *Int32Ty = IntegerType::get(C, 32);
+    IntegerType *Int64Ty = IntegerType::get(C, 64);
+    PointerType *PtrTy = PointerType::get(C, 0);
+
+    ConstantInt *Key = ConstantInt::get(Int32Ty, AArch64PAuth::InitFiniKey);
+    ConstantInt *IntDisc = ConstantInt::get(
+        Int64Ty, AArch64PAuth::InitFiniPointerConstantDiscriminator);
+    Constant *Null = ConstantPointerNull::get(PtrTy);
+    Constant *AddressDisc = Null;
+    if (PtrauthInitFiniAddressDisc) {
+      uint64_t Marker = ConstantPtrAuth::AddrDiscriminator_CtorsDtors;
+      AddressDisc =
+          ConstantExpr::getIntToPtr(ConstantInt::get(Int64Ty, Marker), PtrTy);
+    }
+
+    CV = ConstantPtrAuth::get(const_cast<Constant *>(CV), Key, IntDisc,
+                              AddressDisc, /*DeactivationSymbol=*/Null);
+  }
+
+  // Signed pointers will be lowered by AArch64AsmPrinter::lowerConstantPtrAuth.
   AsmPrinter::emitXXStructor(DL, CV);
 }