[AArch64][PAC] Support 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.

This also 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.

On MachO, the auth stubs are emitted in __DATA,__auth_ptr.

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$addend$auth_ptr$ib$12".  The value is the MCExpr representing that
+  /// pointer, something like "_foo+addend@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

@@ -161,6 +161,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,39 @@ 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 (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 +901,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();
     }

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.