diff options
author | Giampaolo Rodola <g.rodola@gmail.com> | 2020-05-05 12:55:32 -0700 |
---|---|---|
committer | Giampaolo Rodola <g.rodola@gmail.com> | 2020-05-05 12:55:32 -0700 |
commit | 91cd7fb53fe4cc93f837c7f543043c67b7b66674 (patch) | |
tree | f8f7a1034148f34897bd70a1d879d6f40977e0c1 | |
parent | 8ad353bc515d7765e7dc3dfaaa8f9f6db4f85aba (diff) | |
download | psutil-91cd7fb53fe4cc93f837c7f543043c67b7b66674.tar.gz |
Windows: refactor proc username(), split it in 2 functions
-rw-r--r-- | psutil/_psutil_windows.c | 111 | ||||
-rw-r--r-- | psutil/tests/test_testutils.py | 2 |
2 files changed, 59 insertions, 54 deletions
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c index 82fa518e..be4759ac 100644 --- a/psutil/_psutil_windows.c +++ b/psutil/_psutil_windows.c @@ -810,54 +810,34 @@ psutil_proc_open_files(PyObject *self, PyObject *args) { } -/* - * Return process username as a "DOMAIN//USERNAME" string. - */ -static PyObject * -psutil_proc_username(PyObject *self, PyObject *args) { - DWORD pid; - HANDLE processHandle = NULL; - HANDLE tokenHandle = NULL; - PTOKEN_USER user = NULL; - ULONG bufferSize; - WCHAR *name = NULL; - WCHAR *domainName = NULL; - ULONG nameSize; - ULONG domainNameSize; - SID_NAME_USE nameUse; - PyObject *py_username = NULL; - PyObject *py_domain = NULL; - PyObject *py_tuple = NULL; - - if (! PyArg_ParseTuple(args, _Py_PARSE_PID, &pid)) - return NULL; +static PTOKEN_USER +_psutil_user_token_from_pid(DWORD pid) { + HANDLE hProcess = NULL; + HANDLE hToken = NULL; + PTOKEN_USER userToken = NULL; + ULONG bufferSize = 0x100; - processHandle = psutil_handle_from_pid( - pid, PROCESS_QUERY_LIMITED_INFORMATION); - if (processHandle == NULL) + hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION); + if (hProcess == NULL) return NULL; - if (!OpenProcessToken(processHandle, TOKEN_QUERY, &tokenHandle)) { + if (!OpenProcessToken(hProcess, TOKEN_QUERY, &hToken)) { PyErr_SetFromOSErrnoWithSyscall("OpenProcessToken"); goto error; } - CloseHandle(processHandle); - processHandle = NULL; - // Get the user SID. - bufferSize = 0x100; while (1) { - user = malloc(bufferSize); - if (user == NULL) { + userToken = malloc(bufferSize); + if (userToken == NULL) { PyErr_NoMemory(); goto error; } - if (!GetTokenInformation(tokenHandle, TokenUser, user, bufferSize, + if (!GetTokenInformation(hToken, TokenUser, userToken, bufferSize, &bufferSize)) { if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) { - free(user); + free(userToken); continue; } else { @@ -868,15 +848,45 @@ psutil_proc_username(PyObject *self, PyObject *args) { break; } - CloseHandle(tokenHandle); - tokenHandle = NULL; + CloseHandle(hProcess); + CloseHandle(hToken); + return userToken; + +error: + if (hProcess != NULL) + CloseHandle(hProcess); + if (hToken != NULL) + CloseHandle(hToken); + return NULL; +} + + +/* + * Return process username as a "DOMAIN//USERNAME" string. + */ +static PyObject * +psutil_proc_username(PyObject *self, PyObject *args) { + DWORD pid; + PTOKEN_USER userToken = NULL; + WCHAR *userName = NULL; + WCHAR *domainName = NULL; + ULONG nameSize = 0x100; + ULONG domainNameSize = 0x100; + SID_NAME_USE nameUse; + PyObject *py_username = NULL; + PyObject *py_domain = NULL; + PyObject *py_tuple = NULL; + + if (! PyArg_ParseTuple(args, _Py_PARSE_PID, &pid)) + return NULL; + userToken = _psutil_user_token_from_pid(pid); + if (userToken == NULL) + return NULL; // resolve the SID to a name - nameSize = 0x100; - domainNameSize = 0x100; while (1) { - name = malloc(nameSize * sizeof(WCHAR)); - if (name == NULL) { + userName = malloc(nameSize * sizeof(WCHAR)); + if (userName == NULL) { PyErr_NoMemory(); goto error; } @@ -885,11 +895,11 @@ psutil_proc_username(PyObject *self, PyObject *args) { PyErr_NoMemory(); goto error; } - if (!LookupAccountSidW(NULL, user->User.Sid, name, &nameSize, + if (!LookupAccountSidW(NULL, userToken->User.Sid, userName, &nameSize, domainName, &domainNameSize, &nameUse)) { if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) { - free(name); + free(userName); free(domainName); continue; } @@ -904,7 +914,7 @@ psutil_proc_username(PyObject *self, PyObject *args) { py_domain = PyUnicode_FromWideChar(domainName, wcslen(domainName)); if (! py_domain) goto error; - py_username = PyUnicode_FromWideChar(name, wcslen(name)); + py_username = PyUnicode_FromWideChar(userName, wcslen(userName)); if (! py_username) goto error; py_tuple = Py_BuildValue("OO", py_domain, py_username); @@ -913,23 +923,18 @@ psutil_proc_username(PyObject *self, PyObject *args) { Py_DECREF(py_domain); Py_DECREF(py_username); - free(name); + free(userName); free(domainName); - free(user); - + free(userToken); return py_tuple; error: - if (processHandle != NULL) - CloseHandle(processHandle); - if (tokenHandle != NULL) - CloseHandle(tokenHandle); - if (name != NULL) - free(name); + if (userName != NULL) + free(userName); if (domainName != NULL) free(domainName); - if (user != NULL) - free(user); + if (userToken != NULL) + free(userToken); Py_XDECREF(py_domain); Py_XDECREF(py_username); Py_XDECREF(py_tuple); diff --git a/psutil/tests/test_testutils.py b/psutil/tests/test_testutils.py index 383b1470..6395b7c4 100644 --- a/psutil/tests/test_testutils.py +++ b/psutil/tests/test_testutils.py @@ -432,7 +432,7 @@ class TestMemLeakClass(TestMemoryLeak): box = [] self.assertRaisesRegex( - AssertionError, r"1 unclosed fd\(s\) or handle\(s\)", + AssertionError, r"unclosed fd\(s\) or handle\(s\)", self.execute, fun, times=5, warmup_times=5) |