[AArch64][PAC] Lower ptrauth constants in code for MachO. (#97665)

This also adds support for auth stubs on MachO using __DATA,__auth_ptr.

Some of the machinery for auth stubs is already implemented;  this
generalizes that a bit to support MachO, and moves some of the shared
logic into MMIImpls.

In particular, this originally had an AuthStubInfo struct, but we no
longer need it beyond a single MCExpr.  So this provides variants of
the symbol stub helper type declarations and functions for "expr
stubs", where a stub points at an arbitrary MCExpr, rather than
a simple MCSymbol (and a bit).

Author: ahmedbougacha

llvm/include/llvm/CodeGen/MachineModuleInfo.h

@@ -58,12 +58,20 @@ class MachineModuleInfoImpl {
   using StubValueTy = PointerIntPair<MCSymbol *, 1, bool>;
   using SymbolListTy = std::vector<std::pair<MCSymbol *, StubValueTy>>;
 
+  /// A variant of SymbolListTy where the stub is a generalized MCExpr.
+  using ExprStubListTy = std::vector<std::pair<MCSymbol *, const MCExpr *>>;
+
   virtual ~MachineModuleInfoImpl();
 
 protected:
   /// Return the entries from a DenseMap in a deterministic sorted orer.
   /// Clears the map.
   static SymbolListTy getSortedStubs(DenseMap<MCSymbol*, StubValueTy>&);
+
+  /// Return the entries from a DenseMap in a deterministic sorted orer.
+  /// Clears the map.
+  static ExprStubListTy
+  getSortedExprStubs(DenseMap<MCSymbol *, const MCExpr *> &);
 };
 
 //===----------------------------------------------------------------------===//

llvm/include/llvm/CodeGen/MachineModuleInfoImpls.h

@@ -36,6 +36,11 @@ class MachineModuleInfoMachO : public MachineModuleInfoImpl {
   /// bit is true if this GV is external.
   DenseMap<MCSymbol *, StubValueTy> ThreadLocalGVStubs;
 
+  /// Darwin '$auth_ptr' stubs.  The key is the stub symbol, like
+  /// "Lfoo$auth_ptr$ib$12".  The value is the MCExpr representing that
+  /// signed pointer, something like "_foo@AUTH(ib, 12)".
+  DenseMap<MCSymbol *, const MCExpr *> AuthPtrStubs;
+
   virtual void anchor(); // Out of line virtual method.
 
 public:
@@ -51,29 +56,32 @@ class MachineModuleInfoMachO : public MachineModuleInfoImpl {
     return ThreadLocalGVStubs[Sym];
   }
 
+  const MCExpr *&getAuthPtrStubEntry(MCSymbol *Sym) {
+    assert(Sym && "Key cannot be null");
+    return AuthPtrStubs[Sym];
+  }
+
   /// Accessor methods to return the set of stubs in sorted order.
   SymbolListTy GetGVStubList() { return getSortedStubs(GVStubs); }
   SymbolListTy GetThreadLocalGVStubList() {
     return getSortedStubs(ThreadLocalGVStubs);
   }
+
+  ExprStubListTy getAuthGVStubList() {
+    return getSortedExprStubs(AuthPtrStubs);
+  }
 };
 
 /// MachineModuleInfoELF - This is a MachineModuleInfoImpl implementation
 /// for ELF targets.
 class MachineModuleInfoELF : public MachineModuleInfoImpl {
-public:
-  struct AuthStubInfo {
-    const MCExpr *AuthPtrRef;
-  };
-
-private:
   /// GVStubs - These stubs are used to materialize global addresses in PIC
   /// mode.
   DenseMap<MCSymbol *, StubValueTy> GVStubs;
 
   /// AuthPtrStubs - These stubs are used to materialize signed addresses for
   /// extern_weak symbols.
-  DenseMap<MCSymbol *, AuthStubInfo> AuthPtrStubs;
+  DenseMap<MCSymbol *, const MCExpr *> AuthPtrStubs;
 
   virtual void anchor(); // Out of line virtual method.
 
@@ -85,7 +93,7 @@ class MachineModuleInfoELF : public MachineModuleInfoImpl {
     return GVStubs[Sym];
   }
 
-  AuthStubInfo &getAuthPtrStubEntry(MCSymbol *Sym) {
+  const MCExpr *&getAuthPtrStubEntry(MCSymbol *Sym) {
     assert(Sym && "Key cannot be null");
     return AuthPtrStubs[Sym];
   }
@@ -94,10 +102,9 @@ class MachineModuleInfoELF : public MachineModuleInfoImpl {
 
   SymbolListTy GetGVStubList() { return getSortedStubs(GVStubs); }
 
-  using AuthStubPairTy = std::pair<MCSymbol *, AuthStubInfo>;
-  typedef std::vector<AuthStubPairTy> AuthStubListTy;
-
-  AuthStubListTy getAuthGVStubList();
+  ExprStubListTy getAuthGVStubList() {
+    return getSortedExprStubs(AuthPtrStubs);
+  }
 };
 
 /// MachineModuleInfoCOFF - This is a MachineModuleInfoImpl implementation

llvm/lib/CodeGen/MachineModuleInfoImpls.cpp

@@ -43,24 +43,19 @@ MachineModuleInfoImpl::SymbolListTy MachineModuleInfoImpl::getSortedStubs(
   return List;
 }
 
-template <typename MachineModuleInfoTarget>
-static typename MachineModuleInfoTarget::AuthStubListTy getAuthGVStubListHelper(
-    DenseMap<MCSymbol *, typename MachineModuleInfoTarget::AuthStubInfo>
-        &AuthPtrStubs) {
-  typename MachineModuleInfoTarget::AuthStubListTy List(AuthPtrStubs.begin(),
-                                                        AuthPtrStubs.end());
+using ExprStubPairTy = std::pair<MCSymbol *, const MCExpr *>;
+static int SortAuthStubPair(const ExprStubPairTy *LHS,
+                            const ExprStubPairTy *RHS) {
+  return LHS->first->getName().compare(RHS->first->getName());
+}
 
-  if (!List.empty())
-    llvm::sort(List.begin(), List.end(),
-               [](const typename MachineModuleInfoTarget::AuthStubPairTy &LHS,
-                  const typename MachineModuleInfoTarget::AuthStubPairTy &RHS) {
-                 return LHS.first->getName() < RHS.first->getName();
-               });
+MachineModuleInfoImpl::ExprStubListTy MachineModuleInfoImpl::getSortedExprStubs(
+    DenseMap<MCSymbol *, const MCExpr *> &ExprStubs) {
+  MachineModuleInfoImpl::ExprStubListTy List(ExprStubs.begin(),
+                                             ExprStubs.end());
 
-  AuthPtrStubs.clear();
-  return List;
-}
+  array_pod_sort(List.begin(), List.end(), SortAuthStubPair);
 
-MachineModuleInfoELF::AuthStubListTy MachineModuleInfoELF::getAuthGVStubList() {
-  return getAuthGVStubListHelper<MachineModuleInfoELF>(AuthPtrStubs);
+  ExprStubs.clear();
+  return List;
 }

llvm/lib/MC/MCMachOStreamer.cpp

@@ -167,6 +167,10 @@ static bool canGoAfterDWARF(const MCSectionMachO &MSec) {
     return true;
   if (SegName == "__LLVM" && (SecName == "__cg_profile"))
     return true;
+
+  if (SegName == "__DATA" && SecName == "__auth_ptr")
+    return true;
+
   return false;
 }
 

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -848,20 +848,38 @@ void AArch64AsmPrinter::emitHwasanMemaccessSymbols(Module &M) {
   }
 }
 
-template <typename MachineModuleInfoTarget>
-static void emitAuthenticatedPointer(
-    MCStreamer &OutStreamer, MCSymbol *StubLabel,
-    const typename MachineModuleInfoTarget::AuthStubInfo &StubInfo) {
+static void emitAuthenticatedPointer(MCStreamer &OutStreamer,
+                                     MCSymbol *StubLabel,
+                                     const MCExpr *StubAuthPtrRef) {
   // sym$auth_ptr$key$disc:
   OutStreamer.emitLabel(StubLabel);
-  OutStreamer.emitValue(StubInfo.AuthPtrRef, /*size=*/8);
+  OutStreamer.emitValue(StubAuthPtrRef, /*size=*/8);
 }
 
 void AArch64AsmPrinter::emitEndOfAsmFile(Module &M) {
   emitHwasanMemaccessSymbols(M);
 
   const Triple &TT = TM.getTargetTriple();
   if (TT.isOSBinFormatMachO()) {
+    // Output authenticated pointers as indirect symbols, if we have any.
+    MachineModuleInfoMachO &MMIMacho =
+        MMI->getObjFileInfo<MachineModuleInfoMachO>();
+
+    auto Stubs = MMIMacho.getAuthGVStubList();
+
+    if (!Stubs.empty()) {
+      // Switch to the "__auth_ptr" section.
+      OutStreamer->switchSection(
+          OutContext.getMachOSection("__DATA", "__auth_ptr", MachO::S_REGULAR,
+                                     SectionKind::getMetadata()));
+      emitAlignment(Align(8));
+
+      for (const auto &Stub : Stubs)
+        emitAuthenticatedPointer(*OutStreamer, Stub.first, Stub.second);
+
+      OutStreamer->addBlankLine();
+    }
+
     // Funny Darwin hack: This flag tells the linker that no global symbols
     // contain code that falls through to other global symbols (e.g. the obvious
     // implementation of multiple entry points).  If this doesn't occur, the
@@ -882,8 +900,7 @@ void AArch64AsmPrinter::emitEndOfAsmFile(Module &M) {
       emitAlignment(Align(8));
 
       for (const auto &Stub : Stubs)
-        emitAuthenticatedPointer<MachineModuleInfoELF>(*OutStreamer, Stub.first,
-                                                       Stub.second);
+        emitAuthenticatedPointer(*OutStreamer, Stub.first, Stub.second);
 
       OutStreamer->addBlankLine();
     }
@@ -1676,16 +1693,29 @@ void AArch64AsmPrinter::LowerLOADauthptrstatic(const MachineInstr &MI) {
   //
   // Where the $auth_ptr$ symbol is the stub slot containing the signed pointer
   // to symbol.
-  assert(TM.getTargetTriple().isOSBinFormatELF() &&
-         "LOADauthptrstatic is implemented only for ELF");
-  const auto &TLOF =
-      static_cast<const AArch64_ELFTargetObjectFile &>(getObjFileLowering());
-
-  assert(GAOp.getOffset() == 0 &&
-         "non-zero offset for $auth_ptr$ stub slots is not supported");
-  const MCSymbol *GASym = TM.getSymbol(GAOp.getGlobal());
-  MCSymbol *AuthPtrStubSym =
-      TLOF.getAuthPtrSlotSymbol(TM, &MF->getMMI(), GASym, Key, Disc);
+  MCSymbol *AuthPtrStubSym;
+  if (TM.getTargetTriple().isOSBinFormatELF()) {
+    const auto &TLOF =
+        static_cast<const AArch64_ELFTargetObjectFile &>(getObjFileLowering());
+
+    assert(GAOp.getOffset() == 0 &&
+           "non-zero offset for $auth_ptr$ stub slots is not supported");
+    const MCSymbol *GASym = TM.getSymbol(GAOp.getGlobal());
+    AuthPtrStubSym =
+        TLOF.getAuthPtrSlotSymbol(TM, &MF->getMMI(), GASym, Key, Disc);
+  } else {
+    assert(TM.getTargetTriple().isOSBinFormatMachO() &&
+           "LOADauthptrstatic is implemented only for MachO/ELF");
+
+    const auto &TLOF = static_cast<const AArch64_MachoTargetObjectFile &>(
+        getObjFileLowering());
+
+    assert(GAOp.getOffset() == 0 &&
+           "non-zero offset for $auth_ptr$ stub slots is not supported");
+    const MCSymbol *GASym = TM.getSymbol(GAOp.getGlobal());
+    AuthPtrStubSym =
+        TLOF.getAuthPtrSlotSymbol(TM, &MF->getMMI(), GASym, Key, Disc);
+  }
 
   MachineOperand StubMOHi =
       MachineOperand::CreateMCSymbol(AuthPtrStubSym, AArch64II::MO_PAGE);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -9545,8 +9545,7 @@ SDValue AArch64TargetLowering::LowerGlobalTLSAddress(SDValue Op,
 //   Load a signed pointer for symbol 'sym' from a stub slot named
 //   'sym$auth_ptr$key$disc' filled by dynamic linker during relocation
 //   resolving. This usually lowers to adrp+ldr, but also emits an entry into
-//   .data with an
-//   @AUTH relocation. See LowerLOADauthptrstatic.
+//   .data with an @AUTH relocation. See LowerLOADauthptrstatic.
 //
 // All 3 are pseudos that are expand late to longer sequences: this lets us
 // provide integrity guarantees on the to-be-signed intermediate values.
@@ -9599,8 +9598,8 @@ AArch64TargetLowering::LowerPtrAuthGlobalAddress(SDValue Op,
         "constant discriminator in ptrauth global out of range [0, 0xffff]");
 
   // Choosing between 3 lowering alternatives is target-specific.
-  if (!Subtarget->isTargetELF())
-    report_fatal_error("ptrauth global lowering is only implemented for ELF");
+  if (!Subtarget->isTargetELF() && !Subtarget->isTargetMachO())
+    report_fatal_error("ptrauth global lowering only supported on MachO/ELF");
 
   int64_t PtrOffsetC = 0;
   if (Ptr.getOpcode() == ISD::ADD) {

llvm/lib/Target/AArch64/AArch64TargetObjectFile.cpp

@@ -105,15 +105,14 @@ static MCSymbol *getAuthPtrSlotSymbolHelper(
       Twine("$auth_ptr$") + AArch64PACKeyIDToString(Key) + Twine('$') +
       Twine(Discriminator));
 
-  typename MachineModuleInfoTarget::AuthStubInfo &StubInfo =
-      TargetMMI.getAuthPtrStubEntry(StubSym);
+  const MCExpr *&StubAuthPtrRef = TargetMMI.getAuthPtrStubEntry(StubSym);
 
-  if (StubInfo.AuthPtrRef)
+  if (StubAuthPtrRef)
     return StubSym;
 
   const MCExpr *Sym = MCSymbolRefExpr::create(RawSym, Ctx);
 
-  StubInfo.AuthPtrRef =
+  StubAuthPtrRef =
       AArch64AuthMCExpr::create(Sym, Discriminator, Key,
                                 /*HasAddressDiversity=*/false, Ctx);
   return StubSym;
@@ -126,3 +125,11 @@ MCSymbol *AArch64_ELFTargetObjectFile::getAuthPtrSlotSymbol(
   return getAuthPtrSlotSymbolHelper(getContext(), TM, MMI, ELFMMI, RawSym, Key,
                                     Discriminator);
 }
+
+MCSymbol *AArch64_MachoTargetObjectFile::getAuthPtrSlotSymbol(
+    const TargetMachine &TM, MachineModuleInfo *MMI, const MCSymbol *RawSym,
+    AArch64PACKey::ID Key, uint16_t Discriminator) const {
+  auto &MachOMMI = MMI->getObjFileInfo<MachineModuleInfoMachO>();
+  return getAuthPtrSlotSymbolHelper(getContext(), TM, MMI, MachOMMI, RawSym,
+                                    Key, Discriminator);
+}

llvm/lib/Target/AArch64/AArch64TargetObjectFile.h

@@ -60,6 +60,11 @@ class AArch64_MachoTargetObjectFile : public TargetLoweringObjectFileMachO {
 
   void getNameWithPrefix(SmallVectorImpl<char> &OutName, const GlobalValue *GV,
                          const TargetMachine &TM) const override;
+
+  MCSymbol *getAuthPtrSlotSymbol(const TargetMachine &TM,
+                                 MachineModuleInfo *MMI, const MCSymbol *RawSym,
+                                 AArch64PACKey::ID Key,
+                                 uint16_t Discriminator) const;
 };
 
 /// This implementation is used for AArch64 COFF targets.

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

@@ -6636,8 +6636,8 @@ bool AArch64InstructionSelector::selectPtrAuthGlobalValue(
         "constant discriminator in ptrauth global out of range [0, 0xffff]");
 
   // Choosing between 3 lowering alternatives is target-specific.
-  if (!STI.isTargetELF())
-    report_fatal_error("ptrauth global lowering is only implemented for ELF");
+  if (!STI.isTargetELF() && !STI.isTargetMachO())
+    report_fatal_error("ptrauth global lowering only supported on MachO/ELF");
 
   if (!MRI.hasOneDef(Addr))
     return false;