[PATCH v2 2/7] ntdll: Validate length in get_working_set_ex().

[Paul Gofman]

From: Paul Gofman

---
 dlls/ntdll/unix/virtual.c     |  4 +++-
 dlls/psapi/tests/psapi_main.c | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/dlls/ntdll/unix/virtual.c b/dlls/ntdll/unix/virtual.c
index 3981905bcd3..077b68f59f7 100644
--- a/dlls/ntdll/unix/virtual.c
+++ b/dlls/ntdll/unix/virtual.c
@@ -5145,6 +5145,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;
@@ -5241,7 +5243,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);