[SIGAR-66] refactor proc_{exe,args} PEB related code

src/os/win32/peb.c

@@ -27,86 +27,104 @@
 #include "sigar_os.h"
 #include <shellapi.h>
 
-#define PAGE_START    0x00020000
+void dllmod_init_ntdll(sigar_t *sigar);
-#define CWD_OFFSET    PAGE_START + 0x0290
-#define PATH_OFFSET   PAGE_START + 0x0498
-#define START_ADDRESS PAGE_START + 0x0498
 
-static int sigar_peb_get(sigar_t *sigar, HANDLE proc, DWORD *base)
+#define sigar_NtQueryInformationProcess \
+    sigar->ntdll.query_proc_info.func
+
+static int sigar_pbi_get(sigar_t *sigar, HANDLE proc, PEB *peb)
 {
-    MEMORY_BASIC_INFORMATION mbi;
+    int status;
-    DWORD bytes;
+    PROCESS_BASIC_INFORMATION pbi;
-    SIZE_T size = sigar->pagesize;
+    DWORD size=sizeof(pbi);
 
-    if (!sigar->peb) {
+    dllmod_init_ntdll(sigar);
-        sigar->peb = malloc(sigar->pagesize*2);
+
+    if (!sigar_NtQueryInformationProcess) {
+        return SIGAR_ENOTIMPL;
     }
 
-    if (!VirtualQueryEx(proc, (char*)START_ADDRESS, &mbi, sizeof(mbi))) {
+    SIGAR_ZERO(&pbi);
-        return GetLastError();
+    status =
+        sigar_NtQueryInformationProcess(proc,
+                                        ProcessBasicInformation,
+                                        &pbi,
+                                        size, NULL);
+    if (status != ERROR_SUCCESS) {
+        return status;
     }
 
-    if (mbi.RegionSize > sigar->pagesize) {
+    if (!pbi.PebBaseAddress) {
-        /* in the event args crosses the first page boundry.
+        return !SIGAR_OK;
-         * seen with WebSphere.
-         */
-        size *= 2;
     }
 
-    if (!ReadProcessMemory(proc, mbi.BaseAddress, sigar->peb,
+    size = sizeof(*peb);
-                           size, &bytes))
-    {
-        return GetLastError();
-    }
-
-    *base = (DWORD)mbi.BaseAddress;
 
+    if (ReadProcessMemory(proc, pbi.PebBaseAddress, peb, size, NULL)) {
         return SIGAR_OK;
     }
+    else {
+        return GetLastError();
+    }
+}
 
-#define SKIP_NULL(scratch) \
+static int sigar_rtl_get(sigar_t *sigar, HANDLE proc,
-    if (*scratch == '\0') scratch += sizeof(WCHAR)
+                         RTL_USER_PROCESS_PARAMETERS *rtl)
+{
+    PEB peb;
+    int status = sigar_pbi_get(sigar, proc, &peb);
+    DWORD size=sizeof(*rtl);
 
-#define PEB_START(scratch, base, offset) \
+    if (status != SIGAR_OK) {
-    scratch = sigar->peb + ((DWORD)offset - base)
+        return status;
+    }
 
-//point scratch to next string (assumes PEB_FIRST)
+    if (ReadProcessMemory(proc, peb.ProcessParameters, rtl, size, NULL)) {
-#define PEB_NEXT(scratch) \
+        return SIGAR_OK;
-    scratch = scratch + (wcslen((LPWSTR)scratch) + 1) * sizeof(WCHAR); \
+    }
-    SKIP_NULL(scratch)
+    else {
+        return GetLastError();
+    }
+}
+
+#define rtl_bufsize(buf, uc) \
+    ((sizeof(buf) < uc.Length) ? sizeof(buf) : uc.Length)
 
 int sigar_proc_exe_peb_get(sigar_t *sigar, HANDLE proc,
                            sigar_proc_exe_t *procexe)
 {
     int status;
-    LPBYTE scratch;
+    WCHAR buf[MAX_PATH+1];
-    DWORD base;
+    RTL_USER_PROCESS_PARAMETERS rtl;
-    WCHAR buf[MAX_PATH];
+    DWORD size;
 
-    if ((status = sigar_peb_get(sigar, proc, &base)) != SIGAR_OK) {
+    if ((status = sigar_rtl_get(sigar, proc, &rtl)) != SIGAR_OK) {
         return status;
     }
 
-    PEB_START(scratch, base, CWD_OFFSET);
+    size = rtl_bufsize(buf, rtl.ImagePathName);
+    memset(buf, '\0', sizeof(buf));
 
-    wcsncpy(buf, (LPWSTR)scratch, MAX_PATH);
+    if ((size > 0) &&
-    buf[MAX_PATH-1] = L'\0';
+        ReadProcessMemory(proc, rtl.ImagePathName.Buffer, buf, size, NULL))
-
+    {
-    SIGAR_W2A(buf, procexe->cwd, sizeof(procexe->cwd));
+        SIGAR_W2A(buf, procexe->name, sizeof(procexe->name));
-
+    }
-    PEB_START(scratch, base, PATH_OFFSET);
+    else {
-
+        procexe->name[0] = '\0';
-    PEB_NEXT(scratch); //skip PATH
-
-    /* XXX seen on non-english windows, random leading char */
-    if (*(scratch + sizeof(WCHAR)) != L':') {
-        scratch += sizeof(WCHAR);
     }
 
-    wcsncpy(buf, (LPWSTR)scratch, MAX_PATH);
+    size = rtl_bufsize(buf, rtl.CurrentDirectoryName);
-    buf[MAX_PATH-1] = L'\0';
+    memset(buf, '\0', sizeof(buf));
 
-    SIGAR_W2A(buf, procexe->name, sizeof(procexe->name));
+    if ((size > 0) &&
+        ReadProcessMemory(proc, rtl.CurrentDirectoryName.Buffer, buf, size, NULL))
+    {
+        SIGAR_W2A(buf, procexe->cwd, sizeof(procexe->cwd));
+    }
+    else {
+        procexe->cwd[0] = '\0';
+    }
 
     return SIGAR_OK;
 }
@@ -143,22 +161,23 @@ int sigar_proc_args_peb_get(sigar_t *sigar, HANDLE proc,
                             sigar_proc_args_t *procargs)
 {
     int status;
-    LPBYTE scratch;
-    DWORD base;
     WCHAR buf[SIGAR_CMDLINE_MAX];
+    RTL_USER_PROCESS_PARAMETERS rtl;
+    DWORD size;
 
-    if ((status = sigar_peb_get(sigar, proc, &base)) != SIGAR_OK) {
+    if ((status = sigar_rtl_get(sigar, proc, &rtl)) != SIGAR_OK) {
         return status;
     }
 
-    PEB_START(scratch, base, PATH_OFFSET);
+    size = rtl_bufsize(buf, rtl.CommandLine);
-
+    memset(buf, '\0', sizeof(buf));
-    PEB_NEXT(scratch); //skip PATH
-
-    PEB_NEXT(scratch); //skip exe name
-
-    wcsncpy(buf, (LPWSTR)scratch, SIGAR_CMDLINE_MAX);
-    buf[SIGAR_CMDLINE_MAX-1] = L'\0';
             
+    if ((size > 0) &&
+        ReadProcessMemory(proc, rtl.CommandLine.Buffer, buf, size, NULL))
+    {
         return sigar_parse_proc_args(sigar, buf, procargs);
     }
+    else {
+        return SIGAR_OK;
+    }
+}

src/os/win32/sigar_os.h

@@ -219,6 +219,93 @@ typedef struct {
 
 #define SystemProcessorPerformanceInformation 8
 
+/* PEB decls from msdn docs w/ slight mods */
+#define ProcessBasicInformation 0
+
+typedef struct _UNICODE_STRING {
+    USHORT Length;
+    USHORT MaximumLength;
+    PWSTR  Buffer;
+} UNICODE_STRING, *PUNICODE_STRING;
+
+typedef struct _PEB_LDR_DATA {
+    BYTE Reserved1[8];
+    PVOID Reserved2[3];
+    LIST_ENTRY InMemoryOrderModuleList;
+} PEB_LDR_DATA, *PPEB_LDR_DATA;
+
+typedef struct RTL_DRIVE_LETTER_CURDIR {
+    USHORT              Flags;
+    USHORT              Length;
+    ULONG               TimeStamp;
+    UNICODE_STRING      DosPath;
+} RTL_DRIVE_LETTER_CURDIR, *PRTL_DRIVE_LETTER_CURDIR;
+
+/* from: http://source.winehq.org/source/include/winternl.h */
+typedef struct _RTL_USER_PROCESS_PARAMETERS {
+    ULONG               AllocationSize;
+    ULONG               Size;
+    ULONG               Flags;
+    ULONG               DebugFlags;
+    HANDLE              hConsole;
+    ULONG               ProcessGroup;
+    HANDLE              hStdInput;
+    HANDLE              hStdOutput;
+    HANDLE              hStdError;
+    UNICODE_STRING      CurrentDirectoryName;
+    HANDLE              CurrentDirectoryHandle;
+    UNICODE_STRING      DllPath;
+    UNICODE_STRING      ImagePathName;
+    UNICODE_STRING      CommandLine;
+    PWSTR               Environment;
+    ULONG               dwX;
+    ULONG               dwY;
+    ULONG               dwXSize;
+    ULONG               dwYSize;
+    ULONG               dwXCountChars;
+    ULONG               dwYCountChars;
+    ULONG               dwFillAttribute;
+    ULONG               dwFlags;
+    ULONG               wShowWindow;
+    UNICODE_STRING      WindowTitle;
+    UNICODE_STRING      Desktop;
+    UNICODE_STRING      ShellInfo;
+    UNICODE_STRING      RuntimeInfo;
+    RTL_DRIVE_LETTER_CURDIR DLCurrentDirectory[0x20];
+} RTL_USER_PROCESS_PARAMETERS, *PRTL_USER_PROCESS_PARAMETERS;
+
+/* from msdn docs
+typedef struct _RTL_USER_PROCESS_PARAMETERS {
+    BYTE Reserved1[16];
+    PVOID Reserved2[10];
+    UNICODE_STRING ImagePathName;
+    UNICODE_STRING CommandLine;
+} RTL_USER_PROCESS_PARAMETERS, *PRTL_USER_PROCESS_PARAMETERS;
+*/
+
+typedef struct _PEB {
+    BYTE Reserved1[2];
+    BYTE BeingDebugged;
+    BYTE Reserved2[1];
+    PVOID Reserved3[2];
+    PPEB_LDR_DATA Ldr;
+    PRTL_USER_PROCESS_PARAMETERS ProcessParameters;
+    BYTE Reserved4[104];
+    PVOID Reserved5[52];
+    /*PPS_POST_PROCESS_INIT_ROUTINE*/ PVOID PostProcessInitRoutine;
+    BYTE Reserved6[128];
+    PVOID Reserved7[1];
+    ULONG SessionId;
+} PEB, *PPEB;
+
+typedef struct _PROCESS_BASIC_INFORMATION {
+    PVOID Reserved1;
+    PPEB PebBaseAddress;
+    PVOID Reserved2[2];
+    /*ULONG_PTR*/ UINT_PTR UniqueProcessId;
+    PVOID Reserved3;
+} PROCESS_BASIC_INFORMATION;
+
 typedef struct {
     sigar_pid_t pid;
     int ppid;
@@ -330,6 +417,12 @@ typedef DWORD (CALLBACK *ntdll_query_sys_info)(DWORD,
                                                ULONG,
                                                PULONG);
 
+typedef DWORD (CALLBACK *ntdll_query_proc_info)(HANDLE,
+                                                DWORD,
+                                                PVOID,
+                                                ULONG,
+                                                PULONG);
+
 /* psapi.dll */
 typedef BOOL (CALLBACK *psapi_enum_modules)(HANDLE,
                                             HMODULE *,
@@ -405,6 +498,7 @@ typedef struct {
     sigar_dll_handle_t handle;
 
     SIGAR_DLLFUNC(ntdll, query_sys_info);
+    SIGAR_DLLFUNC(ntdll, query_proc_info);
 
     sigar_dll_func_t end;
 } sigar_ntdll_t;
@@ -457,7 +551,6 @@ struct sigar_t {
 
     WORD ws_version;
     int ws_error;
-    LPBYTE peb; //scratch pad for getting peb info
     int ht_enabled;
     int lcpu; //number of logical cpus
     int winnt;

src/os/win32/win32_sigar.c

@@ -284,6 +284,7 @@ static sigar_ntdll_t sigar_ntdll = {
     "ntdll.dll",
     NULL,
     { "NtQuerySystemInformation", NULL },
+    { "NtQueryInformationProcess", NULL },
     { NULL, NULL }
 };
 
@@ -467,11 +468,15 @@ int sigar_os_open(sigar_t **sigar_ptr)
     sigar->pinfo.pid = -1;
     sigar->ws_version = 0;
     sigar->ncpu = 0;
-    sigar->peb = NULL;
 
     return result;
 }
 
+void dllmod_init_ntdll(sigar_t *sigar)
+{
+    DLLMOD_INIT(ntdll, FALSE);
+}
+
 int sigar_os_close(sigar_t *sigar)
 {
     int retval;
@@ -506,10 +511,6 @@ int sigar_os_close(sigar_t *sigar)
         sigar_cache_destroy(sigar->netif_adapters);
     }
 
-    if (sigar->peb) {
-        free(sigar->peb);
-    }
-
     free(sigar);
 
     return retval;