Paul Gofman : ntdll: Validate length in get_working_set_ex().

[Alexandre Julliard]

Module: wine
Branch: master
Commit: 81caa1d72b8cd93219a9303ad04237672f57d988
URL:    https://gitlab.winehq.org/wine/wine/-/commit/81caa1d72b8cd93219a9303ad042376...

Author: Paul Gofman 
Date:   Fri Jun 21 16:42:44 2024 -0600

ntdll: Validate length in get_working_set_ex().

---

dlls/ntdll/unix/virtual.c     |  4 +++-
 dlls/psapi/tests/psapi_main.c | 21 +++++++++++++++++++++
 dlls/wow64/virtual.c          |  4 +++-
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/dlls/ntdll/unix/virtual.c b/dlls/ntdll/unix/virtual.c
index 0c38c57f653..254d4b4a351 100644
--- a/dlls/ntdll/unix/virtual.c
+++ b/dlls/ntdll/unix/virtual.c
@@ -5151,6 +5151,8 @@ static NTSTATUS get_working_set_ex( HANDLE process, LPCVOID addr,
         return STATUS_INVALID_INFO_CLASS;
     }

+    if (len < sizeof(*info)) return STATUS_INFO_LENGTH_MISMATCH;
+
 #if defined(HAVE_LIBPROCSTAT)
     {
         struct procstat *pstat;
@@ -5247,7 +5249,7 @@ static NTSTATUS get_working_set_ex( HANDLE process, LPCVOID addr,
 #endif

if (res_len)
-        *res_len = (UINT_PTR)p - (UINT_PTR)info;
+        *res_len = len;
     return STATUS_SUCCESS;
 }

diff --git a/dlls/psapi/tests/psapi_main.c b/dlls/psapi/tests/psapi_main.c
index 515364f156d..92529447afa 100644
--- a/dlls/psapi/tests/psapi_main.c
+++ b/dlls/psapi/tests/psapi_main.c
@@ -1194,6 +1194,8 @@ static void test_QueryWorkingSetEx(void)
 {
     PSAPI_WORKING_SET_EX_INFORMATION info[4];
     char *addr, *addr2;
+    NTSTATUS status;
+    SIZE_T size;
     DWORD prot;
     BOOL ret;

@@ -1203,6 +1205,25 @@ static void test_QueryWorkingSetEx(void)
         return;
     }

+    size = 0xdeadbeef;
+    memset(info, 0, sizeof(info));
+    status = pNtQueryVirtualMemory(GetCurrentProcess(), NULL, MemoryWorkingSetExInformation, info, 0, &size);
+    ok(status == STATUS_INFO_LENGTH_MISMATCH, "got %#lx.\n", status);
+    ok(size == 0xdeadbeef, "got %Iu.\n", size);
+
+    memset(&info, 0, sizeof(info));
+    ret = pQueryWorkingSetEx(GetCurrentProcess(), info, 0);
+    ok(!ret && GetLastError() == ERROR_BAD_LENGTH, "got ret %d, err %lu.\n", ret, GetLastError());
+
+    size = 0xdeadbeef;
+    memset(info, 0, sizeof(info));
+    status = pNtQueryVirtualMemory(GetCurrentProcess(), NULL, MemoryWorkingSetExInformation, info,
+            sizeof(*info) + sizeof(*info) / 2, &size);
+    ok(!status, "got %#lx.\n", status);
+    ok(!info->VirtualAttributes.Valid, "got %d.\n", info->VirtualAttributes.Valid);
+    ok(size == sizeof(*info) /* wow64 */ || size == sizeof(*info) + sizeof(*info) / 2 /* win64 */,
+            "got %Iu, sizeof(info) %Iu.\n", size, sizeof(info));
+
     addr = (void *)GetModuleHandleA(NULL);
     check_QueryWorkingSetEx(addr, "exe", 1, PAGE_READONLY, 1, FALSE);

diff --git a/dlls/wow64/virtual.c b/dlls/wow64/virtual.c
index 412c79abb14..7ccdd0dc65c 100644
--- a/dlls/wow64/virtual.c
+++ b/dlls/wow64/virtual.c
@@ -622,6 +622,8 @@ NTSTATUS WINAPI wow64_NtQueryVirtualMemory( UINT *args )
         MEMORY_WORKING_SET_EX_INFORMATION *info;
         ULONG i, count = len / sizeof(*info32);

+        if (len < sizeof(*info32)) return STATUS_INFO_LENGTH_MISMATCH;
+
         info = Wow64AllocateTemp( count * sizeof(*info) );
         for (i = 0; i < count; i++) info[i].VirtualAddress = ULongToPtr( info32[i].VirtualAddress );
         if (!(status = NtQueryVirtualMemory( handle, addr, class, info, count * sizeof(*info), &res_len )))
@@ -633,7 +635,7 @@ NTSTATUS WINAPI wow64_NtQueryVirtualMemory( UINT *args )
         break;
     }

-    case MemoryImageInformation: /* MEMORY_IMAEG_INFORMATION */
+    case MemoryImageInformation: /* MEMORY_IMAGE_INFORMATION */
     {
         if (len < sizeof(MEMORY_IMAGE_INFORMATION32)) return STATUS_INFO_LENGTH_MISMATCH;