Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes race condition getting IPv4 TCP table on Windows. #880

Merged
merged 2 commits into from
Sep 11, 2016
Merged

Fixes race condition getting IPv4 TCP table on Windows. #880

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1b9a35c
Fixes race condition getting TCP and UDP tables on Windows.
Sep 1, 2016
7fb15c6
Releases GIL while fetching TCP and UDP tables on Windows.
AndreLouisCaron Sep 3, 2016
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions CREDITS
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -402,3 +402,9 @@ I: 870
N: Yago Jesus
W: https://github.com/YJesus
I: 798

N: Andre Caron
C: Montreal, QC, Canada
E: [email protected]
W: https://github.com/AndreLouisCaron
I: 880
1 change: 1 addition & 0 deletions HISTORY.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ Bug tracker at https://github.com/giampaolo/psutil/issues
- #869: [Windows] Process.wait() may raise TimeoutExpired with wrong timeout
unit (ms instead of sec).
- #870: [Windows] Handle leak inside psutil_get_process_data.
- #880: [Windows] Handle race condition inside psutil_net_connections.


4.3.0 - 2016-06-18
Expand Down
153 changes: 123 additions & 30 deletions psutil/_psutil_windows.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1469,6 +1469,88 @@ psutil_proc_username(PyObject *self, PyObject *args) {
}


typedef DWORD (WINAPI * _GetExtendedTcpTable)(PVOID, PDWORD, BOOL, ULONG,
TCP_TABLE_CLASS, ULONG);


// https://msdn.microsoft.com/library/aa365928.aspx
static DWORD __GetExtendedTcpTable(_GetExtendedTcpTable call,
ULONG address_family,
PVOID * data, DWORD * size)
{
// Due to other processes being active on the machine, it's possible
// that the size of the table increases between the moment where we
// query the size and the moment where we query the data. Therefore, it's
// important to call this in a loop to retry if that happens.
//
// Also, since we may loop a theoretically unbounded number of times here,
// release the GIL while we're doing this.
DWORD error = ERROR_INSUFFICIENT_BUFFER;
*size = 0;
*data = NULL;
Py_BEGIN_ALLOW_THREADS;
error = call(NULL, size, FALSE, address_family,
TCP_TABLE_OWNER_PID_ALL, 0);
while (error == ERROR_INSUFFICIENT_BUFFER)
{
*data = malloc(*size);
if (*data == NULL) {
error = ERROR_NOT_ENOUGH_MEMORY;
continue;
}
error = call(*data, size, FALSE, address_family,
TCP_TABLE_OWNER_PID_ALL, 0);
if (error != NO_ERROR) {
free(*data);
*data = NULL;
}
}
Py_END_ALLOW_THREADS;
return error;
}


typedef DWORD (WINAPI * _GetExtendedUdpTable)(PVOID, PDWORD, BOOL, ULONG,
UDP_TABLE_CLASS, ULONG);


// https://msdn.microsoft.com/library/aa365930.aspx
static DWORD __GetExtendedUdpTable(_GetExtendedUdpTable call,
ULONG address_family,
PVOID * data, DWORD * size)
{
// Due to other processes being active on the machine, it's possible
// that the size of the table increases between the moment where we
// query the size and the moment where we query the data. Therefore, it's
// important to call this in a loop to retry if that happens.
//
// Also, since we may loop a theoretically unbounded number of times here,
// release the GIL while we're doing this.
DWORD error = ERROR_INSUFFICIENT_BUFFER;
*size = 0;
*data = NULL;
Py_BEGIN_ALLOW_THREADS;
error = call(NULL, size, FALSE, address_family,
UDP_TABLE_OWNER_PID, 0);
while (error == ERROR_INSUFFICIENT_BUFFER)
{
*data = malloc(*size);
if (*data == NULL) {
error = ERROR_NOT_ENOUGH_MEMORY;
continue;
}
error = call(*data, size, FALSE, address_family,
UDP_TABLE_OWNER_PID, 0);
if (error != NO_ERROR) {
free(*data);
*data = NULL;
}
}
Py_END_ALLOW_THREADS;
return error;
}


/*
* Return a list of network connections opened by a process
*/
Expand All @@ -1480,14 +1562,11 @@ psutil_net_connections(PyObject *self, PyObject *args) {
_RtlIpv4AddressToStringA rtlIpv4AddressToStringA;
typedef PSTR (NTAPI * _RtlIpv6AddressToStringA)(struct in6_addr *, PSTR);
_RtlIpv6AddressToStringA rtlIpv6AddressToStringA;
typedef DWORD (WINAPI * _GetExtendedTcpTable)(PVOID, PDWORD, BOOL, ULONG,
TCP_TABLE_CLASS, ULONG);
_GetExtendedTcpTable getExtendedTcpTable;
typedef DWORD (WINAPI * _GetExtendedUdpTable)(PVOID, PDWORD, BOOL, ULONG,
UDP_TABLE_CLASS, ULONG);
_GetExtendedUdpTable getExtendedUdpTable;
PVOID table = NULL;
DWORD tableSize;
DWORD error;
PMIB_TCPTABLE_OWNER_PID tcp4Table;
PMIB_UDPTABLE_OWNER_PID udp4Table;
PMIB_TCP6TABLE_OWNER_PID tcp6Table;
Expand Down Expand Up @@ -1570,17 +1649,15 @@ psutil_net_connections(PyObject *self, PyObject *args) {
py_addr_tuple_local = NULL;
py_addr_tuple_remote = NULL;
tableSize = 0;
getExtendedTcpTable(NULL, &tableSize, FALSE, AF_INET,
TCP_TABLE_OWNER_PID_ALL, 0);

table = malloc(tableSize);
if (table == NULL) {
error = __GetExtendedTcpTable(getExtendedTcpTable,
AF_INET, &table, &tableSize);
if (error == ERROR_NOT_ENOUGH_MEMORY) {
PyErr_NoMemory();
goto error;
}

if (getExtendedTcpTable(table, &tableSize, FALSE, AF_INET,
TCP_TABLE_OWNER_PID_ALL, 0) == 0)
if (error == NO_ERROR)
{
tcp4Table = table;

Expand Down Expand Up @@ -1650,8 +1727,14 @@ psutil_net_connections(PyObject *self, PyObject *args) {
Py_DECREF(py_conn_tuple);
}
}
else {
PyErr_SetFromWindowsErr(error);
goto error;
}

free(table);
table = NULL;
tableSize = 0;
}

// TCP IPv6
Expand All @@ -1663,17 +1746,15 @@ psutil_net_connections(PyObject *self, PyObject *args) {
py_addr_tuple_local = NULL;
py_addr_tuple_remote = NULL;
tableSize = 0;
getExtendedTcpTable(NULL, &tableSize, FALSE, AF_INET6,
TCP_TABLE_OWNER_PID_ALL, 0);

table = malloc(tableSize);
if (table == NULL) {
error = __GetExtendedTcpTable(getExtendedTcpTable,
AF_INET6, &table, &tableSize);
if (error == ERROR_NOT_ENOUGH_MEMORY) {
PyErr_NoMemory();
goto error;
}

if (getExtendedTcpTable(table, &tableSize, FALSE, AF_INET6,
TCP_TABLE_OWNER_PID_ALL, 0) == 0)
if (error == NO_ERROR)
{
tcp6Table = table;

Expand Down Expand Up @@ -1743,8 +1824,14 @@ psutil_net_connections(PyObject *self, PyObject *args) {
Py_DECREF(py_conn_tuple);
}
}
else {
PyErr_SetFromWindowsErr(error);
goto error;
}

free(table);
table = NULL;
tableSize = 0;
}

// UDP IPv4
Expand All @@ -1757,17 +1844,14 @@ psutil_net_connections(PyObject *self, PyObject *args) {
py_addr_tuple_local = NULL;
py_addr_tuple_remote = NULL;
tableSize = 0;
getExtendedUdpTable(NULL, &tableSize, FALSE, AF_INET,
UDP_TABLE_OWNER_PID, 0);

table = malloc(tableSize);
if (table == NULL) {
error = __GetExtendedUdpTable(getExtendedUdpTable,
AF_INET, &table, &tableSize);
if (error == ERROR_NOT_ENOUGH_MEMORY) {
PyErr_NoMemory();
goto error;
}

if (getExtendedUdpTable(table, &tableSize, FALSE, AF_INET,
UDP_TABLE_OWNER_PID, 0) == 0)
if (error == NO_ERROR)
{
udp4Table = table;

Expand Down Expand Up @@ -1814,8 +1898,14 @@ psutil_net_connections(PyObject *self, PyObject *args) {
Py_DECREF(py_conn_tuple);
}
}
else {
PyErr_SetFromWindowsErr(error);
goto error;
}

free(table);
table = NULL;
tableSize = 0;
}

// UDP IPv6
Expand All @@ -1828,17 +1918,14 @@ psutil_net_connections(PyObject *self, PyObject *args) {
py_addr_tuple_local = NULL;
py_addr_tuple_remote = NULL;
tableSize = 0;
getExtendedUdpTable(NULL, &tableSize, FALSE,
AF_INET6, UDP_TABLE_OWNER_PID, 0);

table = malloc(tableSize);
if (table == NULL) {
error = __GetExtendedUdpTable(getExtendedUdpTable,
AF_INET6, &table, &tableSize);
if (error == ERROR_NOT_ENOUGH_MEMORY) {
PyErr_NoMemory();
goto error;
}

if (getExtendedUdpTable(table, &tableSize, FALSE, AF_INET6,
UDP_TABLE_OWNER_PID, 0) == 0)
if (error == NO_ERROR)
{
udp6Table = table;

Expand Down Expand Up @@ -1884,8 +1971,14 @@ psutil_net_connections(PyObject *self, PyObject *args) {
Py_DECREF(py_conn_tuple);
}
}
else {
PyErr_SetFromWindowsErr(error);
goto error;
}

free(table);
table = NULL;
tableSize = 0;
}

_psutil_conn_decref_objs();
Expand Down