Skip to content

Commit

Permalink
Merge pull request #880 from AndreLouisCaron/win32-tcp-table-race-con…
Browse files Browse the repository at this point in the history 
…dition

Fixes race condition getting IPv4 TCP table on Windows.
  • Loading branch information
@giampaolo
giampaolo authored Sep 11, 2016
2 parents 273ea68 + 7fb15c6 commit 601f959
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 30 deletions.
6 changes: 6 additions & 0 deletions CREDITS
View file
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
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
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

0 comments on commit 601f959

Please sign in to comment.