Merge pull request #40611 from apple/jgrynspan/coff-elf-double-registration-fix

Fix an issue on COFF/ELF targets where the runtime would register each loaded image twice (at least early on.)

Author: grynspan

stdlib/public/SwiftShims/MetadataSections.h

@@ -39,7 +39,25 @@ typedef struct MetadataSectionRange {
 /// Identifies the address space ranges for the Swift metadata required by the Swift runtime.
 struct MetadataSections {
   __swift_uintptr_t version;
-  __swift_uintptr_t reserved;
+
+  /// The base address of the image where this metadata section was defined, as
+  /// reported when the section was registered with the Swift runtime.
+  ///
+  /// The value of this field is equivalent to the value of
+  /// \c SymbolInfo::baseAddress as returned from \c lookupSymbol() for a symbol
+  /// in the image that contains these sections.
+  ///
+  /// For Mach-O images, set this field to \c __dso_handle (i.e. the Mach header
+  /// for the image.) For ELF images, set it to \c __dso_handle (the runtime
+  /// will adjust it to the start of the ELF image when the image is loaded.)
+  /// For COFF images, set this field to \c __ImageBase.
+  ///
+  /// For platforms that have a single statically-linked image or no dynamic
+  /// loader (i.e. no equivalent of \c __dso_handle or \c __ImageBase), this
+  /// field is ignored and should be set to \c nullptr.
+  ///
+  /// \sa swift_addNewDSOImage()
+  const void *baseAddress;
 
   /// `next` and `prev` are used by the runtime to construct a
   /// circularly doubly linked list to quickly iterate over the metadata

stdlib/public/runtime/AccessibleFunction.cpp

@@ -98,8 +98,8 @@ static void _registerAccessibleFunctions(AccessibleFunctionsState &C,
   C.SectionsToScan.push_back(section);
 }
 
-void swift::addImageAccessibleFunctionsBlockCallbackUnsafe(const void *functions,
-                                                           uintptr_t size) {
+void swift::addImageAccessibleFunctionsBlockCallbackUnsafe(
+  const void *baseAddress, const void *functions, uintptr_t size) {
   assert(
       size % sizeof(AccessibleFunctionRecord) == 0 &&
       "accessible function section not a multiple of AccessibleFunctionRecord");
@@ -108,10 +108,10 @@ void swift::addImageAccessibleFunctionsBlockCallbackUnsafe(const void *functions
   _registerAccessibleFunctions(C, AccessibleFunctionsSection{functions, size});
 }
 
-void swift::addImageAccessibleFunctionsBlockCallback(const void *functions,
-                                                     uintptr_t size) {
+void swift::addImageAccessibleFunctionsBlockCallback(
+  const void *baseAddress, const void *functions, uintptr_t size) {
   Functions.get();
-  addImageAccessibleFunctionsBlockCallbackUnsafe(functions, size);
+  addImageAccessibleFunctionsBlockCallbackUnsafe(baseAddress, functions, size);
 }
 
 static const AccessibleFunctionRecord *

stdlib/public/runtime/ImageInspection.h

@@ -77,22 +77,31 @@ void initializeDynamicReplacementLookup();
 void initializeAccessibleFunctionsLookup();
 
 // Callbacks to register metadata from an image to the runtime.
-void addImageProtocolsBlockCallback(const void *start, uintptr_t size);
-void addImageProtocolsBlockCallbackUnsafe(const void *start, uintptr_t size);
-void addImageProtocolConformanceBlockCallback(const void *start,
+void addImageProtocolsBlockCallback(const void *baseAddress,
+                                    const void *start, uintptr_t size);
+void addImageProtocolsBlockCallbackUnsafe(const void *baseAddress,
+                                          const void *start, uintptr_t size);
+void addImageProtocolConformanceBlockCallback(const void *baseAddress,
+                                              const void *start,
                                               uintptr_t size);
-void addImageProtocolConformanceBlockCallbackUnsafe(const void *start,
+void addImageProtocolConformanceBlockCallbackUnsafe(const void *baseAddress,
+                                                    const void *start,
                                                     uintptr_t size);
-void addImageTypeMetadataRecordBlockCallback(const void *start,
+void addImageTypeMetadataRecordBlockCallback(const void *baseAddress,
+                                             const void *start,
                                              uintptr_t size);
-void addImageTypeMetadataRecordBlockCallbackUnsafe(const void *start,
+void addImageTypeMetadataRecordBlockCallbackUnsafe(const void *baseAddress,
+                                                   const void *start,
                                                    uintptr_t size);
-void addImageDynamicReplacementBlockCallback(const void *start, uintptr_t size,
+void addImageDynamicReplacementBlockCallback(const void *baseAddress,
+                                             const void *start, uintptr_t size,
                                              const void *start2,
                                              uintptr_t size2);
-void addImageAccessibleFunctionsBlockCallback(const void *start,
+void addImageAccessibleFunctionsBlockCallback(const void *baseAddress,
+                                              const void *start,
                                               uintptr_t size);
-void addImageAccessibleFunctionsBlockCallbackUnsafe(const void *start,
+void addImageAccessibleFunctionsBlockCallbackUnsafe(const void *baseAddress,
+                                                    const void *start,
                                                     uintptr_t size);
 
 int lookupSymbol(const void *address, SymbolInfo *info);

stdlib/public/runtime/ImageInspectionCommon.cpp

@@ -28,9 +28,10 @@
 
 namespace swift {
 
+#ifndef NDEBUG
 static swift::MetadataSections *registered = nullptr;
 
-void record(swift::MetadataSections *sections) {
+static void record(swift::MetadataSections *sections) {
   if (registered == nullptr) {
     registered = sections;
     sections->next = sections->prev = sections;
@@ -41,34 +42,67 @@ void record(swift::MetadataSections *sections) {
     registered->prev = sections;
   }
 }
+#endif
+
+static const void *
+getMetadataSectionBaseAddress(swift::MetadataSections *sections) {
+  // If the base address was not set by the caller of swift_addNewDSOImage()
+  // then we can assume that the caller was built against an older version of
+  // the runtime that did not capture a value for this field. Currently nothing
+  // is actively using the image's base address outside of tests that are built
+  // with the runtime/stdlib, so there's no need to try to fix up the value. If
+  // something in the runtime starts using it, we will want to either:
+  // 1. Resolve the address from a known-good address like swift5_protocols when
+  //    the image is first loaded (in this function);
+  // 1. Resolve the address from a known-good address like swift5_protocols when
+  //    the address is first used (and atomically swap the address back so we
+  //    don't incur the cost of lookupSymbol() each time we need it; or
+  // 3. Introduce an ABI-breaking change so that all binaries are rebuilt and
+  //    start supplying a value for this field.
+
+#if defined(__ELF__)
+  // If the base address was set but the image is an ELF image, it is going to
+  // be __dso_handle which is not the value we expect (Dl_info::dli_fbase), so
+  // we need to fix it up.
+  if (auto baseAddress = sections->baseAddress) {
+    swift::SymbolInfo symbolInfo;
+    if (lookupSymbol(baseAddress, &symbolInfo) && symbolInfo.baseAddress) {
+      sections->baseAddress = symbolInfo.baseAddress;
+    }
+  }
+#endif
+
+  return sections->baseAddress;
+}
 }
 
 SWIFT_RUNTIME_EXPORT
-void swift_addNewDSOImage(const void *addr) {
-  // We cast off the const in order to update the linked list
-  // data structure. This is safe to do since we don't touch 
-  // any other fields.
-  swift::MetadataSections *sections =
-      static_cast<swift::MetadataSections *>(const_cast<void *>(addr));
-
+void swift_addNewDSOImage(swift::MetadataSections *sections) {
+#ifndef NDEBUG
   record(sections);
+#endif
+
+  auto baseAddress = swift::getMetadataSectionBaseAddress(sections);
 
   const auto &protocols_section = sections->swift5_protocols;
   const void *protocols = reinterpret_cast<void *>(protocols_section.start);
   if (protocols_section.length)
-    swift::addImageProtocolsBlockCallback(protocols, protocols_section.length);
+    swift::addImageProtocolsBlockCallback(baseAddress,
+                                          protocols, protocols_section.length);
 
   const auto &protocol_conformances = sections->swift5_protocol_conformances;
   const void *conformances =
       reinterpret_cast<void *>(protocol_conformances.start);
   if (protocol_conformances.length)
-    swift::addImageProtocolConformanceBlockCallback(conformances,
+    swift::addImageProtocolConformanceBlockCallback(baseAddress, conformances,
                                              protocol_conformances.length);
 
   const auto &type_metadata = sections->swift5_type_metadata;
   const void *metadata = reinterpret_cast<void *>(type_metadata.start);
   if (type_metadata.length)
-    swift::addImageTypeMetadataRecordBlockCallback(metadata, type_metadata.length);
+    swift::addImageTypeMetadataRecordBlockCallback(baseAddress,
+                                                   metadata,
+                                                   type_metadata.length);
 
   const auto &dynamic_replacements = sections->swift5_replace;
   const auto *replacements =
@@ -77,7 +111,7 @@ void swift_addNewDSOImage(const void *addr) {
     const auto &dynamic_replacements_some = sections->swift5_replac2;
     const auto *replacements_some =
       reinterpret_cast<void *>(dynamic_replacements_some.start);
-    swift::addImageDynamicReplacementBlockCallback(
+    swift::addImageDynamicReplacementBlockCallback(baseAddress,
         replacements, dynamic_replacements.length, replacements_some,
         dynamic_replacements_some.length);
   }
@@ -87,70 +121,22 @@ void swift_addNewDSOImage(const void *addr) {
       reinterpret_cast<void *>(accessible_funcs_section.start);
   if (accessible_funcs_section.length)
     swift::addImageAccessibleFunctionsBlockCallback(
-        functions, accessible_funcs_section.length);
+        baseAddress, functions, accessible_funcs_section.length);
 }
 
 void swift::initializeProtocolLookup() {
-  const swift::MetadataSections *sections = registered;
-  while (true) {
-    const swift::MetadataSectionRange &protocols =
-      sections->swift5_protocols;
-    if (protocols.length)
-      addImageProtocolsBlockCallbackUnsafe(
-          reinterpret_cast<void *>(protocols.start), protocols.length);
-
-    if (sections->next == registered)
-      break;
-    sections = sections->next;
-  }
 }
 
 void swift::initializeProtocolConformanceLookup() {
-  const swift::MetadataSections *sections = registered;
-  while (true) {
-    const swift::MetadataSectionRange &conformances =
-        sections->swift5_protocol_conformances;
-    if (conformances.length)
-      addImageProtocolConformanceBlockCallbackUnsafe(
-          reinterpret_cast<void *>(conformances.start), conformances.length);
-
-    if (sections->next == registered)
-      break;
-    sections = sections->next;
-  }
 }
 
 void swift::initializeTypeMetadataRecordLookup() {
-  const swift::MetadataSections *sections = registered;
-  while (true) {
-    const swift::MetadataSectionRange &type_metadata =
-        sections->swift5_type_metadata;
-    if (type_metadata.length)
-      addImageTypeMetadataRecordBlockCallbackUnsafe(
-          reinterpret_cast<void *>(type_metadata.start), type_metadata.length);
-
-    if (sections->next == registered)
-      break;
-    sections = sections->next;
-  }
 }
 
 void swift::initializeDynamicReplacementLookup() {
 }
 
 void swift::initializeAccessibleFunctionsLookup() {
-  const swift::MetadataSections *sections = registered;
-  while (true) {
-    const swift::MetadataSectionRange &functions =
-        sections->swift5_accessible_functions;
-    if (functions.length)
-      addImageAccessibleFunctionsBlockCallbackUnsafe(
-          reinterpret_cast<void *>(functions.start), functions.length);
-
-    if (sections->next == registered)
-      break;
-    sections = sections->next;
-  }
 }
 
 #ifndef NDEBUG
@@ -173,16 +159,35 @@ const swift::MetadataSections *swift_getMetadataSection(size_t index) {
 }
 
 SWIFT_RUNTIME_EXPORT
-const char *swift_getMetadataSectionName(void *metadata_section) {
+const char *
+swift_getMetadataSectionName(const swift::MetadataSections *section) {
   swift::SymbolInfo info;
-  if (lookupSymbol(metadata_section, &info)) {
+  if (lookupSymbol(section, &info)) {
     if (info.fileName) {
       return info.fileName;
     }
   }
   return "";
 }
 
+SWIFT_RUNTIME_EXPORT
+void swift_getMetadataSectionBaseAddress(const swift::MetadataSections *section,
+                                         void const **out_actual,
+                                         void const **out_expected) {
+  *out_actual = nullptr;
+  *out_expected = section->baseAddress;
+
+  swift::SymbolInfo info;
+  if (lookupSymbol(section, &info)) {
+    *out_actual = info.baseAddress;
+    if (info.fileName) {
+      return info.fileName;
+    }
+  }
+  return "";
+
+}
+
 SWIFT_RUNTIME_EXPORT
 size_t swift_getMetadataSectionCount() {
   if (swift::registered == nullptr)

stdlib/public/runtime/ImageInspectionCommon.h

@@ -52,22 +52,35 @@
 
 namespace swift {
 struct MetadataSections;
-static constexpr const uintptr_t CurrentSectionMetadataVersion = 1;
+static constexpr const uintptr_t CurrentSectionMetadataVersion = 2;
 }
 
 struct SectionInfo {
   uint64_t size;
   const char *data;
 };
 
-// Called by injected constructors when a dynamic library is loaded.
+/// Called by injected constructors when a dynamic library is loaded.
+///
+/// \param sections A structure describing the metadata sections in the
+///     newly-loaded image.
+///
+/// \warning The runtime keeps a reference to \a sections and may mutate it, so
+///   it \em must be mutable and long-lived (that is, statically or dynamically
+///   allocated.) The effect of passing a pointer to a local value is undefined.
 SWIFT_RUNTIME_EXPORT
-void swift_addNewDSOImage(const void *addr);
+void swift_addNewDSOImage(struct swift::MetadataSections *sections);
 
 #ifndef NDEBUG
 
 SWIFT_RUNTIME_EXPORT
-const char *swift_getMetadataSectionName(void *metadata_section);
+const char *
+swift_getMetadataSectionName(const struct swift::MetadataSections *section);
+
+SWIFT_RUNTIME_EXPORT
+void swift_getMetadataSectionBaseAddress(
+  const struct swift::MetadataSections *section,
+  void const **out_actual, void const **out_expected);
 
 SWIFT_RUNTIME_EXPORT
 size_t swift_getMetadataSectionCount();

stdlib/public/runtime/ImageInspectionMachO.cpp

@@ -54,7 +54,8 @@ using mach_header_platform = mach_header;
 #endif
 
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
-         void CONSUME_BLOCK(const void *start, uintptr_t size)>
+         void CONSUME_BLOCK(const void *baseAddress,
+                            const void *start, uintptr_t size)>
 void addImageCallback(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
   assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
@@ -70,17 +71,19 @@ void addImageCallback(const mach_header *mh) {
   if (!section)
     return;
 
-  CONSUME_BLOCK(section, size);
+  CONSUME_BLOCK(mh, section, size);
 }
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
-         void CONSUME_BLOCK(const void *start, uintptr_t size)>
+         void CONSUME_BLOCK(const void *baseAddress,
+                            const void *start, uintptr_t size)>
 void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
   addImageCallback<SEGMENT_NAME, SECTION_NAME, CONSUME_BLOCK>(mh);
 }
 
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
           const char *SEGMENT_NAME2, const char *SECTION_NAME2,
-          void CONSUME_BLOCK(const void *start, uintptr_t size,
+          void CONSUME_BLOCK(const void *baseAddress,
+                             const void *start, uintptr_t size,
                              const void *start2, uintptr_t size2)>
 void addImageCallback2Sections(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
@@ -106,11 +109,12 @@ void addImageCallback2Sections(const mach_header *mh) {
   if (!section2)
     size2 = 0;
 
-  CONSUME_BLOCK(section, size, section2, size2);
+  CONSUME_BLOCK(mh, section, size, section2, size2);
 }
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
           const char *SEGMENT_NAME2, const char *SECTION_NAME2,
-          void CONSUME_BLOCK(const void *start, uintptr_t size,
+          void CONSUME_BLOCK(const void *baseAddress,
+                             const void *start, uintptr_t size,
                              const void *start2, uintptr_t size2)>
 void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
   addImageCallback2Sections<SEGMENT_NAME, SECTION_NAME,