Walk PE files' sections instead of using swift_enumerateAllMetadataSections().

Sources/Testing/ExitTests/ExitTest.swift

@@ -124,7 +124,7 @@ extension ExitTest {
   public static func find(at sourceLocation: SourceLocation) -> Self? {
     var result: Self?
 
-    enumerateTypes(withNamesContaining: _exitTestContainerTypeNameMagic) { type, stop in
+    enumerateTypes(withNamesContaining: _exitTestContainerTypeNameMagic) { _, type, stop in
       if let type = type as? any __ExitTestContainer.Type, type.__sourceLocation == sourceLocation {
         result = ExitTest(
           expectedExitCondition: type.__expectedExitCondition,

Sources/Testing/Test+Discovery.swift

@@ -48,7 +48,7 @@ extension Test {
   private static var _all: some Sequence<Self> {
     get async {
       await withTaskGroup(of: [Self].self) { taskGroup in
-        enumerateTypes(withNamesContaining: _testContainerTypeNameMagic) { type, _ in
+        enumerateTypes(withNamesContaining: _testContainerTypeNameMagic) { _, type, _ in
           if let type = type as? any __TestContainer.Type {
             taskGroup.addTask {
               await type.__tests
@@ -114,11 +114,15 @@ extension Test {
 /// The type of callback called by ``enumerateTypes(withNamesContaining:_:)``.
 ///
 /// - Parameters:
+///   - imageAddress: A pointer to the start of the image. This value is _not_
+///     equal to the value returned from `dlopen()`. On platforms that do not
+///     support dynamic loading (and so do not have loadable images), this
+///     argument is unspecified.
 ///   - type: A Swift type.
 ///   - stop: An `inout` boolean variable indicating whether type enumeration
 ///     should stop after the function returns. Set `stop` to `true` to stop
 ///     type enumeration.
-typealias TypeEnumerator = (_ type: Any.Type, _ stop: inout Bool) -> Void
+typealias TypeEnumerator = (_ imageAddress: UnsafeRawPointer?, _ type: Any.Type, _ stop: inout Bool) -> Void
 
 /// Enumerate all types known to Swift found in the current process whose names
 /// contain a given substring.
@@ -129,11 +133,11 @@ typealias TypeEnumerator = (_ type: Any.Type, _ stop: inout Bool) -> Void
 func enumerateTypes(withNamesContaining nameSubstring: String, _ typeEnumerator: TypeEnumerator) {
   withoutActuallyEscaping(typeEnumerator) { typeEnumerator in
     withUnsafePointer(to: typeEnumerator) { context in
-      swt_enumerateTypes(withNamesContaining: nameSubstring, .init(mutating: context)) { type, stop, context in
+      swt_enumerateTypes(withNamesContaining: nameSubstring, .init(mutating: context)) { imageAddress, type, stop, context in
         let typeEnumerator = context!.load(as: TypeEnumerator.self)
         let type = unsafeBitCast(type, to: Any.Type.self)
         var stop2 = false
-        typeEnumerator(type, &stop2)
+        typeEnumerator(imageAddress, type, &stop2)
         stop.pointee = stop2
       }
     }

Sources/_TestingInternals/Discovery.cpp

@@ -10,22 +10,49 @@
 
 #include "Discovery.h"
 
+#include <algorithm>
+#include <array>
 #include <atomic>
 #include <cstring>
 #include <iterator>
+#include <tuple>
 #include <type_traits>
 #include <vector>
+#include <optional>
 
-#if defined(SWT_NO_DYNAMIC_LINKING)
-#include <algorithm>
-#elif defined(__APPLE__)
+#if defined(__APPLE__) && !defined(SWT_NO_DYNAMIC_LINKING)
 #include <dispatch/dispatch.h>
 #include <mach-o/dyld.h>
 #include <mach-o/getsect.h>
 #include <objc/runtime.h>
 #include <os/lock.h>
 #endif
 
+/// A type that acts as a C++ [Allocator](https://en.cppreference.com/w/cpp/named_req/Allocator)
+/// without using global `operator new` or `operator delete`.
+///
+/// This type is necessary because global `operator new` and `operator delete`
+/// can be overridden in developer-supplied code and cause deadlocks or crashes
+/// when subsequently used while holding a dyld- or libobjc-owned lock. Using
+/// `std::malloc()` and `std::free()` allows the use of C++ container types
+/// without this risk.
+template<typename T>
+struct SWTHeapAllocator {
+  using value_type = T;
+
+  T *allocate(size_t count) {
+    return reinterpret_cast<T *>(std::calloc(count, sizeof(T)));
+  }
+
+  void deallocate(T *ptr, size_t count) {
+    std::free(ptr);
+  }
+};
+
+/// A `std::vector` that uses `SWTHeapAllocator`.
+template <typename T>
+using SWTVector = std::vector<T, SWTHeapAllocator<T>>;
+
 /// Enumerate over all Swift type metadata sections in the current process.
 ///
 /// - Parameters:
@@ -199,39 +226,19 @@ extern "C" const char sectionEnd __asm("section$end$__TEXT$__swift5_types");
 template <typename SectionEnumerator>
 static void enumerateTypeMetadataSections(const SectionEnumerator& body) {
   auto size = std::distance(&sectionBegin, &sectionEnd);
-  body(&sectionBegin, size);
+  bool stop = false;
+  body(nullptr, &sectionBegin, size, &stop);
 }
 
 #elif defined(__APPLE__)
 #pragma mark - Apple implementation
 
-/// A type that acts as a C++ [Allocator](https://en.cppreference.com/w/cpp/named_req/Allocator)
-/// without using global `operator new` or `operator delete`.
-///
-/// This type is necessary because global `operator new` and `operator delete`
-/// can be overridden in developer-supplied code and cause deadlocks or crashes
-/// when subsequently used while holding a dyld- or libobjc-owned lock. Using
-/// `std::malloc()` and `std::free()` allows the use of C++ container types
-/// without this risk.
-template<typename T>
-struct SWTHeapAllocator {
-  using value_type = T;
-
-  T *allocate(size_t count) {
-    return reinterpret_cast<T *>(std::calloc(count, sizeof(T)));
-  }
-
-  void deallocate(T *ptr, size_t count) {
-    std::free(ptr);
-  }
-};
-
 /// A type that acts as a C++ [Container](https://en.cppreference.com/w/cpp/named_req/Container)
 /// and which contains a sequence of Mach headers.
 #if __LP64__
-using SWTMachHeaderList = std::vector<const mach_header_64 *, SWTHeapAllocator<const mach_header_64 *>>;
+using SWTMachHeaderList = SWTVector<const mach_header_64 *>;
 #else
-using SWTMachHeaderList = std::vector<const mach_header *, SWTHeapAllocator<const mach_header *>>;
+using SWTMachHeaderList = SWTVector<const mach_header *>;
 #endif
 
 /// Get a copy of the currently-loaded Mach headers list.
@@ -301,13 +308,118 @@ static void enumerateTypeMetadataSections(const SectionEnumerator& body) {
     unsigned long size = 0;
     const void *section = getsectiondata(mh, SEG_TEXT, "__swift5_types", &size);
     if (section && size > 0) {
-      body(section, size);
+      bool stop = false;
+      body(mh, section, size, &stop);
+      if (stop) {
+        break;
+      }
+    }
+  }
+}
+
+#elif defined(_WIN32)
+#pragma mark - Windows implementation
+
+/// Find the section with the given name in the given module.
+///
+/// - Parameters:
+///   - module: The module to inspect.
+///   - sectionName: The name of the section to look for. Long section names are
+///     not supported.
+///
+/// - Returns: A pointer to the start of the given section along with its size
+///   in bytes, or `std::nullopt` if the section could not be found. If the
+///   section was emitted by the Swift toolchain, be aware it will have leading
+///   and trailing bytes (`sizeof(uintptr_t)` each.)
+static std::optional<std::pair<const void *, size_t>> findSection(HMODULE module, const char *sectionName) {
+  if (!module) {
+    return std::nullopt;
+  }
+
+  // Get the DOS header (to which the HMODULE directly points, conveniently!)
+  // and check it's sufficiently valid for us to walk.
+  auto dosHeader = reinterpret_cast<const PIMAGE_DOS_HEADER>(module);
+  if (dosHeader->e_magic != IMAGE_DOS_SIGNATURE || dosHeader->e_lfanew <= 0) {
+    return std::nullopt;
+  }
+
+  // Check the NT header as well as the optional header.
+  auto ntHeader = reinterpret_cast<const PIMAGE_NT_HEADERS>(reinterpret_cast<uintptr_t>(dosHeader) + dosHeader->e_lfanew);
+  if (!ntHeader || ntHeader->Signature != IMAGE_NT_SIGNATURE) {
+    return std::nullopt;
+  }
+  if (ntHeader->FileHeader.SizeOfOptionalHeader < offsetof(decltype(ntHeader->OptionalHeader), Magic) + sizeof(decltype(ntHeader->OptionalHeader)::Magic)) {
+    return std::nullopt;
+  }
+  if (ntHeader->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR_MAGIC) {
+    return std::nullopt;
+  }
+
+  auto sectionCount = ntHeader->FileHeader.NumberOfSections;
+  auto section = IMAGE_FIRST_SECTION(ntHeader);
+  for (size_t i = 0; i < sectionCount; i++, section += 1) {
+    if (section->VirtualAddress == 0) {
+      continue;
+    }
+
+    auto start = reinterpret_cast<const void *>(reinterpret_cast<uintptr_t>(dosHeader) + section->VirtualAddress);
+    size_t size = std::min(section->Misc.VirtualSize, section->SizeOfRawData);
+    if (start && size > 0) {
+      // FIXME: Handle longer names ("/%u") from string table
+      auto thisSectionName = reinterpret_cast<const char *>(section->Name);
+      if (0 == std::strncmp(sectionName, thisSectionName, IMAGE_SIZEOF_SHORT_NAME)) {
+        return std::make_pair(start, size);
+      }
+    }
+  }
+
+  return std::nullopt;
+}
+
+template <typename SectionEnumerator>
+static void enumerateTypeMetadataSections(const SectionEnumerator& body) {
+  // Find all the modules loaded in the current process. We assume there aren't
+  // more than 1024 loaded modules (as does Microsoft sample code.)
+  std::array<HMODULE, 1024> hModules;
+  DWORD byteCountNeeded = 0;
+  if (!EnumProcessModules(GetCurrentProcess(), &hModules[0], hModules.size() * sizeof(HMODULE), &byteCountNeeded)) {
+    return;
+  }
+  DWORD hModuleCount = std::min(hModules.size(), byteCountNeeded / sizeof(HMODULE));
+
+  // Look in all the loaded modules for Swift type metadata sections and store
+  // them in a side table.
+  //
+  // This two-step process is less algorithmically efficient than a single loop,
+  // but it is safer: the callback will eventually invoke developer code that
+  // could theoretically unload a module from the list we're enumerating. (Swift
+  // modules do not support unloading, so we'll just not worry about them.)
+  using SWTSectionList = SWTVector<std::tuple<HMODULE, const void *, size_t>>;
+  SWTSectionList sectionList;
+  for (DWORD i = 0; i < hModuleCount; i++) {
+    if (auto section = findSection(hModules[i], ".sw5tymd")) {
+      sectionList.emplace_back(hModules[i], section->first, section->second);
+    }
+  }
+
+  // Pass the loaded module and section info back to the body callback.
+  // Note we ignore the leading and trailing uintptr_t values: they're both
+  // always set to zero so we'll skip them in the callback, and in the future
+  // the toolchain might not emit them at all in which case we don't want to
+  // skip over real section data.
+  bool stop = false;
+  for (const auto& section : sectionList) {
+    // TODO: Use C++17 unstructured binding here.
+    body(get<0>(section), get<1>(section), get<2>(section), &stop);
+    if (stop) {
+      break;
     }
   }
 }
 
-#elif defined(__linux__) || defined(__FreeBSD__) || defined(_WIN32) || defined(__wasi__) || defined(__ANDROID__)
-#pragma mark - Linux/Windows implementation
+
+#elif defined(__linux__) || defined(__FreeBSD__) || defined(__wasi__) || defined(__ANDROID__)
+#pragma mark - ELF implementation
 
 /// Specifies the address range corresponding to a section.
 struct MetadataSectionRange {
@@ -352,7 +464,11 @@ static void enumerateTypeMetadataSections(const SectionEnumerator& body) {
     const auto& body = *reinterpret_cast<const SectionEnumerator *>(context);
     MetadataSectionRange section = sections->swift5_type_metadata;
     if (section.start && section.length > 0) {
-      body(reinterpret_cast<const void *>(section.start), section.length);
+      bool stop = false;
+      body(sections->baseAddress.load(), reinterpret_cast<const void *>(section.start), section.length, &stop);
+      if (stop) {
+        return false;
+      }
     }
     return true;
   }, const_cast<SectionEnumerator *>(&body));
@@ -366,12 +482,11 @@ static void enumerateTypeMetadataSections(const SectionEnumerator& body) {}
 #pragma mark -
 
 void swt_enumerateTypesWithNamesContaining(const char *nameSubstring, void *context, SWTTypeEnumerator body) {
-  enumerateTypeMetadataSections([=] (const void *section, size_t size) {
+  enumerateTypeMetadataSections([=] (const void *imageAddress, const void *section, size_t size, bool *stop) {
     auto records = reinterpret_cast<const SWTTypeMetadataRecord *>(section);
     size_t recordCount = size / sizeof(SWTTypeMetadataRecord);
 
-    bool stop = false;
-    for (size_t i = 0; i < recordCount && !stop; i++) {
+    for (size_t i = 0; i < recordCount && !*stop; i++) {
       const auto& record = records[i];
 
       auto contextDescriptor = record.getContextDescriptor();
@@ -394,7 +509,7 @@ void swt_enumerateTypesWithNamesContaining(const char *nameSubstring, void *cont
       }
 
       if (void *typeMetadata = contextDescriptor->getMetadata()) {
-        body(typeMetadata, &stop, context);
+        body(imageAddress, typeMetadata, stop, context);
       }
     }
   });

Sources/_TestingInternals/include/Discovery.h

@@ -19,13 +19,17 @@ SWT_ASSUME_NONNULL_BEGIN
 /// The type of callback called by `swt_enumerateTypes()`.
 ///
 /// - Parameters:
+///   - imageAddress: A pointer to the start of the image. This value is _not_
+///     equal to the value returned from `dlopen()`. On platforms that do not
+///     support dynamic loading (and so do not have loadable images), this
+///     argument is unspecified.
 ///   - typeMetadata: A type metadata pointer that can be bitcast to `Any.Type`.
 ///   - stop: A pointer to a boolean variable indicating whether type
 ///     enumeration should stop after the function returns. Set `*stop` to
 ///     `true` to stop type enumeration.
 ///   - context: An arbitrary pointer passed by the caller to
 ///     `swt_enumerateTypes()`.
-typedef void (* SWTTypeEnumerator)(void *typeMetadata, bool *stop, void *_Null_unspecified context);
+typedef void (* SWTTypeEnumerator)(const void *_Null_unspecified imageAddress, void *typeMetadata, bool *stop, void *_Null_unspecified context);
 
 /// Enumerate all types known to Swift found in the current process.
 ///