Translate NTSTATUS into the last-error slot instead of discarding it

MemLib/windows.py

@@ -154,6 +154,61 @@ def GetLastError() -> int:
     """
     return _GetLastError()
 
+# noinspection PyPep8Naming
+# pylint: disable=invalid-name
+def SetLastError(error_code: int) -> None:
+    """
+    Sets the last-error code for the calling thread.
+
+    Args:
+        error_code (int): The last-error code to publish.
+
+    See also:
+        https://learn.microsoft.com/en-us/windows/win32/api/errhandlingapi/nf-errhandlingapi-setlasterror
+    """
+    _SetLastError(error_code)
+
+# noinspection PyPep8Naming
+# pylint: disable=invalid-name
+def RtlNtStatusToDosError(nt_status: int) -> int:
+    """
+    Converts an NTSTATUS code into the equivalent Win32 error code.
+
+    Args:
+        nt_status (int): The NTSTATUS value to translate.
+
+    Returns:
+        int: The matching Win32 error code, or ERROR_MR_MID_NOT_FOUND (317) if the
+            status has no Win32 equivalent.
+
+    See also:
+        https://learn.microsoft.com/en-us/windows/win32/api/winternl/nf-winternl-rtlntstatustodoserror
+    """
+    return _RtlNtStatusToDosError(nt_status)
+
+def _nt_ok(nt_status: int) -> bool:
+    """
+    Converts an NTSTATUS return value into a bool, publishing the failure via SetLastError.
+
+    NTSTATUS routines report failure in their return value and do not call SetLastError.
+    Collapsing one to a bool therefore discards the only copy of the error: a following
+    ``Win32Exception()`` reads whatever the *previous* call left in the thread-local
+    last-error slot, which is 0 on a fresh thread ("The operation completed successfully")
+    and a stale, unrelated code otherwise. Translating the status keeps the real failure
+    reachable at the call site.
+
+    Args:
+        nt_status (int): The NTSTATUS value returned by the routine.
+
+    Returns:
+        bool: True if the status is STATUS_SUCCESS, otherwise False.
+    """
+    if nt_status == STATUS_SUCCESS:
+        return True
+
+    SetLastError(RtlNtStatusToDosError(nt_status))
+    return False
+
 # noinspection PyPep8Naming
 # pylint: disable=invalid-name
 def FormatMessage(
@@ -1075,7 +1130,7 @@ def NtMapViewOfSection(
         win32_protect
     )
 
-    return nt_status == STATUS_SUCCESS
+    return _nt_ok(nt_status)
 
 # noinspection PyPep8Naming
 # pylint: disable=invalid-name
@@ -1095,7 +1150,7 @@ def NtUnmapViewOfSection(process_handle: int, base_address: int) -> bool:
         https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wdm/nf-wdm-zwunmapviewofsection
     """
     nt_status: int = _NtUnmapViewOfSection(process_handle, base_address)
-    return nt_status == STATUS_SUCCESS
+    return _nt_ok(nt_status)
 
 # noinspection PyPep8Naming
 # pylint: disable=invalid-name
@@ -1127,7 +1182,7 @@ def NtQueryInformationProcess(
         process_information_length,
         0
     )
-    return nt_status == STATUS_SUCCESS
+    return _nt_ok(nt_status)
 
 # noinspection PyPep8Naming
 # pylint: disable=invalid-name
@@ -1145,7 +1200,7 @@ def NtSuspendProcess(process_handle: int) -> bool:
         https://cyberstoph.org/posts/2021/05/fun-with-processes-suspend-and-resume/
     """
     nt_status: int = _NtSuspendProcess(process_handle)
-    return nt_status == STATUS_SUCCESS
+    return _nt_ok(nt_status)
 
 # noinspection PyPep8Naming
 # pylint: disable=invalid-name
@@ -1163,7 +1218,7 @@ def NtResumeProcess(process_handle: int) -> bool:
         https://cyberstoph.org/posts/2021/05/fun-with-processes-suspend-and-resume/
     """
     nt_status: int = _NtResumeProcess(process_handle)
-    return nt_status == STATUS_SUCCESS
+    return _nt_ok(nt_status)
 
 # noinspection PyPep8Naming
 # pylint: disable=invalid-name
@@ -1732,6 +1787,14 @@ def MessageBoxW(window_handle: int, text: str, caption: str, type_flags: int) ->
 _GetLastError.argtypes = []
 _GetLastError.restype = DWORD
 
+_SetLastError = windll.kernel32.SetLastError
+_SetLastError.argtypes = [DWORD]
+_SetLastError.restype = None
+
+_RtlNtStatusToDosError = windll.ntdll.RtlNtStatusToDosError
+_RtlNtStatusToDosError.argtypes = [ULONG]
+_RtlNtStatusToDosError.restype = ULONG
+
 _FormatMessageA = windll.kernel32.FormatMessageA
 _FormatMessageA.argtypes = [DWORD, LPVOID, DWORD, DWORD, LPSTR, DWORD, LPVOID]
 _FormatMessageA.restype = DWORD

tests/test_ntstatus_last_error.py

@@ -0,0 +1,64 @@
+"""NTSTATUS routines report failure in their return value, not via SetLastError.
+
+Without translation, a `Win32Exception()` raised after one of them reads whatever the
+previous call left behind: 0 on a fresh thread ("The operation completed successfully"),
+or a stale unrelated code otherwise. These tests pin the translated behaviour.
+"""
+
+import ctypes
+
+from MemLib.Constants import STATUS_SUCCESS
+from MemLib.Process import Process
+from MemLib.windows import (
+    GetLastError, NtUnmapViewOfSection, RtlNtStatusToDosError, SetLastError, Win32Exception, _nt_ok,
+)
+
+# What NtUnmapViewOfSection returns for an address that is not the base of a mapped
+# view. Translates to ERROR_INVALID_ADDRESS (487).
+STATUS_NOT_MAPPED_VIEW: int = 0xC0000019
+ERROR_INVALID_ADDRESS: int = 487
+UNMAPPED_ADDRESS: int = 0xDEAD0000
+
+
+def test_status_translates_to_the_matching_win32_code():
+    assert RtlNtStatusToDosError(STATUS_NOT_MAPPED_VIEW) == ERROR_INVALID_ADDRESS
+    assert RtlNtStatusToDosError(0xC0000008) == 6  # STATUS_INVALID_HANDLE -> ERROR_INVALID_HANDLE
+    assert RtlNtStatusToDosError(0xC000000D) == 87  # STATUS_INVALID_PARAMETER -> ERROR_INVALID_PARAMETER
+
+
+def test_nt_ok_leaves_last_error_alone_on_success():
+    SetLastError(1234)
+
+    assert _nt_ok(STATUS_SUCCESS) is True
+    # A successful Win32 call does not clear the last error either, so neither does this.
+    assert GetLastError() == 1234
+
+
+def test_nt_ok_publishes_the_translated_failure():
+    SetLastError(0)
+
+    assert _nt_ok(STATUS_NOT_MAPPED_VIEW) is False
+    assert GetLastError() == ERROR_INVALID_ADDRESS
+
+
+def test_failure_is_reported_even_from_a_clean_last_error():
+    """Regression: GetLastError() == 0 here made Win32Exception() say
+    'The operation completed successfully.' inside a raised exception."""
+    process = Process(ctypes.windll.kernel32.GetCurrentProcessId())
+    SetLastError(0)
+
+    assert NtUnmapViewOfSection(process.handle, UNMAPPED_ADDRESS) is False
+
+    error = Win32Exception()
+    assert error.code == ERROR_INVALID_ADDRESS
+    assert "completed successfully" not in error.message
+
+
+def test_failure_is_not_masked_by_a_stale_last_error():
+    """Regression: the stale code from an earlier call was reported instead, which is
+    wrong but plausible -- harder to spot than the 'success' message."""
+    process = Process(ctypes.windll.kernel32.GetCurrentProcessId())
+    SetLastError(1234)
+
+    assert NtUnmapViewOfSection(process.handle, UNMAPPED_ADDRESS) is False
+    assert GetLastError() == ERROR_INVALID_ADDRESS