Win32 section lookup

Win32SectionLookup.cpp

/// Enumerate the sections in a given module.
///
/// - Parameters:
///   - module: The module to enumerate.
///   - body: A function to invoke once per section in `module`. Each
///     section's name, base address, and length (in bytes) are passed.
template <typename SectionEnumerator>
static void enumerateSections(HMODULE module, const SectionEnumerator& body) {
  auto dosHeader = reinterpret_cast<const PIMAGE_DOS_HEADER>(module);
  if (dosHeader->e_lfanew <= 0) {
    return;
  }

  auto ntHeader = reinterpret_cast<const PIMAGE_NT_HEADERS>(reinterpret_cast<const char *>(dosHeader) + dosHeader->e_lfanew);
  if (!ntHeader || ntHeader->Signature != IMAGE_NT_SIGNATURE) {
    return;
  }

  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 address = reinterpret_cast<const void *>(reinterpret_cast<const char *>(module) + section->VirtualAddress);
    size_t length = std::min(section->Misc.VirtualSize, section->SizeOfRawData);
    if (address && length > 0) {
      char name[IMAGE_SIZEOF_SHORT_NAME + 1] = {};
      strncpy_s(name, reinterpret_cast<const char *>(section->Name), IMAGE_SIZEOF_SHORT_NAME);
      body(name, address, length);
    }
  }
}

template <typename ImageEnumerator>
static void enumerateImages(const ImageEnumerator& body) {
  std::vector<HMODULE, SWTHeapAllocator<HMODULE>> modules;

  modules.resize(1);
  DWORD modulesByteCount = sizeof(HMODULE);
  while (!EnumProcessModules(GetCurrentProcess(), &modules[0], modulesByteCount, &modulesByteCount)) {
    modules.resize(modulesByteCount / sizeof(HMODULE));
  }

  for (HMODULE module : modules) {
    SWTMetadataSections sections = {};
    sections.baseAddress.store(module, std::memory_order_relaxed);

    enumerateSections(module, [&] (const char *name, const void *address, size_t length) {
      if (0 == strcmp(name, ".sw5tymd")) {
        sections.swift5_type_metadata = { reinterpret_cast<uintptr_t>(address), length };
      } else if (0 == strcmp(name, ".sw5test")) {
        sections.version = SWT_METADATA_SECTION_MINIMUM_VERSION_WITH_TESTS;
        sections.swift5_tests = { reinterpret_cast<uintptr_t>(address), length };
      }
    });

    body(sections);
  }
}