[ExternalGenericMetadataBuilder] Support ARM64e fixups.

Add ptrauth attributes to the witness tables and descriptor pointers, and have the test support ARM64e.

rdar://122660111

Author: mikeash

include/swift/Runtime/GenericMetadataBuilder.h

@@ -289,9 +289,13 @@ class GenericMetadataBuilder {
           extraDataPattern->SizeInWords, extraDataPattern->OffsetInWords);
       auto patternPointers =
           patternBuffer.resolvePointer(&extraDataPattern->Pattern);
+      if (!patternPointers)
+        return *patternPointers.getError();
       for (unsigned i = 0; i < extraDataPattern->SizeInWords; i++) {
         auto patternPointer =
             patternPointers->resolvePointer(&patternPointers->ptr[i]);
+        if (!patternPointer)
+          return *patternPointer.getError();
         data.writePointer(
             &metadataExtraData[i + extraDataPattern->OffsetInWords],
             patternPointer->template cast<const StoredPointer>());
@@ -303,6 +307,8 @@ class GenericMetadataBuilder {
     // necessary.
     auto valueWitnesses =
         patternBuffer.resolvePointer(&pattern->ValueWitnesses);
+    if (!valueWitnesses)
+      return *valueWitnesses.getError();
     METADATA_BUILDER_LOG("Setting initial value witnesses");
     data.writePointer(&fullMetadata->ValueWitnesses, *valueWitnesses);
 

lib/ExternalGenericMetadataBuilder/ExternalGenericMetadataBuilder.cpp

@@ -117,12 +117,62 @@ struct TypedExternal {
   template <typename T, bool Nullable = true, typename Offset = int32_t>
   using CompactFunctionPointer = TypedInteger<T, int32_t, Nullable>;
 
+  template <unsigned discriminator>
+  struct ValueWitnessFunctionPointer {
+    StoredPointer pointer;
+  };
+
   StoredPointer
   getStrippedSignedPointer(const StoredSignedPointer pointer) const {
     return swift_ptrauth_strip(pointer);
   }
 };
 
+using ExternalRuntime32 =
+    swift::TypedExternal<swift::WithObjCInterop<swift::RuntimeTarget<4>>>;
+using ExternalRuntime64 =
+    swift::TypedExternal<swift::WithObjCInterop<swift::RuntimeTarget<8>>>;
+
+// Declare a specialized version of the value witness types that use a wrapper
+// on the functions that captures the ptrauth discriminator.
+template <>
+class TargetValueWitnessTypes<ExternalRuntime64> {
+public:
+  using StoredPointer = typename ExternalRuntime64::StoredPointer;
+
+#define WANT_ALL_VALUE_WITNESSES
+#define DATA_VALUE_WITNESS(lowerId, upperId, type)
+#define FUNCTION_VALUE_WITNESS(lowerId, upperId, returnType, paramTypes)       \
+  typedef ExternalRuntime64::ValueWitnessFunctionPointer<                      \
+      SpecialPointerAuthDiscriminators::upperId>                               \
+      lowerId;
+#define MUTABLE_VALUE_TYPE TargetPointer<ExternalRuntime64, OpaqueValue>
+#define IMMUTABLE_VALUE_TYPE ConstTargetPointer<ExternalRuntime64, OpaqueValue>
+#define MUTABLE_BUFFER_TYPE TargetPointer<ExternalRuntime64, ValueBuffer>
+#define IMMUTABLE_BUFFER_TYPE ConstTargetPointer<ExternalRuntime64, ValueBuffer>
+#define TYPE_TYPE ConstTargetPointer<ExternalRuntime64, Metadata>
+#define SIZE_TYPE StoredSize
+#define INT_TYPE int
+#define UINT_TYPE unsigned
+#define VOID_TYPE void
+#include "swift/ABI/ValueWitness.def"
+
+  // Handle the data witnesses explicitly so we can use more specific
+  // types for the flags enums.
+  typedef size_t size;
+  typedef size_t stride;
+  typedef TargetValueWitnessFlags<typename ExternalRuntime64::StoredSize> flags;
+  typedef uint32_t extraInhabitantCount;
+};
+
+enum class PtrauthKey : int8_t {
+  None = -1,
+  IA = 0,
+  IB = 1,
+  DA = 2,
+  DB = 3,
+};
+
 template <typename T>
 inline MetadataKind getEnumeratedMetadataKind(T kind,
                                               decltype(kind.value) dummy = 0) {
@@ -188,6 +238,26 @@ class ReaderWriter;
 
 template <typename Runtime>
 class ExternalGenericMetadataBuilderContext {
+  // Structs that provide the ptrauth info for pointers to specific tyeps.
+  template <typename Target>
+  struct PtrauthInfo;
+
+  template <>
+  struct PtrauthInfo<const TargetValueTypeDescriptor<Runtime>> {
+    static constexpr PtrauthKey key = PtrauthKey::DA;
+    static constexpr bool addressDiversified = true;
+    static constexpr unsigned discriminator =
+        SpecialPointerAuthDiscriminators::TypeDescriptor;
+  };
+
+  template <>
+  struct PtrauthInfo<const TargetValueWitnessTable<Runtime>> {
+    static constexpr PtrauthKey key = PtrauthKey::DA;
+    static constexpr bool addressDiversified = true;
+    static constexpr unsigned discriminator =
+        SpecialPointerAuthDiscriminators::ValueWitnessTable;
+  };
+
   struct Atom;
 
   struct FileTarget {
@@ -207,6 +277,10 @@ class ExternalGenericMetadataBuilderContext {
     unsigned offset;
     unsigned size;
 
+    PtrauthKey ptrauthKey;
+    bool addressDiversified;
+    unsigned discriminator;
+
     std::variant<FileTarget, AtomTarget> fileOrAtom;
   };
 
@@ -512,7 +586,10 @@ class ExternalGenericMetadataBuilderContext {
     }
 
     template <typename U>
-    void writePointerImpl(void *where, Buffer<U> value) {
+    void writePointerImpl(void *where, Buffer<U> value,
+                          PtrauthKey ptrauthKey = PtrauthKey::None,
+                          bool addressDiversified = true,
+                          unsigned discriminator = 0) {
       if (!value) {
         memset(where, 0, sizeof(StoredPointer));
         return;
@@ -530,6 +607,10 @@ class ExternalGenericMetadataBuilderContext {
 
       target.size = sizeof(StoredPointer);
 
+      target.ptrauthKey = ptrauthKey;
+      target.addressDiversified = addressDiversified;
+      target.discriminator = discriminator;
+
       if (auto *file = value.file) {
         auto contents = value.section.getContents();
         if (!contents) {
@@ -624,7 +705,9 @@ class ExternalGenericMetadataBuilderContext {
       checkPtr(where);
       where->SignedValue.value = ~(uintptr_t)value.ptr;
 
-      writePointerImpl(where, value);
+      writePointerImpl(where, value, PtrauthInfo<U>::key,
+                       PtrauthInfo<U>::addressDiversified,
+                       PtrauthInfo<U>::discriminator);
     }
 
     template <typename U>
@@ -638,6 +721,15 @@ class ExternalGenericMetadataBuilderContext {
     void writeFunctionPointer(void *where, Buffer<const char> target) {
       writePointer(reinterpret_cast<StoredPointer *>(where), target);
     }
+
+    template <unsigned discriminator>
+    void writeFunctionPointer(
+        ExternalRuntime64::ValueWitnessFunctionPointer<discriminator> *where,
+        Buffer<const char> target) {
+      checkPtr(where);
+
+      writePointerImpl(where, target, PtrauthKey::IA, true, discriminator);
+    }
   };
 
   ExternalGenericMetadataBuilderContext() {
@@ -660,7 +752,11 @@ class ExternalGenericMetadataBuilderContext {
     return allocate<T>(count * sizeof(T));
   }
 
-  void setArch(const char *arch) { this->arch = arch; }
+  void setArch(const char *arch) {
+    this->arch = arch;
+    this->usePtrauth = this->arch == "arm64e";
+  }
+
   void setNamesToBuild(const std::vector<std::string> &names) {
     this->mangledNamesToBuild = names;
   }
@@ -723,6 +819,9 @@ class ExternalGenericMetadataBuilderContext {
   // The architecture being targeted.
   std::string arch;
 
+  // Does this target use pointer authentication?
+  bool usePtrauth = false;
+
   // The readerWriter and builder helper objects.
   std::unique_ptr<ReaderWriter<Runtime>> readerWriter;
   std::unique_ptr<Builder> builder;
@@ -2127,6 +2226,14 @@ void ExternalGenericMetadataBuilderContext<Runtime>::writeAtomContentsJSON(
         J.attribute("addend", atomTarget.offset);
         J.attribute("kind", ptrTargetKind);
       }
+
+      if (usePtrauth && targetsCursor->ptrauthKey != PtrauthKey::None) {
+        J.attributeObject("authPtr", [&] {
+          J.attribute("key", static_cast<uint8_t>(targetsCursor->ptrauthKey));
+          J.attribute("addr", targetsCursor->addressDiversified);
+          J.attribute("diversity", targetsCursor->discriminator);
+        });
+      }
     });
 
     bufferCursor = targetsCursor->offset + targetsCursor->size;
@@ -2286,16 +2393,11 @@ BuilderErrorOr<unsigned> getPointerWidth(std::string arch) {
 
 } // namespace swift
 
-using ExternalRuntime32 =
-    swift::TypedExternal<swift::WithObjCInterop<swift::RuntimeTarget<4>>>;
-using ExternalRuntime64 =
-    swift::TypedExternal<swift::WithObjCInterop<swift::RuntimeTarget<8>>>;
-
 struct SwiftExternalMetadataBuilder {
   using Builder32 =
-      swift::ExternalGenericMetadataBuilderContext<ExternalRuntime32>;
+      swift::ExternalGenericMetadataBuilderContext<swift::ExternalRuntime32>;
   using Builder64 =
-      swift::ExternalGenericMetadataBuilderContext<ExternalRuntime64>;
+      swift::ExternalGenericMetadataBuilderContext<swift::ExternalRuntime64>;
 
   std::variant<Builder32, Builder64> context;
 

stdlib/public/runtime/GenericMetadataBuilder.cpp

@@ -485,13 +485,6 @@ bool swift::compareGenericMetadata(const Metadata *original,
         equal = false;
       }
 
-      if (memcmp(original, newMetadata, genericArgumentsStart)) {
-        validationLog(
-            true,
-            "Metadatas do not match in the part before generic arguments");
-        equal = false;
-      }
-
       for (unsigned i = 0; i < genericContextHeader.NumKeyArguments; i++) {
         auto *originalArg =
             originalDescriptor->getGenericArguments(original)[i];

stdlib/public/runtime/LibPrespecialized.cpp

@@ -36,8 +36,11 @@ static const LibPrespecializedData<InProcess> *findLibPrespecialized() {
   auto path = runtime::environment::SWIFT_DEBUG_LIB_PRESPECIALIZED_PATH();
   if (path && path[0]) {
     void *handle = dlopen(path, RTLD_LAZY);
-    if (!handle)
+    if (!handle) {
+      swift::warning(0, "Failed to load prespecializations library: %s\n",
+                     dlerror());
       return nullptr;
+    }
 
     dataPtr = dlsym(handle, LIB_PRESPECIALIZED_TOP_LEVEL_SYMBOL_NAME);
   }

test/ExternalGenericMetadataBuilder/Inputs/json2c.swift

@@ -7,9 +7,16 @@ enum ParseError: Error {
   case invalidHex(Substring)
 }
 
+struct Ptrauth: Decodable {
+  var key: Int8
+  var addr: Bool
+  var diversity: UInt
+}
+
 struct Fixup: Decodable {
   var target: String
   var addend: Int64
+  var authPtr: Ptrauth?
 }
 
 enum AtomContent: Decodable {
@@ -99,8 +106,18 @@ func printDeclarations(_ document: Document) {
       switch content {
         case .bytes(let bytes):
           print("  uint8_t \(name)[\(bytes.count)];")
-        case .pointer(_):
-          print("  const char *\(name);")
+        case .pointer(let fixup):
+          let ptrauthQualifier: String
+          if let ptrauth = fixup.authPtr {
+            ptrauthQualifier = """
+              __ptrauth(\(ptrauth.key),
+                        \(ptrauth.addr ? 1 : 0),
+                        \(ptrauth.diversity))
+            """
+          } else {
+            ptrauthQualifier = ""
+          }
+          print("  const char *\(ptrauthQualifier) \(name);")
       }
     }
     print("};")

test/ExternalGenericMetadataBuilder/VerifyExternalMetadata.swift

@@ -1,17 +1,18 @@
 // RUN: %empty-directory(%t)
-// RUN: %target-build-swift -Xfrontend -disable-availability-checking -I %swift-lib-dir -I %swift_src_root/lib/ExternalGenericMetadataBuilder -L%swift-lib-dir -lswiftGenericMetadataBuilder -enable-experimental-feature Extern %s -o %t/VerifyExternalMetadata
+// RUN: %target-build-swift -Xfrontend -disable-availability-checking -I %swift-lib-dir -I %swift_src_root/lib/ExternalGenericMetadataBuilder -enable-experimental-feature Extern %s -o %t/VerifyExternalMetadata
 // RUN: %target-codesign %t/VerifyExternalMetadata
 //
-// RUN: %target-build-swift -I %swift-lib-dir -I %swift_src_root/lib/ExternalGenericMetadataBuilder -L%swift-lib-dir -lswiftGenericMetadataBuilder -enable-experimental-feature Extern %S/Inputs/buildMetadataJSON.swift -o %t/buildMetadataJSON
+// RUN: %host-build-swift -Xfrontend -disable-availability-checking -I %swift-lib-dir -I %swift_src_root/lib/ExternalGenericMetadataBuilder -L%swift-lib-dir -lswiftGenericMetadataBuilder -Xlinker -rpath -Xlinker %swift-lib-dir -enable-experimental-feature Extern %S/Inputs/buildMetadataJSON.swift -o %t/buildMetadataJSON
+// no: %target-build-swift -I %swift-lib-dir -I %swift_src_root/lib/ExternalGenericMetadataBuilder -L%swift-lib-dir -lswiftGenericMetadataBuilder -enable-experimental-feature Extern %S/Inputs/buildMetadataJSON.swift -o %t/buildMetadataJSON
 // RUN: %target-codesign %t/buildMetadataJSON
 //
-// RUN: %target-build-swift %S/Inputs/json2c.swift -o %t/json2c
+// RUN: %target-build-swift -Xfrontend -disable-availability-checking %S/Inputs/json2c.swift -o %t/json2c
 // RUN: %target-codesign %t/json2c
 //
 // RUN: %target-run %t/VerifyExternalMetadata getJSON > %t/names.json
-// RUN: %target-run %t/buildMetadataJSON arm64 %t/VerifyExternalMetadata %stdlib_dir/libswiftCore.dylib < %t/names.json > %t/libswiftPrespecialized.json
+// RUN: %target-run %t/buildMetadataJSON %target-arch %t/VerifyExternalMetadata %stdlib_dir/libswiftCore.dylib < %t/names.json > %t/libswiftPrespecialized.json
 // RUN: %target-run %t/json2c %t/libswiftPrespecialized.json > %t/libswiftPrespecialized.c
-// RUN: %clang -isysroot %sdk -bundle %t/libswiftPrespecialized.c -L%stdlib_dir -lswiftCore -bundle_loader %t/VerifyExternalMetadata -o %t/libswiftPrespecialized.bundle
+// RUN: %clang -isysroot %sdk -target %target-triple -bundle %t/libswiftPrespecialized.c -L%stdlib_dir -lswiftCore -bundle_loader %t/VerifyExternalMetadata -o %t/libswiftPrespecialized.bundle
 //
 //
 // RUN: env SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING=y SWIFT_DEBUG_LIB_PRESPECIALIZED_PATH=%t/libswiftPrespecialized.bundle %target-run %t/VerifyExternalMetadata