Skip to content

Commit

Permalink
Merge 08b4099 into 96cc7ea
Browse files Browse the repository at this point in the history
  • Loading branch information
@giampaolo
giampaolo committed Feb 1, 2020
2 parents 96cc7ea + 08b4099 commit 8ba1992
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 93 deletions.
164 changes: 72 additions & 92 deletions psutil/arch/windows/socks.c
View file
Expand Up @@ -16,87 +16,81 @@


#define BYTESWAP_USHORT(x) ((((USHORT)(x) << 8) | ((USHORT)(x) >> 8)) & 0xffff)
typedef DWORD (WINAPI * TYPE_GetExtendedTcpTable)();
typedef DWORD (WINAPI * TYPE_GetExtendedUdpTable)();


static DWORD __GetExtendedTcpTable(TYPE_GetExtendedTcpTable call,
ULONG 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.
// See https://github.com/giampaolo/psutil/pull/1335 concerning 0xC0000001 error
// and https://github.com/giampaolo/psutil/issues/1294
DWORD error = ERROR_INSUFFICIENT_BUFFER;
*size = 0;
*data = NULL;
error = call(NULL, size, FALSE, family, TCP_TABLE_OWNER_PID_ALL, 0);
while (error == ERROR_INSUFFICIENT_BUFFER || error == 0xC0000001)
{
*data = malloc(*size);
if (*data == NULL) {
error = ERROR_NOT_ENOUGH_MEMORY;
continue;
#define STATUS_UNSUCCESSFUL 0xC0000001


// Note about GetExtended[Tcp|Udp]Table syscalls: 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 we call them in a loop to retry if
// that happens. See:
// https://github.com/giampaolo/psutil/pull/1335
// https://github.com/giampaolo/psutil/issues/1294


static PVOID __GetExtendedTcpTable(ULONG family) {
DWORD err;
PVOID table;
ULONG size = 0;

while (1) {
// get table size
GetExtendedTcpTable(NULL, &size, FALSE, family,
TCP_TABLE_OWNER_PID_ALL, 0);

table = malloc(size);
if (table == NULL) {
PyErr_NoMemory();
return NULL;
}
error = call(*data, size, FALSE, family, TCP_TABLE_OWNER_PID_ALL, 0);
if (error != NO_ERROR) {
free(*data);
*data = NULL;

// get connections
err = GetExtendedTcpTable(table, &size, FALSE, family,
TCP_TABLE_OWNER_PID_ALL, 0);
if (err == NO_ERROR)
return table;

free(table);
if (err == ERROR_INSUFFICIENT_BUFFER || err == STATUS_UNSUCCESSFUL) {
psutil_debug("GetExtendedTcpTable: retry with different bufsize");
continue;
}
}
if (error == ERROR_NOT_ENOUGH_MEMORY) {
PyErr_NoMemory();
return 1;
}
if (error != NO_ERROR) {
PyErr_SetString(PyExc_RuntimeError, "GetExtendedTcpTable failed");
return 1;
return NULL;
}
return error;
}


static DWORD __GetExtendedUdpTable(TYPE_GetExtendedUdpTable call,
ULONG 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.
// See https://github.com/giampaolo/psutil/pull/1335 concerning 0xC0000001 error
// and https://github.com/giampaolo/psutil/issues/1294
DWORD error = ERROR_INSUFFICIENT_BUFFER;
*size = 0;
*data = NULL;
error = call(NULL, size, FALSE, family, UDP_TABLE_OWNER_PID, 0);
while (error == ERROR_INSUFFICIENT_BUFFER || error == 0xC0000001)
{
*data = malloc(*size);
if (*data == NULL) {
error = ERROR_NOT_ENOUGH_MEMORY;
continue;
static PVOID __GetExtendedUdpTable(ULONG family) {
DWORD err;
PVOID table;
ULONG size = 0;

while (1) {
// get table size
GetExtendedUdpTable(NULL, &size, FALSE, family,
UDP_TABLE_OWNER_PID, 0);

table = malloc(size);
if (table == NULL) {
PyErr_NoMemory();
return NULL;
}
error = call(*data, size, FALSE, family, UDP_TABLE_OWNER_PID, 0);
if (error != NO_ERROR) {
free(*data);
*data = NULL;

// get connections
err = GetExtendedUdpTable(table, &size, FALSE, family,
UDP_TABLE_OWNER_PID, 0);
if (err == NO_ERROR)
return table;

free(table);
if (err == ERROR_INSUFFICIENT_BUFFER || err == STATUS_UNSUCCESSFUL) {
psutil_debug("GetExtendedUdpTable: retry with different bufsize");
continue;
}
}
if (error == ERROR_NOT_ENOUGH_MEMORY) {
PyErr_NoMemory();
return 1;
}
if (error != NO_ERROR) {
PyErr_SetString(PyExc_RuntimeError, "GetExtendedUdpTable failed");
return 1;
return NULL;
}
return 0;
}


Expand All @@ -116,8 +110,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {
DWORD pid;
int pid_return;
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 @@ -176,11 +168,9 @@ psutil_net_connections(PyObject *self, PyObject *args) {
py_conn_tuple = NULL;
py_addr_tuple_local = NULL;
py_addr_tuple_remote = NULL;
tableSize = 0;

error = __GetExtendedTcpTable(GetExtendedTcpTable,
AF_INET, &table, &tableSize);
if (error != 0)
table = __GetExtendedTcpTable(AF_INET);
if (table == NULL)
goto error;
tcp4Table = table;
for (i = 0; i < tcp4Table->dwNumEntries; i++) {
Expand Down Expand Up @@ -250,7 +240,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {

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

// TCP IPv6
Expand All @@ -262,11 +251,9 @@ psutil_net_connections(PyObject *self, PyObject *args) {
py_conn_tuple = NULL;
py_addr_tuple_local = NULL;
py_addr_tuple_remote = NULL;
tableSize = 0;

error = __GetExtendedTcpTable(GetExtendedTcpTable,
AF_INET6, &table, &tableSize);
if (error != 0)
table = __GetExtendedTcpTable(AF_INET6);
if (table == NULL)
goto error;
tcp6Table = table;
for (i = 0; i < tcp6Table->dwNumEntries; i++)
Expand Down Expand Up @@ -337,7 +324,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {

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

// UDP IPv4
Expand All @@ -349,10 +335,8 @@ psutil_net_connections(PyObject *self, PyObject *args) {
py_conn_tuple = NULL;
py_addr_tuple_local = NULL;
py_addr_tuple_remote = NULL;
tableSize = 0;
error = __GetExtendedUdpTable(GetExtendedUdpTable,
AF_INET, &table, &tableSize);
if (error != 0)
table = __GetExtendedUdpTable(AF_INET);
if (table == NULL)
goto error;
udp4Table = table;
for (i = 0; i < udp4Table->dwNumEntries; i++)
Expand Down Expand Up @@ -400,7 +384,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {

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

// UDP IPv6
Expand All @@ -413,10 +396,8 @@ psutil_net_connections(PyObject *self, PyObject *args) {
py_conn_tuple = NULL;
py_addr_tuple_local = NULL;
py_addr_tuple_remote = NULL;
tableSize = 0;
error = __GetExtendedUdpTable(GetExtendedUdpTable,
AF_INET6, &table, &tableSize);
if (error != 0)
table = __GetExtendedUdpTable(AF_INET6);
if (table == NULL)
goto error;
udp6Table = table;
for (i = 0; i < udp6Table->dwNumEntries; i++) {
Expand Down Expand Up @@ -463,7 +444,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {

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

psutil_conn_decref_objs();
Expand Down
2 changes: 1 addition & 1 deletion scripts/internal/print_access_denied.py
View file
Expand Up @@ -76,7 +76,7 @@ def main():
templ = "%-20s %-5s %-9s %s"
s = templ % ("API", "AD", "Percent", "Outcome")
print(hilite(s, ok=None, bold=True))
for methname, ads in sorted(d.items(), key=lambda x: x[1]):
for methname, ads in sorted(d.items(), key=lambda x: (x[1], x[0])):
perc = (ads / tot_procs) * 100
outcome = "SUCCESS" if not ads else "ACCESS DENIED"
s = templ % (methname, ads, "%6.1f%%" % perc, outcome)
Expand Down

0 comments on commit 8ba1992

Please sign in to comment.