Merge pull request #835 from jhomann/master

[Windows] Fix for net_io_counters wrapping after 4.3GB due to MIB_IFROW
giampaolo · Jun 12, 2016 · 49f8555 · 49f8555
2 parents 029fbdf + 7769dbc
commit 49f8555
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 3 deletions.
diff --git a/CREDITS b/CREDITS
@@ -386,3 +386,7 @@ I: 776
 N: Farhan Khan
 E: [email protected]
 I: 823
+
+N: Jake Omann
+E: https://github.com/jhomann
+I: 816
diff --git a/HISTORY.rst b/HISTORY.rst
@@ -9,6 +9,8 @@ Bug tracker at https://github.com/giampaolo/psutil/issues
 - #823: [NetBSD] virtual_memory() raises TypeError on Python 3.
 - #829: [UNIX] psutil.disk_usage() percent field takes root reserved space
   into account.
+- #816: [Windows] fixed net_io_counter() values wrapping after 4.3GB in
+  Windows Vista (NT 6.0) and above using 64bit values from newer win APIs.
 
 
 4.2.0 - 2016-05-14

diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
@@ -18,9 +18,11 @@
 #include <tchar.h>
 #include <tlhelp32.h>
 #include <winsock2.h>
+#if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
+#include <ws2tcpip.h>
+#endif
 #include <iphlpapi.h>
 #include <wtsapi32.h>
-#include <ws2tcpip.h>
 #include <Winsvc.h>
 
 // Link with Iphlpapi.lib
@@ -85,6 +87,7 @@ typedef struct _DISK_PERFORMANCE_WIN_2008 {
 
 // --- network connections mingw32 support
 #ifndef _IPRTRMIB_H
+#if (_WIN32_WINNT < 0x0600) // Windows XP
 typedef struct _MIB_TCP6ROW_OWNER_PID {
     UCHAR ucLocalAddr[16];
     DWORD dwLocalScopeId;
@@ -101,6 +104,7 @@ typedef struct _MIB_TCP6TABLE_OWNER_PID {
     MIB_TCP6ROW_OWNER_PID table[ANY_SIZE];
 } MIB_TCP6TABLE_OWNER_PID, *PMIB_TCP6TABLE_OWNER_PID;
 #endif
+#endif
 
 #ifndef __IPHLPAPI_H__
 typedef struct in6_addr {
@@ -128,6 +132,7 @@ typedef struct _MIB_UDPTABLE_OWNER_PID {
 } MIB_UDPTABLE_OWNER_PID, *PMIB_UDPTABLE_OWNER_PID;
 #endif
 
+#if (_WIN32_WINNT < 0x0600) // Windows XP
 typedef struct _MIB_UDP6ROW_OWNER_PID {
     UCHAR ucLocalAddr[16];
     DWORD dwLocalScopeId;
@@ -139,6 +144,7 @@ typedef struct _MIB_UDP6TABLE_OWNER_PID {
     DWORD dwNumEntries;
     MIB_UDP6ROW_OWNER_PID table[ANY_SIZE];
 } MIB_UDP6TABLE_OWNER_PID, *PMIB_UDP6TABLE_OWNER_PID;
+#endif
 
 PIP_ADAPTER_ADDRESSES
 psutil_get_nic_addresses() {
@@ -1636,7 +1642,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {
     }
 
     // TCP IPv6
-
     if ((PySequence_Contains(py_af_filter, _AF_INET6) == 1) &&
             (PySequence_Contains(py_type_filter, _SOCK_STREAM) == 1))
     {
@@ -2168,7 +2173,13 @@ psutil_disk_usage(PyObject *self, PyObject *args) {
 static PyObject *
 psutil_net_io_counters(PyObject *self, PyObject *args) {
     DWORD dwRetVal = 0;
+
+#if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
+    MIB_IF_ROW2 *pIfRow = NULL;
+#else // Windows XP
     MIB_IFROW *pIfRow = NULL;
+#endif
+
     PIP_ADAPTER_ADDRESSES pAddresses = NULL;
     PIP_ADAPTER_ADDRESSES pCurrAddresses = NULL;
     PyObject *py_retdict = PyDict_New();
@@ -2185,20 +2196,44 @@ psutil_net_io_counters(PyObject *self, PyObject *args) {
     while (pCurrAddresses) {
         py_nic_name = NULL;
         py_nic_info = NULL;
+
+#if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
+        pIfRow = (MIB_IF_ROW2 *) malloc(sizeof(MIB_IF_ROW2));
+#else // Windows XP
         pIfRow = (MIB_IFROW *) malloc(sizeof(MIB_IFROW));
+#endif
 
         if (pIfRow == NULL) {
             PyErr_NoMemory();
             goto error;
         }
 
+#if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
+        SecureZeroMemory((PVOID)pIfRow, sizeof(MIB_IF_ROW2));
+        pIfRow->InterfaceIndex = pCurrAddresses->IfIndex;
+        dwRetVal = GetIfEntry2(pIfRow);
+#else // Windows XP
         pIfRow->dwIndex = pCurrAddresses->IfIndex;
         dwRetVal = GetIfEntry(pIfRow);
+#endif
+
         if (dwRetVal != NO_ERROR) {
-            PyErr_SetString(PyExc_RuntimeError, "GetIfEntry() failed.");
+            PyErr_SetString(PyExc_RuntimeError,
+                            "GetIfEntry() or GetIfEntry2() failed.");
             goto error;
         }
 
+#if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
+        py_nic_info = Py_BuildValue("(KKKKKKKK)",
+                                    pIfRow->OutOctets,
+                                    pIfRow->InOctets,
+                                    pIfRow->OutUcastPkts,
+                                    pIfRow->InUcastPkts,
+                                    pIfRow->InErrors,
+                                    pIfRow->OutErrors,
+                                    pIfRow->InDiscards,
+                                    pIfRow->OutDiscards);
+#else // Windows XP
         py_nic_info = Py_BuildValue("(kkkkkkkk)",
                                     pIfRow->dwOutOctets,
                                     pIfRow->dwInOctets,
@@ -2208,6 +2243,8 @@ psutil_net_io_counters(PyObject *self, PyObject *args) {
                                     pIfRow->dwOutErrors,
                                     pIfRow->dwInDiscards,
                                     pIfRow->dwOutDiscards);
+#endif
+
         if (!py_nic_info)
             goto error;