[JITLink][arm64] Support arm64e JIT'd code (initially enabled for MachO only).

Adds two new JITLink passes to create and populate a pointer-signing function
that can be called via an allocation-action attached to the LinkGraph:

* createEmptyPointerSigningFunction creates a pointer signing function in a
custome section, reserving sufficient space for the signing code. It should
be run as a post-prune pass (to ensure that memory is reserved prior to
allocation).
* lowerPointer64AuthEdgesToSigningFunction pass populates the signing function
by walking the graph, decoding the ptrauth info (encoded in the edge addend) and
writing an instruction sequence to sign all ptrauth fixup locations.

rdar://61956998

Author: lhames

llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h

@@ -31,6 +31,36 @@ enum EdgeKind_aarch64 : Edge::Kind {
   ///
   Pointer64 = Edge::FirstRelocation,
 
+  /// An arm64e authenticated pointer relocation. The addend contains a 64-bit
+  /// struct containing the authentication parameters:
+  ///
+  ///   Addend encoding:
+  ///     int32_t  addend;
+  ///     uint16_t diversityData;
+  ///     uint16_t hasAddressDiversity : 1;
+  ///     uint16_t key : 2;
+  ///     uint16_t zeroes : 12;
+  ///     uint16_t authenticated : 1;
+  ///
+  /// Note: This means that the addend cannot be interpreted as a plain offset
+  ///       prior to lowering.
+  ///
+  /// Authenticated pointer edges cannot be fixed up directly by JITLink as the
+  /// signing keys are held in the executing process. They can be removed from
+  /// the graph by a combination of the createEmptyPointerSigningFunction pass
+  /// (post-prune) and the lowerPointer64AuthEdgesToSigningFunction pass
+  /// (pre-fixup). Together these passes construct a signing function that will
+  /// be run in the executing process to write the signed pointers to the fixup
+  /// locations.
+  ///
+  /// Fixup expression:
+  ///   NONE
+  ///
+  /// Errors:
+  ///   - Failure to handle edges of this kind prior to the fixup phase will
+  ///     result in an unsupported error during the fixup phase.
+  Pointer64Authenticated,
+
   /// A plain 32-bit pointer value relocation.
   ///
   /// Fixup expression:
@@ -832,6 +862,29 @@ class PLTTableManager : public TableManager<PLTTableManager> {
   Section *StubsSection = nullptr;
 };
 
+/// Returns the name of the pointer signing function section.
+const char *getPointerSigningFunctionSectionName();
+
+/// Creates a pointer signing function section, block, and symbol to reserve
+/// space for a signing function for this LinkGraph. Clients should insert this
+/// pass in the post-prune phase, and add the paired
+/// lowerPointer64AuthEdgesToSigningFunction pass to the pre-fixup phase.
+///
+/// No new Pointer64Auth edges can be inserted into the graph between when this
+/// pass is run and when the pass below runs (since there will not be sufficient
+/// space reserved in the signing function to write the signing code for them).
+Error createEmptyPointerSigningFunction(LinkGraph &G);
+
+/// Given a LinkGraph containing Pointer64Authenticated edges, transform those
+/// edges to Pointer64 and add signing code to the pointer signing function
+/// (which must already have been created by the
+/// createEmptyPointerSigningFunction pass above).
+///
+/// This function will add a $__ptrauth_sign section with finalization-lifetime
+/// containing an anonymous function that will sign all pointers in the graph.
+/// An allocation action will be added to run this function during finalization.
+Error lowerPointer64AuthEdgesToSigningFunction(LinkGraph &G);
+
 } // namespace aarch64
 } // namespace jitlink
 } // namespace llvm

llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp

@@ -28,8 +28,8 @@ class MachOLinkGraphBuilder_arm64 : public MachOLinkGraphBuilder {
 public:
   MachOLinkGraphBuilder_arm64(const object::MachOObjectFile &Obj,
                               SubtargetFeatures Features)
-      : MachOLinkGraphBuilder(Obj, Triple("arm64-apple-darwin"),
-                              std::move(Features), aarch64::getEdgeKindName),
+      : MachOLinkGraphBuilder(Obj, getObjectTriple(Obj), std::move(Features),
+                              aarch64::getEdgeKindName),
         NumSymbols(Obj.getSymtabLoadCommand().nsyms) {}
 
 private:
@@ -38,6 +38,7 @@ class MachOLinkGraphBuilder_arm64 : public MachOLinkGraphBuilder {
     MachOPointer32,
     MachOPointer64,
     MachOPointer64Anon,
+    MachOPointer64Authenticated,
     MachOPage21,
     MachOPageOffset12,
     MachOGOTPage21,
@@ -53,6 +54,18 @@ class MachOLinkGraphBuilder_arm64 : public MachOLinkGraphBuilder {
     MachONegDelta64,
   };
 
+  static Triple getObjectTriple(const object::MachOObjectFile &Obj) {
+    // Get the CPU sub-type from the header.
+    // jitLink_MachO should already have validated that the buffer is big enough
+    // to cover a mach_header64 so this is safe.
+    uint32_t CPUSubType =
+        *(const support::ulittle32_t *)(Obj.getData().data() + 8);
+    CPUSubType &= ~MachO::CPU_SUBTYPE_MASK;
+    if (CPUSubType == MachO::CPU_SUBTYPE_ARM64E)
+      return Triple("arm64e-apple-darwin");
+    return Triple("arm64-apple-darwin");
+  }
+
   static Expected<MachOARM64RelocationKind>
   getRelocationKind(const MachO::relocation_info &RI) {
     switch (RI.r_type) {
@@ -103,6 +116,10 @@ class MachOLinkGraphBuilder_arm64 : public MachOLinkGraphBuilder {
       if (!RI.r_pcrel && !RI.r_extern && RI.r_length == 2)
         return MachOPairedAddend;
       break;
+    case MachO::ARM64_RELOC_AUTHENTICATED_POINTER:
+      if (!RI.r_pcrel && RI.r_extern && RI.r_length == 3)
+        return MachOPointer64Authenticated;
+      break;
     case MachO::ARM64_RELOC_TLVP_LOAD_PAGE21:
       if (RI.r_pcrel && RI.r_extern && RI.r_length == 2)
         return MachOTLVPage21;
@@ -366,12 +383,15 @@ class MachOLinkGraphBuilder_arm64 : public MachOLinkGraphBuilder {
           Kind = aarch64::Pointer32;
           break;
         case MachOPointer64:
+        case MachOPointer64Authenticated:
           if (auto TargetSymbolOrErr = findSymbolByIndex(RI.r_symbolnum))
             TargetSymbol = TargetSymbolOrErr->GraphSymbol;
           else
             return TargetSymbolOrErr.takeError();
           Addend = *(const ulittle64_t *)FixupContent;
-          Kind = aarch64::Pointer64;
+          Kind = *MachORelocKind == MachOPointer64
+                     ? aarch64::Pointer64
+                     : aarch64::Pointer64Authenticated;
           break;
         case MachOPointer64Anon: {
           orc::ExecutorAddr TargetAddress(*(const ulittle64_t *)FixupContent);
@@ -493,6 +513,8 @@ class MachOLinkGraphBuilder_arm64 : public MachOLinkGraphBuilder {
       return "MachOPointer64";
     case MachOPointer64Anon:
       return "MachOPointer64Anon";
+    case MachOPointer64Authenticated:
+      return "MachOPointer64Authenticated";
     case MachOPage21:
       return "MachOPage21";
     case MachOPageOffset12:
@@ -601,6 +623,14 @@ void link_MachO_arm64(std::unique_ptr<LinkGraph> G,
 
     // Add an in-place GOT/Stubs pass.
     Config.PostPrunePasses.push_back(buildTables_MachO_arm64);
+
+    // If this is an arm64e graph then add pointer signing passes.
+    if (G->getTargetTriple().isArm64e()) {
+      Config.PostPrunePasses.push_back(
+          aarch64::createEmptyPointerSigningFunction);
+      Config.PreFixupPasses.push_back(
+          aarch64::lowerPointer64AuthEdgesToSigningFunction);
+    }
   }
 
   if (auto Err = Ctx->modifyPassConfig(*G, Config))