port NtQueryObject

3 files changed, 24 insertions, 20 deletions

diff --git a/psutil/arch/windows/global.c b/psutil/arch/windows/global.c
index 950c291d..f76d41d5 100644
--- a/psutil/arch/windows/global.c
+++ b/psutil/arch/windows/global.c
@@ -73,6 +73,11 @@ psutil_load_globals() {
     if (! psutil_WinStationQueryInformationW)
         return 1;
 
+    psutil_NtQueryObject = ps_GetProcAddressFromLib(
+        "ntdll.dll", "NtQueryObject");
+    if (! psutil_NtQueryObject)
+        return 1;
+
     psutil_rtlIpv4AddressToStringA = ps_GetProcAddressFromLib(
         "ntdll.dll", "RtlIpv4AddressToStringA");
     if (! psutil_rtlIpv4AddressToStringA)
@@ -93,7 +98,8 @@ psutil_load_globals() {
     if (! psutil_GetExtendedUdpTable)
         return 1;
 
-    // Optionals.
+    // Optional.
+
     psutil_GetTickCount64 = ps_GetProcAddress(
         "kernel32", "GetTickCount64");
 
diff --git a/psutil/arch/windows/global.h b/psutil/arch/windows/global.h
index da2d7293..b4d7f7d7 100644
--- a/psutil/arch/windows/global.h
+++ b/psutil/arch/windows/global.h
@@ -16,6 +16,14 @@ typedef DWORD (WINAPI * _GetExtendedUdpTable)(PVOID, PDWORD, BOOL, ULONG,
                                               UDP_TABLE_CLASS, ULONG);
 typedef DWORD (CALLBACK *_GetActiveProcessorCount)(WORD);
 typedef ULONGLONG (CALLBACK *_GetTickCount64)(void);
+typedef NTSTATUS (NTAPI *_NtQueryObject)(
+    HANDLE ObjectHandle,
+    ULONG ObjectInformationClass,
+    PVOID ObjectInformation,
+    ULONG ObjectInformationLength,
+    PULONG ReturnLength
+);
+
 
 // probably unnecessary?
 /*
@@ -54,5 +62,7 @@ _GetActiveProcessorCount \
 _GetTickCount64 \
     psutil_GetTickCount64;
 
+_NtQueryObject \
+    psutil_NtQueryObject;
 
 int psutil_load_globals();
diff --git a/psutil/arch/windows/process_handles.c b/psutil/arch/windows/process_handles.c
index 33a9c778..4910f4ef 100644
--- a/psutil/arch/windows/process_handles.c
+++ b/psutil/arch/windows/process_handles.c
@@ -10,8 +10,6 @@
 #include "global.h"
 #include "../../_psutil_common.h"
 
-static _NtQueryObject __NtQueryObject = NULL;
-
 CRITICAL_SECTION g_cs;
 BOOL g_initialized = FALSE;
 NTSTATUS g_status;
@@ -45,10 +43,6 @@ psutil_get_open_files_init(BOOL threaded) {
     if (g_initialized == TRUE)
         return;
 
-    // Resolve the Windows API calls
-    __NtQueryObject = psutil_GetProcAddressFromLib(
-        "ntdll.dll", "NtQueryObject");
-
     // Create events for signalling work between threads
     if (threaded == TRUE) {
         g_hEvtStart = CreateEvent(NULL, FALSE, FALSE, NULL);
@@ -80,8 +74,7 @@ psutil_get_open_files_ntqueryobject(long dwPid, HANDLE hProcess) {
     // to psutil_get_open_files() is running
     EnterCriticalSection(&g_cs);
 
-    if (__NtQueryObject == NULL ||
-        g_hEvtStart == NULL ||
+    if (g_hEvtStart == NULL ||
         g_hEvtFinish == NULL)
 
     {
@@ -307,11 +300,12 @@ psutil_NtQueryObjectThread(LPVOID lpvParam) {
     while (TRUE) {
         WaitForSingleObject(g_hEvtStart, INFINITE);
 
-        g_status = __NtQueryObject(g_hFile,
-                                   ObjectNameInformation,
-                                   g_pNameBuffer,
-                                   g_dwSize,
-                                   &g_dwLength);
+        g_status = psutil_NtQueryObject(
+            g_hFile,
+            ObjectNameInformation,
+            g_pNameBuffer,
+            g_dwSize,
+            &g_dwLength);
         SetEvent(g_hEvtFinish);
     }
 }
@@ -337,12 +331,6 @@ psutil_get_open_files_getmappedfilename(long dwPid, HANDLE hProcess) {
     if (g_initialized == FALSE)
         psutil_get_open_files_init(FALSE);
 
-    if (__NtQueryObject == NULL) {
-        PyErr_SetFromWindowsErr(0);
-        error = TRUE;
-        goto cleanup;
-    }
-
     // Py_BuildValue raises an exception if NULL is returned
     py_retlist = PyList_New(0);
     if (py_retlist == NULL) {