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

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))