port GetExtendedUpdTablewin-loadlib-refact

3 files changed, 18 insertions, 28 deletions

diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 18edd37a..afc9c223 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -1556,13 +1556,8 @@ static DWORD __GetExtendedTcpTable(ULONG address_family,
 }
 
 
-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,
+static DWORD __GetExtendedUdpTable(ULONG address_family,
                                    PVOID * data, DWORD * size)
 {
     // Due to other processes being active on the machine, it's possible
@@ -1574,8 +1569,8 @@ static DWORD __GetExtendedUdpTable(_GetExtendedUdpTable call,
     DWORD error = ERROR_INSUFFICIENT_BUFFER;
     *size = 0;
     *data = NULL;
-    error = call(NULL, size, FALSE, address_family,
-                 UDP_TABLE_OWNER_PID, 0);
+    error = psutil_GetExtendedUdpTable(
+        NULL, size, FALSE, address_family, UDP_TABLE_OWNER_PID, 0);
     while (error == ERROR_INSUFFICIENT_BUFFER || error == 0xC0000001)
     {
         *data = malloc(*size);
@@ -1583,8 +1578,8 @@ static DWORD __GetExtendedUdpTable(_GetExtendedUdpTable call,
             error = ERROR_NOT_ENOUGH_MEMORY;
             continue;
         }
-        error = call(*data, size, FALSE, address_family,
-                     UDP_TABLE_OWNER_PID, 0);
+        error = psutil_GetExtendedUdpTable(
+            *data, size, FALSE, address_family, UDP_TABLE_OWNER_PID, 0);
         if (error != NO_ERROR) {
             free(*data);
             *data = NULL;
@@ -1602,7 +1597,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {
     static long null_address[4] = { 0, 0, 0, 0 };
     unsigned long pid;
     int pid_return;
-    _GetExtendedUdpTable getExtendedUdpTable;
     PVOID table = NULL;
     DWORD tableSize;
     DWORD error;
@@ -1625,12 +1619,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {
     PyObject *_SOCK_STREAM = PyLong_FromLong((long)SOCK_STREAM);
     PyObject *_SOCK_DGRAM = PyLong_FromLong((long)SOCK_DGRAM);
 
-    // Import some functions.
-    getExtendedUdpTable = psutil_GetProcAddressFromLib(
-        "iphlpapi.dll", "GetExtendedUdpTable");
-    if (getExtendedUdpTable == NULL)
-        goto error;
-
     if (! PyArg_ParseTuple(args, "lOO", &pid, &py_af_filter, &py_type_filter))
         goto error;
 
@@ -1652,13 +1640,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {
         }
     }
 
-    if (getExtendedUdpTable == NULL) {
-        PyErr_SetString(PyExc_NotImplementedError,
-                        "feature not supported on this Windows version");
-        _psutil_conn_decref_objs();
-        return NULL;
-    }
-
     py_retlist = PyList_New(0);
     if (py_retlist == NULL) {
         _psutil_conn_decref_objs();
@@ -1868,8 +1849,7 @@ psutil_net_connections(PyObject *self, PyObject *args) {
         py_addr_tuple_local = NULL;
         py_addr_tuple_remote = NULL;
         tableSize = 0;
-        error = __GetExtendedUdpTable(getExtendedUdpTable,
-                                      AF_INET, &table, &tableSize);
+        error = __GetExtendedUdpTable(AF_INET, &table, &tableSize);
         if (error == ERROR_NOT_ENOUGH_MEMORY) {
             PyErr_NoMemory();
             goto error;
@@ -1942,8 +1922,7 @@ psutil_net_connections(PyObject *self, PyObject *args) {
         py_addr_tuple_local = NULL;
         py_addr_tuple_remote = NULL;
         tableSize = 0;
-        error = __GetExtendedUdpTable(getExtendedUdpTable,
-                                      AF_INET6, &table, &tableSize);
+        error = __GetExtendedUdpTable(AF_INET6, &table, &tableSize);
         if (error == ERROR_NOT_ENOUGH_MEMORY) {
             PyErr_NoMemory();
             goto error;
diff --git a/psutil/arch/windows/global.c b/psutil/arch/windows/global.c
index 9ead780a..b8bf8162 100644
--- a/psutil/arch/windows/global.c
+++ b/psutil/arch/windows/global.c
@@ -88,6 +88,11 @@ psutil_load_globals() {
     if (! psutil_GetExtendedTcpTable)
         return 1;
 
+    psutil_GetExtendedUdpTable = ps_GetProcAddressFromLib(
+        "iphlpapi.dll", "GetExtendedUdpTable");
+    if (! psutil_GetExtendedUdpTable)
+        return 1;
+
     // Optionals.
 
     psutil_GetActiveProcessorCount = ps_GetProcAddress(
diff --git a/psutil/arch/windows/global.h b/psutil/arch/windows/global.h
index 323d2c6d..69a1fd6c 100644
--- a/psutil/arch/windows/global.h
+++ b/psutil/arch/windows/global.h
@@ -19,6 +19,8 @@ typedef DWORD (CALLBACK *_GetActiveProcessorCount)(WORD);
 typedef ULONGLONG (CALLBACK *_GetTickCount64)(void);
 typedef DWORD (WINAPI * _GetExtendedTcpTable)(PVOID, PDWORD, BOOL, ULONG,
                                               TCP_TABLE_CLASS, ULONG);
+typedef DWORD (WINAPI * _GetExtendedUdpTable)(PVOID, PDWORD, BOOL, ULONG,
+                                              UDP_TABLE_CLASS, ULONG);
 
 
 _RtlIpv4AddressToStringA \
@@ -48,4 +50,8 @@ _GetTickCount64 \
 _GetExtendedTcpTable \
     psutil_GetExtendedTcpTable;
 
+_GetExtendedUdpTable \
+    psutil_GetExtendedUdpTable;
+
+
 int psutil_load_globals();