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

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,9 +159,10 @@ 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;
     }
@@ -183,6 +170,24 @@ const char *swift_getMetadataSectionName(void *metadata_section) {
   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)