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

On Windows, the platform provides sufficient API for us to walk images
at runtime looking for those that contain Swift metadata (similar to
what we do on Darwin.) Therefore, we don't need to use the
`swift_enumerateAllMetadataSections()` function from the Swift runtime
on Windows.

This change also plumbs through an `imageAddress` argument to the
discovery callback function. We're not currently using it, but it should
be useful in the future for diagnostics (e.g. indicating that some data
from a given image is malformed, or for passing to downstream testing
tools that can make use of it.)

This change _also_ also plumbs the `stop` argument _down_ into the
platform implementations so that we don't keep iterating images after
the caller has told us to stop. In practice this has not been a
correctness problem, but it does impact performance of type discovery.

### Checklist:

- [x] Code and documentation should follow the style of the [Style
Guide](https://github.com/apple/swift-testing/blob/main/Documentation/StyleGuide.md).
- [x] If public symbols are renamed or modified, DocC references should
be updated.

Author: grynspan

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.
 ///