Paul Gofman : ntdll: Return sufficient info size at once from NtQuerySystemInformation(SystemProcessInformation).

[Alexandre Julliard]

Module: wine
Branch: master
Commit: b5f3ddd185808e45e702490348493919dac95ad2
URL:    https://source.winehq.org/git/wine.git/?a=commit;h=b5f3ddd185808e45e70249034...

Author: Paul Gofman <pgofman(a)codeweavers.com>
Date:   Mon Oct 25 21:46:19 2021 +0300

ntdll: Return sufficient info size at once from NtQuerySystemInformation(SystemProcessInformation).

Signed-off-by: Paul Gofman <pgofman(a)codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard(a)winehq.org>

---

 dlls/ntdll/tests/info.c        | 43 +++++++++++++++++++++++++++++-------------
 dlls/ntdll/unix/system.c       | 17 +++++++++++++----
 include/wine/server_protocol.h |  4 +++-
 server/process.c               |  4 ++++
 server/protocol.def            |  2 ++
 server/request.h               |  4 +++-
 server/trace.c                 |  2 ++
 7 files changed, 57 insertions(+), 19 deletions(-)

diff --git a/dlls/ntdll/tests/info.c b/dlls/ntdll/tests/info.c
index 6e71f0848ec..635c551383c 100644
--- a/dlls/ntdll/tests/info.c
+++ b/dlls/ntdll/tests/info.c
@@ -470,25 +470,30 @@ static void test_query_process(void)
         SYSTEM_THREAD_INFORMATION ti[1];
     } SYSTEM_PROCESS_INFORMATION_PRIVATE;
 
-    ULONG SystemInformationLength = sizeof(SYSTEM_PROCESS_INFORMATION_PRIVATE);
-    SYSTEM_PROCESS_INFORMATION_PRIVATE *spi, *spi_buf = HeapAlloc(GetProcessHeap(), 0, SystemInformationLength);
+    SYSTEM_PROCESS_INFORMATION_PRIVATE *spi, *spi_buf;
+    BOOL current_process_found = FALSE;
 
     /* test ReturnLength */
     ReturnLength = 0;
     status = pNtQuerySystemInformation(SystemProcessInformation, NULL, 0, &ReturnLength);
     ok( status == STATUS_INFO_LENGTH_MISMATCH, "Expected STATUS_INFO_LENGTH_MISMATCH got %08x\n", status);
-    ok( ReturnLength > 0, "got 0 length\n");
+    ok( ReturnLength > 0, "got 0 length\n" );
 
-    /* W2K3 and later returns the needed length, the rest returns 0, so we have to loop */
-    for (;;)
+    /* W2K3 and later returns the needed length, the rest returns 0. */
+    if (!ReturnLength)
     {
-        status = pNtQuerySystemInformation(SystemProcessInformation, spi_buf, SystemInformationLength, &ReturnLength);
-
-        if (status != STATUS_INFO_LENGTH_MISMATCH) break;
-        
-        spi_buf = HeapReAlloc(GetProcessHeap(), 0, spi_buf , SystemInformationLength *= 2);
+        win_skip( "Zero return length, skipping tests." );
+        return;
     }
-    ok( status == STATUS_SUCCESS, "Expected STATUS_SUCCESS, got %08x\n", status);
+
+    spi_buf = HeapAlloc(GetProcessHeap(), 0, ReturnLength);
+    status = pNtQuerySystemInformation(SystemProcessInformation, spi_buf, ReturnLength, &ReturnLength);
+
+    /* Sometimes new process or threads appear between the call and increase the size,
+     * otherwise the previously returned buffer size should be sufficient. */
+    ok( status == STATUS_SUCCESS || status == STATUS_INFO_LENGTH_MISMATCH,
+        "Expected STATUS_SUCCESS, got %08x\n", status );
+
     spi = spi_buf;
 
     for (;;)
@@ -496,9 +501,16 @@ static void test_query_process(void)
         DWORD_PTR tid;
         DWORD j;
 
+        winetest_push_context( "i %u (%s)", i, debugstr_w(spi->ProcessName.Buffer) );
+
         i++;
 
         last_pid = (DWORD_PTR)spi->UniqueProcessId;
+        ok( !(last_pid & 3), "Unexpected PID low bits: %p\n", spi->UniqueProcessId );
+
+        if (last_pid == GetCurrentProcessId())
+            current_process_found = TRUE;
+
         ok(!(last_pid & 3), "Unexpected PID low bits: %p\n", spi->UniqueProcessId);
         for (j = 0; j < spi->dwThreadCount; j++)
         {
@@ -511,12 +523,17 @@ static void test_query_process(void)
             ok(!(tid & 3), "Unexpected TID low bits: %p\n", spi->ti[j].ClientId.UniqueThread);
         }
 
-        if (!spi->NextEntryOffset) break;
-
+        if (!spi->NextEntryOffset)
+        {
+            winetest_pop_context();
+            break;
+        }
         one_before_last_pid = last_pid;
 
         spi = (SYSTEM_PROCESS_INFORMATION_PRIVATE*)((char*)spi + spi->NextEntryOffset);
+        winetest_pop_context();
     }
+    ok( current_process_found, "Test process not found.\n" );
     if (winetest_debug > 1) trace("%u processes, %u threads\n", i, k);
 
     if (one_before_last_pid == 0) one_before_last_pid = last_pid;
diff --git a/dlls/ntdll/unix/system.c b/dlls/ntdll/unix/system.c
index c6bbabd85c5..33e7083d18b 100644
--- a/dlls/ntdll/unix/system.c
+++ b/dlls/ntdll/unix/system.c
@@ -2446,10 +2446,13 @@ NTSTATUS WINAPI NtQuerySystemInformation( SYSTEM_INFORMATION_CLASS class,
 
     case SystemProcessInformation:  /* 5 */
     {
-        unsigned int process_count, i, j;
+        unsigned int process_count, total_thread_count, total_name_len, i, j;
         char *buffer = NULL;
         unsigned int pos = 0;
 
+C_ASSERT( sizeof(struct thread_info) <= sizeof(SYSTEM_THREAD_INFORMATION) );
+C_ASSERT( sizeof(struct process_info) <= sizeof(SYSTEM_PROCESS_INFORMATION) );
+
         if (size && !(buffer = malloc( size )))
         {
             ret = STATUS_NO_MEMORY;
@@ -2460,19 +2463,25 @@ NTSTATUS WINAPI NtQuerySystemInformation( SYSTEM_INFORMATION_CLASS class,
         {
             wine_server_set_reply( req, buffer, size );
             ret = wine_server_call( req );
-            len = reply->info_size;
+            total_thread_count = reply->total_thread_count;
+            total_name_len = reply->total_name_len;
             process_count = reply->process_count;
         }
         SERVER_END_REQ;
 
+        len = 0;
+
         if (ret)
         {
+            if (ret == STATUS_INFO_LENGTH_MISMATCH)
+                len = sizeof(SYSTEM_PROCESS_INFORMATION) * process_count
+                      + (total_name_len + process_count) * sizeof(WCHAR)
+                      + total_thread_count * sizeof(SYSTEM_THREAD_INFORMATION);
+
             free( buffer );
             break;
         }
 
-        len = 0;
-
         for (i = 0; i < process_count; i++)
         {
             SYSTEM_PROCESS_INFORMATION *nt_process = (SYSTEM_PROCESS_INFORMATION *)((char *)info + len);
diff --git a/include/wine/server_protocol.h b/include/wine/server_protocol.h
index 8a85e0bc4ff..6e38d0b075d 100644
--- a/include/wine/server_protocol.h
+++ b/include/wine/server_protocol.h
@@ -2029,6 +2029,8 @@ struct list_processes_reply
     struct reply_header __header;
     data_size_t     info_size;
     int             process_count;
+    int             total_thread_count;
+    data_size_t     total_name_len;
     /* VARARG(data,process_info,info_size); */
 };
 
@@ -6259,7 +6261,7 @@ union generic_reply
 
 /* ### protocol_version begin ### */
 
-#define SERVER_PROTOCOL_VERSION 733
+#define SERVER_PROTOCOL_VERSION 734
 
 /* ### protocol_version end ### */
 
diff --git a/server/process.c b/server/process.c
index d8d09698558..b6ad4e47bd4 100644
--- a/server/process.c
+++ b/server/process.c
@@ -1868,6 +1868,8 @@ DECL_HANDLER(list_processes)
     char *buffer;
 
     reply->process_count = 0;
+    reply->total_thread_count = 0;
+    reply->total_name_len = 0;
     reply->info_size = 0;
 
     LIST_FOR_EACH_ENTRY( process, &process_list, struct process, entry )
@@ -1877,6 +1879,8 @@ DECL_HANDLER(list_processes)
         reply->info_size = (reply->info_size + 7) & ~7;
         reply->info_size += process->running_threads * sizeof(struct thread_info);
         reply->process_count++;
+        reply->total_thread_count += process->running_threads;
+        reply->total_name_len += process->imagelen;
     }
 
     if (reply->info_size > get_reply_max_size())
diff --git a/server/protocol.def b/server/protocol.def
index 0f58dc7c85a..7d27fa382d2 100644
--- a/server/protocol.def
+++ b/server/protocol.def
@@ -1624,6 +1624,8 @@ struct process_info
 @REPLY
     data_size_t     info_size;
     int             process_count;
+    int             total_thread_count;
+    data_size_t     total_name_len;
     VARARG(data,process_info,info_size);
 @END
 
diff --git a/server/request.h b/server/request.h
index 778e7272bf5..33d4237f49f 100644
--- a/server/request.h
+++ b/server/request.h
@@ -1127,7 +1127,9 @@ C_ASSERT( sizeof(struct get_mapping_filename_reply) == 16 );
 C_ASSERT( sizeof(struct list_processes_request) == 16 );
 C_ASSERT( FIELD_OFFSET(struct list_processes_reply, info_size) == 8 );
 C_ASSERT( FIELD_OFFSET(struct list_processes_reply, process_count) == 12 );
-C_ASSERT( sizeof(struct list_processes_reply) == 16 );
+C_ASSERT( FIELD_OFFSET(struct list_processes_reply, total_thread_count) == 16 );
+C_ASSERT( FIELD_OFFSET(struct list_processes_reply, total_name_len) == 20 );
+C_ASSERT( sizeof(struct list_processes_reply) == 24 );
 C_ASSERT( FIELD_OFFSET(struct create_debug_obj_request, access) == 12 );
 C_ASSERT( FIELD_OFFSET(struct create_debug_obj_request, flags) == 16 );
 C_ASSERT( sizeof(struct create_debug_obj_request) == 24 );
diff --git a/server/trace.c b/server/trace.c
index 5a2afac504a..6e78f8281f0 100644
--- a/server/trace.c
+++ b/server/trace.c
@@ -2287,6 +2287,8 @@ static void dump_list_processes_reply( const struct list_processes_reply *req )
 {
     fprintf( stderr, " info_size=%u", req->info_size );
     fprintf( stderr, ", process_count=%d", req->process_count );
+    fprintf( stderr, ", total_thread_count=%d", req->total_thread_count );
+    fprintf( stderr, ", total_name_len=%u", req->total_name_len );
     dump_varargs_process_info( ", data=", min(cur_size,req->info_size) );
 }