unify cmdline and cwd reading some more

Rename psutil_get_parameter to psutil_get_process_data
and let it return only the data to the caller.  This reduces
the needed local variables and makes the function slightly
simpler.

1 files changed, 134 insertions, 196 deletions

diff --git a/psutil/arch/windows/process_info.c b/psutil/arch/windows/process_info.c
index e66166c0..df7ae9fb 100644
--- a/psutil/arch/windows/process_info.c
+++ b/psutil/arch/windows/process_info.c
@@ -308,15 +308,20 @@ typedef struct {
 } PEB64;
 #endif
 
-/* Get one or more parameters of the process with the given pid:
-
-   pcmdline: the command line as a Python list
-   pcwd: the current working directory as a Python string
-
-   On success 0 is returned.  On error the given output parameters are not
-   touched, -1 is returned, and an appropriate Python exception is set. */
-static int psutil_get_parameters(long pid, PyObject **pcmdline,
-                                 PyObject **pcwd) {
+enum psutil_process_data_kind {
+    KIND_CMDLINE,
+    KIND_CWD,
+};
+
+/* Get data from the process with the given pid.  The data is returned in the
+   pdata output member as a nul terminated string which must be freed on
+   success.
+
+   On success 0 is returned.  On error the output parameter is not touched, -1
+   is returned, and an appropriate Python exception is set. */
+static int psutil_get_process_data(long pid,
+                                   enum psutil_process_data_kind kind,
+                                   WCHAR **pdata) {
     /* This function is quite complex because there are several cases to be
        considered:
 
@@ -346,19 +351,17 @@ static int psutil_get_parameters(long pid, PyObject **pcmdline,
     static _NtQueryInformationProcess NtWow64QueryInformationProcess64 = NULL;
     static _NtWow64ReadVirtualMemory64 NtWow64ReadVirtualMemory64 = NULL;
 #endif
-    int nArgs, i;
-    LPWSTR *szArglist = NULL;
     HANDLE hProcess = NULL;
-    WCHAR *commandLineContents = NULL;
-    WCHAR *currentDirectoryContent = NULL;
+    LPCVOID src;
+    SIZE_T size;
+    WCHAR *buffer = NULL;
 #ifdef _WIN64
     LPVOID ppeb32 = NULL;
 #else
+    PVOID64 src64;
     BOOL weAreWow64;
     BOOL theyAreWow64;
 #endif
-    PyObject *py_unicode = NULL;
-    PyObject *py_retlist = NULL;
 
     hProcess = psutil_handle_from_pid(pid);
     if (hProcess == NULL)
@@ -380,29 +383,20 @@ static int psutil_get_parameters(long pid, PyObject **pcmdline,
         PyErr_SetFromWindowsErr(0);
         goto error;
     }
-#else
-    /* 32 bit case.  Check if the target is also 32 bit. */
-    if (!IsWow64Process(GetCurrentProcess(), &weAreWow64) ||
-        !IsWow64Process(hProcess, &theyAreWow64)) {
-        PyErr_SetFromWindowsErr(0);
-        goto error;
-    }
-#endif
 
-#ifdef _WIN64
     if (ppeb32 != NULL) {
         /* We are 64 bit.  Target process is 32 bit running in WoW64 mode. */
         PEB32 peb32;
         RTL_USER_PROCESS_PARAMETERS32 procParameters32;
 
         // read PEB
-        if(!ReadProcessMemory(hProcess, ppeb32, &peb32, sizeof(peb32), NULL)) {
+        if (!ReadProcessMemory(hProcess, ppeb32, &peb32, sizeof(peb32), NULL)) {
             PyErr_SetFromWindowsErr(0);
             goto error;
         }
 
         // read process parameters
-        if(!ReadProcessMemory(hProcess,
+        if (!ReadProcessMemory(hProcess,
                               UlongToPtr(peb32.ProcessParameters),
                               &procParameters32,
                               sizeof(procParameters32),
@@ -411,47 +405,25 @@ static int psutil_get_parameters(long pid, PyObject **pcmdline,
             goto error;
         }
 
-        if (pcmdline != NULL) {
-            // read command line aguments
-            commandLineContents = 
-                calloc(procParameters32.CommandLine.Length + 2, 1);
-            if (commandLineContents == NULL) {
-                PyErr_NoMemory();
-                goto error;
-            }
-
-            if(!ReadProcessMemory(
-                        hProcess,
-                        UlongToPtr(procParameters32.CommandLine.Buffer),
-                        commandLineContents,
-                        procParameters32.CommandLine.Length,
-                        NULL)) {
-                PyErr_SetFromWindowsErr(0);
-                goto error;
-            }
-        }
-
-        if (pcwd != NULL) {
-            // read cwd
-            currentDirectoryContent =
-                calloc(procParameters32.CurrentDirectoryPath.Length + 2, 1);
-            if (currentDirectoryContent == NULL) {
-                PyErr_NoMemory();
-                goto error;
-            }
-
-            if (!ReadProcessMemory(
-                    hProcess,
-                    UlongToPtr(procParameters32.CurrentDirectoryPath.Buffer),
-                    currentDirectoryContent,
-                    procParameters32.CurrentDirectoryPath.Length,
-                    NULL)) {
-                PyErr_SetFromWindowsErr(0);
-                goto error;
-            }
+        switch (kind) {
+            case KIND_CMDLINE:
+                src = UlongToPtr(procParameters32.CommandLine.Buffer),
+                size = procParameters32.CommandLine.Length;
+                break;
+            case KIND_CWD:
+                src = UlongToPtr(procParameters32.CurrentDirectoryPath.Buffer);
+                size = procParameters32.CurrentDirectoryPath.Length;
+                break;
         }
     } else
 #else
+    /* 32 bit case.  Check if the target is also 32 bit. */
+    if (!IsWow64Process(GetCurrentProcess(), &weAreWow64) ||
+        !IsWow64Process(hProcess, &theyAreWow64)) {
+        PyErr_SetFromWindowsErr(0);
+        goto error;
+    }
+
     if (weAreWow64 && !theyAreWow64) {
         /* We are 32 bit running in WoW64 mode.  Target process is 64 bit. */
         PROCESS_BASIC_INFORMATION64 pbi64;
@@ -502,44 +474,15 @@ static int psutil_get_parameters(long pid, PyObject **pcmdline,
             goto error;
         }
 
-        if (pcmdline != NULL) {
-            // read command line aguments
-            commandLineContents =
-                calloc(procParameters64.CommandLine.Length + 2, 1);
-            if(!commandLineContents) {
-                PyErr_NoMemory();
-                goto error;
-            }
-
-            if (!NT_SUCCESS(NtWow64ReadVirtualMemory64(
-                            hProcess,
-                            procParameters64.CommandLine.Buffer,
-                            commandLineContents,
-                            procParameters64.CommandLine.Length,
-                            NULL))) {
-                PyErr_SetFromWindowsErr(0);
-                goto error;
-            }
-        }
-
-        if (pcwd != NULL) {
-            // read cwd
-            currentDirectoryContent =
-                calloc(procParameters64.CurrentDirectoryPath.Length + 2, 1);
-            if (!currentDirectoryContent) {
-                PyErr_NoMemory();
-                goto error;
-            }
-
-            if (!NT_SUCCESS(NtWow64ReadVirtualMemory64(
-                            hProcess,
-                            procParameters64.CurrentDirectoryPath.Buffer,
-                            currentDirectoryContent,
-                            procParameters64.CurrentDirectoryPath.Length,
-                            NULL))) {
-                PyErr_SetFromWindowsErr(0);
-                goto error;
-            }
+        switch (kind) {
+            case KIND_CMDLINE:
+                src64 = procParameters64.CommandLine.Buffer;
+                size = procParameters64.CommandLine.Length;
+                break;
+            case KIND_CWD:
+                src64 = procParameters64.CurrentDirectoryPath.Buffer,
+                size = procParameters64.CurrentDirectoryPath.Length;
+                break;
         }
     } else
 #endif
@@ -560,123 +503,70 @@ static int psutil_get_parameters(long pid, PyObject **pcmdline,
         }
 
         // read peb
-        if(!ReadProcessMemory(hProcess,
-                              pbi.PebBaseAddress,
-                              &peb,
-                              sizeof(peb),
-                              NULL)) {
+        if (!ReadProcessMemory(hProcess,
+                               pbi.PebBaseAddress,
+                               &peb,
+                               sizeof(peb),
+                               NULL)) {
             PyErr_SetFromWindowsErr(0);
             goto error;
         }
 
         // read process parameters
-        if(!ReadProcessMemory(hProcess,
-                              peb.ProcessParameters,
-                              &procParameters,
-                              sizeof(procParameters),
-                              NULL)) {
+        if (!ReadProcessMemory(hProcess,
+                               peb.ProcessParameters,
+                               &procParameters,
+                               sizeof(procParameters),
+                               NULL)) {
             PyErr_SetFromWindowsErr(0);
             goto error;
         }
 
-        if (pcmdline != NULL) {
-            // read command line aguments
-            commandLineContents =
-                calloc(procParameters.CommandLine.Length + 2, 1);
-            if(!commandLineContents) {
-                PyErr_NoMemory();
-                goto error;
-            }
-
-            if(!ReadProcessMemory(hProcess,
-                                  procParameters.CommandLine.Buffer,
-                                  commandLineContents,
-                                  procParameters.CommandLine.Length,
-                                  NULL)) {
-                PyErr_SetFromWindowsErr(0);
-                goto error;
-            }
+        switch (kind) {
+            case KIND_CMDLINE:
+                src = procParameters.CommandLine.Buffer;
+                size = procParameters.CommandLine.Length;
+                break;
+            case KIND_CWD: 
+                src = procParameters.CurrentDirectoryPath.Buffer;
+                size = procParameters.CurrentDirectoryPath.Length;
+                break;
         }
+    }
 
-        if (pcwd != NULL) {
-            // read cwd
-            currentDirectoryContent =
-                calloc(procParameters.CurrentDirectoryPath.Length + 2, 1);
-            if (!currentDirectoryContent) {
-                PyErr_NoMemory();
-                goto error;
-            }
+    buffer = calloc(size + 2, 1);
 
-            if (!ReadProcessMemory(hProcess,
-                                   procParameters.CurrentDirectoryPath.Buffer,
-                                   currentDirectoryContent,
-                                   procParameters.CurrentDirectoryPath.Length,
-                                   NULL)) {
-                PyErr_SetFromWindowsErr(0);
-                goto error;
-            }
-        }
+    if (buffer == NULL) {
+        PyErr_NoMemory();
+        goto error;
     }
 
-    if (pcmdline != NULL) {
-        // attempt to parse the command line using Win32 API
-        szArglist = CommandLineToArgvW(commandLineContents, &nArgs);
-        if (szArglist == NULL) {
+#ifndef _WIN64
+    if (weAreWow64 && !theyAreWow64) {
+        if (!NT_SUCCESS(NtWow64ReadVirtualMemory64(hProcess,
+                                                   src64,
+                                                   buffer,
+                                                   size,
+                                                   NULL))) {
             PyErr_SetFromWindowsErr(0);
             goto error;
         }
-        else {
-            // arglist parsed as array of UNICODE_STRING, so convert each to
-            // Python string object and add to arg list
-            py_retlist = Py_BuildValue("[]");
-            if (py_retlist == NULL)
-                goto error;
-            for (i = 0; i < nArgs; i++) {
-                py_unicode = PyUnicode_FromWideChar(
-                    szArglist[i], wcslen(szArglist[i]));
-                if (py_unicode == NULL)
-                    goto error;
-                if (PyList_Append(py_retlist, py_unicode))
-                    goto error;
-                Py_CLEAR(py_unicode);
-            }
-        }
-
-        if (szArglist != NULL)
-            LocalFree(szArglist);
-        free(commandLineContents);
-
-        *pcmdline = py_retlist;
-        py_retlist = NULL;
-    }
-
-    if (pcwd != NULL) {
-        // convert wchar array to a Python unicode string
-        py_unicode = PyUnicode_FromWideChar(currentDirectoryContent,
-                                            wcslen(currentDirectoryContent));
-        if (py_unicode == NULL)
-            goto error;
-        CloseHandle(hProcess);
-        free(currentDirectoryContent);
-        *pcwd = py_unicode;
-        py_unicode = NULL;
+    } else
+#endif
+    if (!ReadProcessMemory(hProcess, src, buffer, size, NULL)) {
+        PyErr_SetFromWindowsErr(0);
+        goto error;
     }
 
-    CloseHandle(hProcess);
+    *pdata = buffer;
 
     return 0;
 
 error:
-    Py_XDECREF(py_unicode);
-    Py_XDECREF(py_retlist);
     if (hProcess != NULL)
         CloseHandle(hProcess);
-    if (commandLineContents != NULL)
-        free(commandLineContents);
-    if (szArglist != NULL)
-        LocalFree(szArglist);
-    if (currentDirectoryContent != NULL)
-        free(currentDirectoryContent);
+    if (buffer != NULL)
+        free(buffer);
     return -1;
 }
 
@@ -687,13 +577,61 @@ error:
 PyObject *
 psutil_get_cmdline(long pid) {
     PyObject *ret = NULL;
-    psutil_get_parameters(pid, &ret, NULL);
+    WCHAR *data = NULL;
+    PyObject *py_retlist = NULL;
+    PyObject *py_unicode = NULL;
+    LPWSTR *szArglist = NULL;
+    int nArgs, i;
+
+    if (psutil_get_process_data(pid, KIND_CMDLINE, &data) != 0)
+        goto out;
+
+    // attempt to parse the command line using Win32 API
+    szArglist = CommandLineToArgvW(data, &nArgs);
+    if (szArglist == NULL) {
+        PyErr_SetFromWindowsErr(0);
+        goto out;
+    }
+
+    // arglist parsed as array of UNICODE_STRING, so convert each to
+    // Python string object and add to arg list
+    py_retlist = PyList_New(nArgs);
+    if (py_retlist == NULL)
+        goto out;
+    for (i = 0; i < nArgs; i++) {
+        py_unicode = PyUnicode_FromWideChar(szArglist[i],
+                                            wcslen(szArglist[i]));
+        if (py_unicode == NULL)
+            goto out;
+        PyList_SET_ITEM(py_retlist, i, py_unicode);
+        py_unicode = NULL;
+    }
+
+    ret = py_retlist;
+    py_retlist = NULL;
+
+out:
+    LocalFree(szArglist);
+    free(data);
+    Py_XDECREF(py_unicode);
+    Py_XDECREF(py_retlist);
+
     return ret;
 }
 
 PyObject *psutil_get_cwd(long pid) {
     PyObject *ret = NULL;
-    psutil_get_parameters(pid, NULL, &ret);
+    WCHAR *data = NULL;
+
+    if (psutil_get_process_data(pid, KIND_CWD, &data) != 0)
+        goto out;
+
+    // convert wchar array to a Python unicode string
+    ret = PyUnicode_FromWideChar(data, wcslen(data));
+
+out:
+    free(data);
+
     return ret;
 }