Skip to content

Commit

Permalink
#1398 / win / cmdline: call NtQueryInformationProcess twice, the firs…
Browse files Browse the repository at this point in the history 
…t time to get the right buf size (ProcessHacker does this)
  • Loading branch information
@giampaolo
giampaolo committed Feb 28, 2019
1 parent e8bd07f commit 59e3c5e
Showing 1 changed file with 22 additions and 4 deletions.
26 changes: 22 additions & 4 deletions psutil/arch/windows/process_info.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,8 @@ typedef struct {
// ====================================================================


#define STATUS_BUFFER_OVERFLOW ((NTSTATUS)0x80000005L)

/*
* Return 1 if PID exists, 0 if not, -1 on error.
*/
Expand Down Expand Up @@ -746,6 +748,7 @@ psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) {
WCHAR * cmdline_buffer_wchar = NULL;
PUNICODE_STRING tmp = NULL;
DWORD string_size;
int ProcessCommandLineInformation = 60;

cmdline_buffer = calloc(ret_length, 1);
if (cmdline_buffer == NULL) {
Expand All @@ -756,27 +759,42 @@ psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) {
hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
if (hProcess == NULL)
goto error;

// get the right buf size
status = psutil_NtQueryInformationProcess(
hProcess,
60, // ProcessCommandLineInformation
ProcessCommandLineInformation,
NULL,
0,
&ret_length);
if (status != STATUS_BUFFER_OVERFLOW && \
status != STATUS_BUFFER_TOO_SMALL && \
status != STATUS_INFO_LENGTH_MISMATCH) {
PyErr_SetFromOSErrnoWithSyscall("NtQueryInformationProcess(0)");
goto error;
}

// get the cmdline
status = psutil_NtQueryInformationProcess(
hProcess,
ProcessCommandLineInformation,
cmdline_buffer,
ret_length,
&ret_length
);
if (! NT_SUCCESS(status)) {
PyErr_SetFromOSErrnoWithSyscall("NtQueryInformationProcess");
PyErr_SetFromOSErrnoWithSyscall("NtQueryInformationProcess(withlen)");
goto error;
}

// build the string
tmp = (PUNICODE_STRING)cmdline_buffer;
string_size = wcslen(tmp->Buffer) + 1;
cmdline_buffer_wchar = (WCHAR *)calloc(string_size, sizeof(WCHAR));

if (cmdline_buffer_wchar == NULL) {
PyErr_NoMemory();
goto error;
}

wcscpy_s(cmdline_buffer_wchar, string_size, tmp->Buffer);
*pdata = cmdline_buffer_wchar;
*psize = string_size * sizeof(WCHAR);
Expand Down

4 comments on commit 59e3c5e

@EccoTheFlintstone
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm, I think you have to choose :

either you call NtQueryInformationProcess twice like this, but you have to allocate the necessary memory between the 2 calls using the size in ret_length

or you keep the simple first version which had a 4096 bytes buffer and call NtQueryInformationProcess once and hope 4096 is enough

As is, you have a 4096 bytes buffer (enough for 2048 bytes minus a few of command line, might not be enough for edge cases), ask for the correct length, but don't do anything with the information (cmdline_buffer is already allocated at the beginning)

The first choice is, IMHO, the best one.
The version I provided was indeed crude and should have requested the correct length

@giampaolo
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right. I misread ProcessHacker code which indeed it allocates memory:
https://github.com/processhacker/processhacker/blob/master/phlib/native.c#L495
It uses RtlAllocateHeap though, which is kinda painful to implement. Unless you have other ideas to proper allocate memory I'd rather go back to your original implementation using a fixed size.

@EccoTheFlintstone
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's a really tiny fix actually (move the calloc after the first NtQueryInformationProcess call, set ret_length initially to 0 to be sure):

ULONG ret_length = 0;
...
    hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
    if (hProcess == NULL)
        goto error;

    // get the right buf size
    status = psutil_NtQueryInformationProcess(
        hProcess,
        ProcessCommandLineInformation,
        NULL,
        0,
        &ret_length);
       
    if (status != STATUS_BUFFER_OVERFLOW && \
            status != STATUS_BUFFER_TOO_SMALL && \
            status != STATUS_INFO_LENGTH_MISMATCH) {
        PyErr_SetFromOSErrnoWithSyscall("NtQueryInformationProcess(0)");
        goto error;
    }

    // allocate enough memory for the buffer
    cmdline_buffer = calloc(ret_length, 1);
    if (cmdline_buffer == NULL) {
        PyErr_NoMemory();
        goto error;
    }


    // get the cmdline
    status = psutil_NtQueryInformationProcess(
        hProcess,
        ProcessCommandLineInformation,
        cmdline_buffer,
        ret_length,
        &ret_length
    );
    if (! NT_SUCCESS(status)) {
        PyErr_SetFromOSErrnoWithSyscall("NtQueryInformationProcess(withlen)");
        goto error;
    }
   ...

@giampaolo
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in #1446. Thanks @EccoTheFlintstone.

Please sign in to comment.