From 33d320fcb2da6dbc9f2daeeda0877f57064d0317 Mon Sep 17 00:00:00 2001
From: Doug MacEachern <dougm@hyperic.com>
Date: Fri, 13 May 2005 20:56:17 +0000
Subject: [PATCH] use psapi instead of CreateToolSnapshot.  it is possible for
 the latter to hang on windows 2000 and the former also works on NT

---
 src/os/win32/sigar_os.h    |   9 +++
 src/os/win32/win32_sigar.c | 109 +++++++++++++------------------------
 2 files changed, 47 insertions(+), 71 deletions(-)

diff --git a/src/os/win32/sigar_os.h b/src/os/win32/sigar_os.h
index 308c1466..3a51c15e 100644
--- a/src/os/win32/sigar_os.h
+++ b/src/os/win32/sigar_os.h
@@ -89,6 +89,12 @@ typedef DWORD (CALLBACK *LPGETUDPEXTABLE)(PMIB_UDPEXTABLE *, BOOL, HANDLE,
 
 typedef DWORD (CALLBACK *LPSYSINFO)(DWORD, PVOID, ULONG, PULONG);
 
+typedef BOOL (CALLBACK *LPENUMMODULES)(HANDLE, HMODULE*,
+                                       DWORD, LPDWORD);
+
+typedef DWORD (CALLBACK *LPGETMODULENAME)(HANDLE, HMODULE,
+                                          LPTSTR, DWORD);
+
 /* no longer in the standard header files */
 typedef struct {
     LARGE_INTEGER IdleTime;
@@ -122,6 +128,7 @@ struct sigar_t {
     DWORD perfbuf_size;
     HINSTANCE ip_handle;
     HINSTANCE nt_handle;
+    HINSTANCE ps_handle;
     LPGETIFTABLE get_if_table;
     LPGETIPFORWARDTABLE get_ipforward_table;
     LPGETTCPTABLE get_tcp_table;
@@ -129,6 +136,8 @@ struct sigar_t {
     LPGETUDPTABLE get_udp_table;
     LPGETUDPEXTABLE get_udpx_table;
     LPSYSINFO get_ntsys_info;
+    LPENUMMODULES enum_modules;
+    LPGETMODULENAME get_module_name;
     sigar_win32_pinfo_t pinfo;
     WORD ws_version;
     int ws_error;
diff --git a/src/os/win32/win32_sigar.c b/src/os/win32/win32_sigar.c
index ae87f86d..41a64555 100644
--- a/src/os/win32/win32_sigar.c
+++ b/src/os/win32/win32_sigar.c
@@ -243,6 +243,17 @@ int sigar_os_open(sigar_t **sigar)
         (*sigar)->nt_handle = NULL;
     }
 
+    if ((h = LoadLibrary("psapi.dll"))) {
+        (*sigar)->enum_modules =
+            (LPENUMMODULES)GetProcAddress(h, "EnumProcessModules");
+        (*sigar)->get_module_name =
+            (LPGETMODULENAME)GetProcAddress(h, "GetModuleFileNameExA");
+        (*sigar)->ps_handle = h;
+    }
+    else {
+        (*sigar)->ps_handle = NULL;
+    }
+
     (*sigar)->pinfo.pid = -1;
     (*sigar)->ws_version = 0;
     (*sigar)->ncpu = 0;
@@ -269,6 +280,10 @@ int sigar_os_close(sigar_t *sigar)
         FreeLibrary(sigar->nt_handle);
     }
 
+    if (sigar->ps_handle) {
+        FreeLibrary(sigar->ps_handle);
+    }
+
     if (sigar->ws_version != 0) {
         WSACleanup();
     }
@@ -1263,97 +1278,49 @@ SIGAR_DECLARE(int) sigar_proc_exe_get(sigar_t *sigar, sigar_pid_t pid,
     return status;
 }
 
-typedef HANDLE (CALLBACK *LPCREATESNAPSHOT)(DWORD, DWORD);
-typedef BOOL (CALLBACK *LPMODULEITER)(HANDLE, LPMODULEENTRY32);
-
-/* not available on NT */
-static int sigar_proc_modules_get_toolhelp(sigar_t *sigar,
-                                           sigar_pid_t pid,
-                                           sigar_proc_modules_t *procmods)
+SIGAR_DECLARE(int) sigar_proc_modules_get(sigar_t *sigar, sigar_pid_t pid,
+                                          sigar_proc_modules_t *procmods)
 {
-    HINSTANCE k32_handle;
-    HANDLE snap_shot;
-    MODULEENTRY32 module;
-    LPCREATESNAPSHOT create_snapshot;
-    LPMODULEITER module_first, module_next;
+    HANDLE proc; 
+    HMODULE modules[1024];
+    DWORD size = 0;
+    unsigned int i;
 
-    /* XXX: cache this stuff within sigar_t */
-    k32_handle = LoadLibrary("kernel32.dll");
-    if (!k32_handle) {
+    if (!sigar->ps_handle) {
+        return SIGAR_ENOTIMPL;
+    }
+
+    if (!(proc = open_process(pid))) {
         return GetLastError();
     }
 
-    create_snapshot =
-        (LPCREATESNAPSHOT)GetProcAddress(k32_handle,
-                                         "CreateToolhelp32Snapshot");
-
-    if (!create_snapshot) {
-        FreeLibrary(k32_handle);
+    if (!sigar->enum_modules(proc, modules, sizeof(modules), &size)) {
+        CloseHandle(proc);
         return GetLastError();
     }
 
-    module_first =
-        (LPMODULEITER)GetProcAddress(k32_handle, "Module32First");
+    for (i=0; i<(size/sizeof(HMODULE)); i++) {
+        int status;
+        char name[MAX_PATH];
 
-    if (!module_first) {
-        FreeLibrary(k32_handle);
-        return GetLastError();
-    }
+        if (!sigar->get_module_name(proc, modules[i], name, sizeof(name))) {
+            continue;
+        }
 
-    module_next =
-        (LPMODULEITER)GetProcAddress(k32_handle, "Module32Next");
-
-    if (!module_next) {
-        FreeLibrary(k32_handle);
-        return GetLastError();
-    }
-
-    snap_shot = create_snapshot(TH32CS_SNAPMODULE, (DWORD)pid);
-
-    if (snap_shot == INVALID_HANDLE_VALUE) {
-        FreeLibrary(k32_handle);
-        return GetLastError();
-    }
-
-    module.dwSize = sizeof(MODULEENTRY32);
-    if (!module_first(snap_shot, &module)) {
-        CloseHandle(snap_shot);
-        FreeLibrary(k32_handle);
-        return SIGAR_OK;
-    }
-
-    do {
-        int status =
-            procmods->module_getter(procmods->data,
-                                    module.szExePath,
-                                    strlen(module.szExePath));
+        status = procmods->module_getter(procmods->data,
+                                         name, strlen(name));
 
         if (status != SIGAR_OK) {
             /* not an error; just stop iterating */
             break;
         }
+    }
 
-        module.dwSize = sizeof(MODULEENTRY32);
-    } while (module_next(snap_shot, &module));
-
-    CloseHandle(snap_shot);
-    FreeLibrary(k32_handle);
+    CloseHandle(proc);
 
     return SIGAR_OK;
 }
 
-SIGAR_DECLARE(int) sigar_proc_modules_get(sigar_t *sigar, sigar_pid_t pid,
-                                          sigar_proc_modules_t *procmods)
-{
-    if (sigar->winnt) {
-        /* XXX need to use psapi.dll for NT */
-        return SIGAR_ENOTIMPL;
-    }
-    return sigar_proc_modules_get_toolhelp(sigar,
-                                           pid,
-                                           procmods);
-}
-
 #define FT2INT64(ft) \
   ((__int64)((__int64)(ft).dwHighDateTime << 32 | \
              (__int64)(ft).dwLowDateTime))