1 files changed, 68 insertions, 49 deletions

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_win.cc b/compiler-rt/lib/sanitizer_common/sanitizer_win.cc
index 66448ea950c..e246d5aeca1 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_win.cc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_win.cc
@@ -209,76 +209,46 @@ u32 GetUid() {
 
 namespace {
 struct ModuleInfo {
-  HMODULE handle;
+  const char *filepath;
   uptr base_address;
   uptr end_address;
 };
 
 int CompareModulesBase(const void *pl, const void *pr) {
-  const ModuleInfo &l = *(ModuleInfo *)pl, &r = *(ModuleInfo *)pr;
-  if (l.base_address < r.base_address)
+  const ModuleInfo *l = (ModuleInfo *)pl, *r = (ModuleInfo *)pr;
+  if (l->base_address < r->base_address)
     return -1;
-  return l.base_address > r.base_address;
+  return l->base_address > r->base_address;
 }
 }  // namespace
 
 #ifndef SANITIZER_GO
 void DumpProcessMap() {
   Report("Dumping process modules:\n");
-  HANDLE cur_process = GetCurrentProcess();
-
-  // Query the list of modules.  Start by assuming there are no more than 256
-  // modules and retry if that's not sufficient.
-  ModuleInfo *modules;
-  size_t num_modules;
-  {
-    HMODULE *hmodules = 0;
-    uptr modules_buffer_size = sizeof(HMODULE) * 256;
-    DWORD bytes_required;
-    while (!hmodules) {
-      hmodules = (HMODULE *)MmapOrDie(modules_buffer_size, __FUNCTION__);
-      CHECK(EnumProcessModules(cur_process, hmodules, modules_buffer_size,
-                               &bytes_required));
-      if (bytes_required > modules_buffer_size) {
-        // Either there turned out to be more than 256 hmodules, or new hmodules
-        // could have loaded since the last try.  Retry.
-        UnmapOrDie(hmodules, modules_buffer_size);
-        hmodules = 0;
-        modules_buffer_size = bytes_required;
-      }
-    }
+  InternalScopedBuffer<LoadedModule> modules(kMaxNumberOfModules);
+  uptr num_modules =
+      GetListOfModules(modules.data(), kMaxNumberOfModules, nullptr);
 
-    num_modules = bytes_required / sizeof(HMODULE);
-    modules =
-        (ModuleInfo *)MmapOrDie(num_modules * sizeof(ModuleInfo), __FUNCTION__);
-    for (size_t i = 0; i < num_modules; ++i) {
-      modules[i].handle = hmodules[i];
-      MODULEINFO mi;
-      if (!GetModuleInformation(cur_process, hmodules[i], &mi, sizeof(mi)))
-        continue;
-      modules[i].base_address = (uptr)mi.lpBaseOfDll;
-      modules[i].end_address = (uptr)mi.lpBaseOfDll + mi.SizeOfImage;
-    }
-    UnmapOrDie(hmodules, modules_buffer_size);
+  InternalScopedBuffer<ModuleInfo> module_infos(num_modules);
+  for (size_t i = 0; i < num_modules; ++i) {
+    module_infos[i].filepath = modules[i].full_name();
+    module_infos[i].base_address = modules[i].base_address();
+    module_infos[i].end_address = modules[i].ranges().next()->end;
   }
-
-  qsort(modules, num_modules, sizeof(ModuleInfo), CompareModulesBase);
+  qsort(module_infos.data(), num_modules, sizeof(ModuleInfo),
+        CompareModulesBase);
 
   for (size_t i = 0; i < num_modules; ++i) {
-    const ModuleInfo &mi = modules[i];
-    char module_name[MAX_PATH];
-    bool got_module_name = GetModuleFileNameA(
-        mi.handle, module_name, sizeof(module_name));
+    const ModuleInfo &mi = module_infos[i];
     if (mi.end_address != 0) {
       Printf("\t%p-%p %s\n", mi.base_address, mi.end_address,
-             got_module_name ? module_name : "[no name]");
-    } else if (got_module_name) {
-      Printf("\t??\?-??? %s\n", module_name);
+             mi.filepath[0] ? mi.filepath : "[no name]");
+    } else if (mi.filepath) {
+      Printf("\t??\?-??? %s\n", mi.filepath);
     } else {
       Printf("\t???\n");
     }
   }
-  UnmapOrDie(modules, num_modules * sizeof(ModuleInfo));
 }
 #endif
 
@@ -349,7 +319,56 @@ void Abort() {
 
 uptr GetListOfModules(LoadedModule *modules, uptr max_modules,
                       string_predicate_t filter) {
-  UNIMPLEMENTED();
+  HANDLE cur_process = GetCurrentProcess();
+
+  // Query the list of modules.  Start by assuming there are no more than 256
+  // modules and retry if that's not sufficient.
+  HMODULE *hmodules = 0;
+  uptr modules_buffer_size = sizeof(HMODULE) * 256;
+  DWORD bytes_required;
+  while (!hmodules) {
+    hmodules = (HMODULE *)MmapOrDie(modules_buffer_size, __FUNCTION__);
+    CHECK(EnumProcessModules(cur_process, hmodules, modules_buffer_size,
+                             &bytes_required));
+    if (bytes_required > modules_buffer_size) {
+      // Either there turned out to be more than 256 hmodules, or new hmodules
+      // could have loaded since the last try.  Retry.
+      UnmapOrDie(hmodules, modules_buffer_size);
+      hmodules = 0;
+      modules_buffer_size = bytes_required;
+    }
+  }
+
+  // |num_modules| is the number of modules actually present,
+  // |count| is the number of modules we return.
+  size_t nun_modules = bytes_required / sizeof(HMODULE),
+         count = 0;
+  for (size_t i = 0; i < nun_modules && count < max_modules; ++i) {
+    HMODULE handle = hmodules[i];
+    MODULEINFO mi;
+    if (!GetModuleInformation(cur_process, handle, &mi, sizeof(mi)))
+      continue;
+
+    char module_name[MAX_PATH];
+    bool got_module_name =
+        GetModuleFileNameA(handle, module_name, sizeof(module_name));
+    if (!got_module_name)
+      module_name[0] = '\0';
+
+    if (filter && !filter(module_name))
+      continue;
+
+    uptr base_address = (uptr)mi.lpBaseOfDll;
+    uptr end_address = (uptr)mi.lpBaseOfDll + mi.SizeOfImage;
+    LoadedModule *cur_module = &modules[count];
+    cur_module->set(module_name, base_address);
+    // We add the whole module as one single address range.
+    cur_module->addAddressRange(base_address, end_address, /*executable*/ true);
+    count++;
+  }
+  UnmapOrDie(hmodules, modules_buffer_size);
+
+  return count;
 };
 
 #ifndef SANITIZER_GO