From 46f473c050fcfec26bea64a40db03e6286fce83a Mon Sep 17 00:00:00 2001
From: Giampaolo Rodola <g.rodola@gmail.com>
Date: Thu, 4 Apr 2019 08:18:31 -0700
Subject: define a utility function which converts NTSTATUS errors to Windows
 errors; fix ionice() which wasn't checking returning code

---
 HISTORY.rst                    |  2 ++
 psutil/_psutil_windows.c       | 15 ++++++++++++---
 psutil/arch/windows/global.c   | 30 ++++++++++++++++++++++++++++++
 psutil/arch/windows/global.h   |  4 ++++
 psutil/arch/windows/ntextapi.h |  4 ++++
 5 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index 79963855..e16b8f40 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -21,6 +21,8 @@
   exist.  (patch by Cedric Lamoriniere)
 - 1471_: [SunOS] Process name() and cmdline() can return SystemError.  (patch
   by Daniel Beer)
+- XXXX_: [Windows] wrong or absent error handling for private NTSTATUS Windows
+  APIs. Affected process methods are: ionice() (get and set), ...
 
 5.6.1
 =====
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 4dfae2d5..1a647c31 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -1859,19 +1859,24 @@ psutil_proc_io_priority_get(PyObject *self, PyObject *args) {
     long pid;
     HANDLE hProcess;
     DWORD IoPriority;
+    NTSTATUS status;
 
     if (! PyArg_ParseTuple(args, "l", &pid))
         return NULL;
     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
     if (hProcess == NULL)
         return NULL;
-    psutil_NtQueryInformationProcess(
+    status = psutil_NtQueryInformationProcess(
         hProcess,
         ProcessIoPriority,
         &IoPriority,
         sizeof(DWORD),
         NULL
     );
+    if (! NT_SUCCESS(status)) {
+        CloseHandle(hProcess);
+        return psutil_SetFromNTStatusErr(status);
+    }
     CloseHandle(hProcess);
     return Py_BuildValue("i", IoPriority);
 }
@@ -1885,6 +1890,7 @@ psutil_proc_io_priority_set(PyObject *self, PyObject *args) {
     long pid;
     DWORD prio;
     HANDLE hProcess;
+    NTSTATUS status;
     DWORD access = PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION;
 
     if (! PyArg_ParseTuple(args, "li", &pid, &prio))
@@ -1893,13 +1899,16 @@ psutil_proc_io_priority_set(PyObject *self, PyObject *args) {
     if (hProcess == NULL)
         return NULL;
 
-    psutil_NtSetInformationProcess(
+    status = psutil_NtSetInformationProcess(
         hProcess,
         ProcessIoPriority,
         (PVOID)&prio,
         sizeof(DWORD)
     );
-
+    if (! NT_SUCCESS(status)) {
+        CloseHandle(hProcess);
+        return psutil_SetFromNTStatusErr(status);
+    }
     CloseHandle(hProcess);
     Py_RETURN_NONE;
 }
diff --git a/psutil/arch/windows/global.c b/psutil/arch/windows/global.c
index 9ef92092..a5258491 100644
--- a/psutil/arch/windows/global.c
+++ b/psutil/arch/windows/global.c
@@ -18,6 +18,15 @@
 int PSUTIL_WINVER;
 SYSTEM_INFO PSUTIL_SYSTEM_INFO;
 
+#define FACILITY_WIN32 0x0007
+#define NT_FACILITY_MASK 0xfff
+#define NT_FACILITY_SHIFT 16
+#define NT_FACILITY(Status) \
+    ((((ULONG)(Status)) >> NT_FACILITY_SHIFT) & NT_FACILITY_MASK)
+#define NT_NTWIN32(status) (NT_FACILITY(Status) == FACILITY_WIN32)
+#define WIN32_FROM_NTSTATUS(Status) (((ULONG)(Status)) & 0xffff)
+
+
 // A wrapper around GetModuleHandle and GetProcAddress.
 PVOID
 psutil_GetProcAddress(LPCSTR libname, LPCSTR procname) {
@@ -60,6 +69,22 @@ psutil_GetProcAddressFromLib(LPCSTR libname, LPCSTR procname) {
 }
 
 
+/*
+ * Convert a NTSTATUS value to a Win32 error code and set the proper
+ * Python exception.
+ */
+PVOID
+psutil_SetFromNTStatusErr(NTSTATUS Status) {
+    ULONG err;
+
+    if (NT_NTWIN32(Status))
+        err = WIN32_FROM_NTSTATUS(Status);
+    else
+        err = psutil_RtlNtStatusToDosErrorNoTeb(Status);
+    return PyErr_SetFromWindowsErr(err);
+}
+
+
 static int
 psutil_loadlibs() {
     /*
@@ -127,6 +152,11 @@ psutil_loadlibs() {
     if (! psutil_NtQueryVirtualMemory)
         return 1;
 
+    psutil_RtlNtStatusToDosErrorNoTeb = psutil_GetProcAddressFromLib(
+        "ntdll", "RtlNtStatusToDosErrorNoTeb");
+    if (! psutil_RtlNtStatusToDosErrorNoTeb)
+        return 1;
+
     /*
      * Optional.
      */
diff --git a/psutil/arch/windows/global.h b/psutil/arch/windows/global.h
index fb24bac9..31e20293 100644
--- a/psutil/arch/windows/global.h
+++ b/psutil/arch/windows/global.h
@@ -23,6 +23,7 @@ extern SYSTEM_INFO PSUTIL_SYSTEM_INFO;
 int psutil_load_globals();
 PVOID psutil_GetProcAddress(LPCSTR libname, LPCSTR procname);
 PVOID psutil_GetProcAddressFromLib(LPCSTR libname, LPCSTR procname);
+PVOID psutil_SetFromNTStatusErr(NTSTATUS Status);
 
 _NtQuerySystemInformation \
     psutil_NtQuerySystemInformation;
@@ -71,3 +72,6 @@ _NtResumeProcess \
 
 _NtQueryVirtualMemory \
     psutil_NtQueryVirtualMemory;
+
+_RtlNtStatusToDosErrorNoTeb \
+    psutil_RtlNtStatusToDosErrorNoTeb;
diff --git a/psutil/arch/windows/ntextapi.h b/psutil/arch/windows/ntextapi.h
index 178f9866..b6f23d99 100644
--- a/psutil/arch/windows/ntextapi.h
+++ b/psutil/arch/windows/ntextapi.h
@@ -500,4 +500,8 @@ typedef NTSTATUS (NTAPI *_NtQueryVirtualMemory) (
     PSIZE_T ReturnLength
 );
 
+typedef ULONG (WINAPI *_RtlNtStatusToDosErrorNoTeb) (
+    NTSTATUS status
+);
+
 #endif // __NTEXTAPI_H__
-- 
cgit v1.2.1