wtsapi32: Implement WTSEnumerateProcessesW().

Based on a patch by Sebastian Lackner.

Wine-Bug: https://bugs.winehq.org/show_bug.cgi?id=29903
Signed-off-by: Zebediah Figura <z.figura12@gmail.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>
(cherry picked from commit 5f8c7c2563)
Signed-off-by: Michael Stefaniuc <mstefani@winehq.org>

dlls/wtsapi32/tests/Makefile.in

@@ -1,5 +1,5 @@
 TESTDLL   = wtsapi32.dll
-IMPORTS   = wtsapi32 advapi32
+IMPORTS   = wtsapi32 advapi32 psapi
 
 C_SRCS = \
 	wtsapi.c

dlls/wtsapi32/tests/wtsapi.c

@@ -16,6 +16,8 @@
  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
  */
 
+#include <ntstatus.h>
+#define WIN32_NO_STATUS
 #include <stdarg.h>
 #include <stdlib.h>
 #include <windef.h>
@@ -23,25 +25,93 @@
 #include <winternl.h>
 #include <lmcons.h>
 #include <wtsapi32.h>
+#define PSAPI_VERSION 1
+#include <psapi.h>
 
 #include "wine/test.h"
 
+static BOOL (WINAPI *pWTSEnumerateProcessesExW)(HANDLE server, DWORD *level, DWORD session, WCHAR **info, DWORD *count);
+static BOOL (WINAPI *pWTSFreeMemoryExW)(WTS_TYPE_CLASS class, void *memory, ULONG count);
+
+static const SYSTEM_PROCESS_INFORMATION *find_nt_process_info(const SYSTEM_PROCESS_INFORMATION *head, DWORD pid)
+{
+    for (;;)
+    {
+        if ((DWORD)(DWORD_PTR)head->UniqueProcessId == pid)
+            return head;
+        if (!head->NextEntryOffset)
+            break;
+        head = (SYSTEM_PROCESS_INFORMATION *)((char *)head + head->NextEntryOffset);
+    }
+    return NULL;
+}
+
+static void check_wts_process_info(const WTS_PROCESS_INFOW *info, DWORD count)
+{
+    ULONG nt_length = 1024;
+    SYSTEM_PROCESS_INFORMATION *nt_info = malloc(nt_length);
+    WCHAR process_name[MAX_PATH], *process_filepart;
+    BOOL ret, found = FALSE;
+    NTSTATUS status;
+    DWORD i;
+
+    GetModuleFileNameW(NULL, process_name, MAX_PATH);
+    process_filepart = wcsrchr(process_name, '\\') + 1;
+
+    while ((status = NtQuerySystemInformation(SystemProcessInformation, nt_info,
+            nt_length, NULL)) == STATUS_INFO_LENGTH_MISMATCH)
+    {
+        nt_length *= 2;
+        nt_info = realloc(nt_info, nt_length);
+    }
+    ok(!status, "got %#x\n", status);
+
+    for (i = 0; i < count; i++)
+    {
+        char sid_buffer[50];
+        SID_AND_ATTRIBUTES *sid = (SID_AND_ATTRIBUTES *)sid_buffer;
+        const SYSTEM_PROCESS_INFORMATION *nt_process;
+        HANDLE process, token;
+        DWORD size;
+
+        nt_process = find_nt_process_info(nt_info, info[i].ProcessId);
+        ok(!!nt_process, "failed to find pid %#x\n", info[i].ProcessId);
+
+        winetest_push_context("pid %#x", info[i].ProcessId);
+
+        ok(info[i].SessionId == nt_process->SessionId, "expected session id %#x, got %#x\n",
+                nt_process->SessionId, info[i].SessionId);
+
+        ok(!memcmp(info[i].pProcessName, nt_process->ProcessName.Buffer, nt_process->ProcessName.Length),
+                "expected process name %s, got %s\n",
+                debugstr_w(nt_process->ProcessName.Buffer), debugstr_w(info[i].pProcessName));
+
+        if ((process = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, info[i].ProcessId)))
+        {
+            ret = OpenProcessToken(process, TOKEN_QUERY, &token);
+            ok(ret, "failed to open token, error %u\n", GetLastError());
+            ret = GetTokenInformation(token, TokenUser, sid_buffer, sizeof(sid_buffer), &size);
+            ok(ret, "failed to get token user, error %u\n", GetLastError());
+            ok(EqualSid(info[i].pUserSid, sid->Sid), "SID did not match\n");
+            CloseHandle(token);
+            CloseHandle(process);
+        }
+
+        winetest_pop_context();
+
+        found = found || !wcscmp(info[i].pProcessName, process_filepart);
+    }
+
+    ok(found, "did not find current process\n");
+
+    free(nt_info);
+}
+
 static void test_WTSEnumerateProcessesW(void)
 {
-    BOOL found = FALSE, ret;
-    DWORD count, i;
     PWTS_PROCESS_INFOW info;
-    WCHAR *pname, nameW[MAX_PATH];
-
-    GetModuleFileNameW(NULL, nameW, MAX_PATH);
-    for (pname = nameW + lstrlenW(nameW); pname > nameW; pname--)
-    {
-        if(*pname == '/' || *pname == '\\')
-        {
-            pname++;
-            break;
-        }
-    }
+    DWORD count, level;
+    BOOL ret;
 
     info = NULL;
     SetLastError(0xdeadbeef);
@@ -80,15 +150,41 @@ static void test_WTSEnumerateProcessesW(void)
     info = NULL;
     SetLastError(0xdeadbeef);
     ret = WTSEnumerateProcessesW(WTS_CURRENT_SERVER_HANDLE, 0, 1, &info, &count);
-    ok(ret || broken(!ret), /* fails on Win2K with error ERROR_APP_WRONG_OS */
-        "expected WTSEnumerateProcessesW to succeed; failed with %d\n", GetLastError());
-    for(i = 0; ret && i < count; i++)
-    {
-        found = found || !lstrcmpW(pname, info[i].pProcessName);
-    }
-    todo_wine
-    ok(found || broken(!ret), "process name %s not found\n", wine_dbgstr_w(pname));
+    ok(ret, "expected success\n");
+    ok(!GetLastError(), "got error %u\n", GetLastError());
+    check_wts_process_info(info, count);
     WTSFreeMemory(info);
+
+    if (!pWTSEnumerateProcessesExW)
+    {
+        skip("WTSEnumerateProcessesEx is not available\n");
+        return;
+    }
+
+    level = 0;
+
+    SetLastError(0xdeadbeef);
+    count = 0xdeadbeef;
+    ret = pWTSEnumerateProcessesExW(WTS_CURRENT_SERVER_HANDLE, &level, WTS_ANY_SESSION, NULL, &count);
+    ok(!ret, "expected failure\n");
+    ok(GetLastError() == ERROR_INVALID_PARAMETER, "got error %u\n", GetLastError());
+    ok(count == 0xdeadbeef, "got count %u\n", count);
+
+    info = (void *)0xdeadbeef;
+    SetLastError(0xdeadbeef);
+    ret = pWTSEnumerateProcessesExW(WTS_CURRENT_SERVER_HANDLE, &level, WTS_ANY_SESSION, (WCHAR **)&info, NULL);
+    ok(!ret, "expected failure\n");
+    ok(GetLastError() == ERROR_INVALID_PARAMETER, "got error %u\n", GetLastError());
+    ok(info == (void *)0xdeadbeef, "got info %p\n", info);
+
+    info = NULL;
+    count = 0;
+    SetLastError(0xdeadbeef);
+    ret = pWTSEnumerateProcessesExW(WTS_CURRENT_SERVER_HANDLE, &level, WTS_ANY_SESSION, (WCHAR **)&info, &count);
+    ok(ret, "expected success\n");
+    ok(!GetLastError(), "got error %u\n", GetLastError());
+    check_wts_process_info(info, count);
+    pWTSFreeMemoryExW(WTSTypeProcessInfoLevel0, info, count);
 }
 
 static void test_WTSQuerySessionInformation(void)
@@ -201,6 +297,9 @@ static void test_WTSQueryUserToken(void)
 
 START_TEST (wtsapi)
 {
+    pWTSEnumerateProcessesExW = (void *)GetProcAddress(GetModuleHandleA("wtsapi32"), "WTSEnumerateProcessesExW");
+    pWTSFreeMemoryExW = (void *)GetProcAddress(GetModuleHandleA("wtsapi32"), "WTSFreeMemoryExW");
+
     test_WTSEnumerateProcessesW();
     test_WTSQuerySessionInformation();
     test_WTSQueryUserToken();

dlls/wtsapi32/wtsapi32.c

@@ -15,10 +15,13 @@
  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
  */
 
+#include "ntstatus.h"
+#define WIN32_NO_STATUS
 #include <stdarg.h>
 #include <stdlib.h>
 #include "windef.h"
 #include "winbase.h"
+#include "winternl.h"
 #include "winnls.h"
 #include "lmcons.h"
 #include "wtsapi32.h"
@@ -76,11 +79,127 @@ BOOL WINAPI WTSEnableChildSessions(BOOL enable)
 /************************************************************
  *                WTSEnumerateProcessesExW  (WTSAPI32.@)
  */
-BOOL WINAPI WTSEnumerateProcessesExW(HANDLE server, DWORD *level, DWORD session_id, WCHAR **info, DWORD *count)
+BOOL WINAPI WTSEnumerateProcessesExW(HANDLE server, DWORD *level, DWORD session_id,
+        WCHAR **ret_info, DWORD *ret_count)
 {
-    FIXME("Stub %p %p %d %p %p\n", server, level, session_id, info, count);
-    if (count) *count = 0;
+    SYSTEM_PROCESS_INFORMATION *nt_info, *nt_process;
+    WTS_PROCESS_INFOW *info;
+    ULONG nt_size = 4096;
+    DWORD count, size;
+    NTSTATUS status;
+    char *p;
+
+    TRACE("server %p, level %u, session_id %#x, ret_info %p, ret_count %p\n",
+            server, *level, session_id, ret_info, ret_count);
+
+    if (!ret_info || !ret_count)
+    {
+        SetLastError(ERROR_INVALID_PARAMETER);
         return FALSE;
+    }
+
+    if (session_id != WTS_ANY_SESSION)
+        FIXME("ignoring session id %#x\n", session_id);
+
+    if (*level)
+    {
+        FIXME("unhandled level %u\n", *level);
+        SetLastError(ERROR_CALL_NOT_IMPLEMENTED);
+        return FALSE;
+    }
+
+    if (!(nt_info = malloc(nt_size)))
+    {
+        SetLastError(ERROR_OUTOFMEMORY);
+        return FALSE;
+    }
+
+    while ((status = NtQuerySystemInformation(SystemProcessInformation, nt_info,
+            nt_size, NULL)) == STATUS_INFO_LENGTH_MISMATCH)
+    {
+        SYSTEM_PROCESS_INFORMATION *new_info;
+
+        nt_size *= 2;
+        if (!(new_info = realloc(nt_info, nt_size)))
+        {
+            free(nt_info);
+            SetLastError(ERROR_OUTOFMEMORY);
+            return FALSE;
+        }
+        nt_info = new_info;
+    }
+    if (status)
+    {
+        free(nt_info);
+        SetLastError(RtlNtStatusToDosError(status));
+        return FALSE;
+    }
+
+    size = 0;
+    count = 0;
+    nt_process = nt_info;
+    for (;;)
+    {
+        size += sizeof(WTS_PROCESS_INFOW) + nt_process->ProcessName.Length + sizeof(WCHAR);
+        size += offsetof(SID, SubAuthority[SID_MAX_SUB_AUTHORITIES]);
+        ++count;
+
+        if (!nt_process->NextEntryOffset)
+            break;
+        nt_process = (SYSTEM_PROCESS_INFORMATION *)((char *)nt_process + nt_process->NextEntryOffset);
+    }
+
+    if (!(info = heap_alloc(size)))
+    {
+        free(nt_info);
+        SetLastError(ERROR_OUTOFMEMORY);
+        return FALSE;
+    }
+    p = (char *)(info + count);
+
+    count = 0;
+    nt_process = nt_info;
+    for (;;)
+    {
+        HANDLE process, token;
+
+        info[count].SessionId = nt_process->SessionId;
+        info[count].ProcessId = (DWORD_PTR)nt_process->UniqueProcessId;
+
+        info[count].pProcessName = (WCHAR *)p;
+        memcpy(p, nt_process->ProcessName.Buffer, nt_process->ProcessName.Length);
+        info[count].pProcessName[nt_process->ProcessName.Length / sizeof(WCHAR)] = 0;
+        p += nt_process->ProcessName.Length + sizeof(WCHAR);
+
+        info[count].pUserSid = NULL;
+        if ((process = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, info[count].ProcessId)))
+        {
+            if (OpenProcessToken(process, TOKEN_QUERY, &token))
+            {
+                char buffer[sizeof(TOKEN_USER) + offsetof(SID, SubAuthority[SID_MAX_SUB_AUTHORITIES])];
+                TOKEN_USER *user = (TOKEN_USER *)buffer;
+                DWORD size;
+
+                GetTokenInformation(token, TokenUser, buffer, sizeof(buffer), &size);
+                info[count].pUserSid = p;
+                size = GetLengthSid(user->User.Sid);
+                memcpy(p, user->User.Sid, size);
+                p += size;
+                CloseHandle(token);
+            }
+            CloseHandle(process);
+        }
+
+        ++count;
+        if (!nt_process->NextEntryOffset)
+            break;
+        nt_process = (SYSTEM_PROCESS_INFORMATION *)((char *)nt_process + nt_process->NextEntryOffset);
+    }
+
+    *ret_info = (WCHAR *)info;
+    *ret_count = count;
+    SetLastError(0);
+    return TRUE;
 }
 
 /************************************************************
@@ -113,22 +232,20 @@ BOOL WINAPI WTSEnumerateProcessesA(HANDLE hServer, DWORD Reserved, DWORD Version
 /************************************************************
  *                WTSEnumerateProcessesW  (WTSAPI32.@)
  */
-BOOL WINAPI WTSEnumerateProcessesW(HANDLE hServer, DWORD Reserved, DWORD Version,
-    PWTS_PROCESS_INFOW* ppProcessInfo, DWORD* pCount)
+BOOL WINAPI WTSEnumerateProcessesW(HANDLE server, DWORD reserved, DWORD version,
+        WTS_PROCESS_INFOW **info, DWORD *count)
 {
-    FIXME("Stub %p 0x%08x 0x%08x %p %p\n", hServer, Reserved, Version,
-          ppProcessInfo, pCount);
+    DWORD level = 0;
 
-    if (!ppProcessInfo || !pCount || Reserved != 0 || Version != 1)
+    TRACE("server %p, reserved %#x, version %u, info %p, count %p\n", server, reserved, version, info, count);
+
+    if (reserved || version != 1)
     {
         SetLastError(ERROR_INVALID_PARAMETER);
         return FALSE;
     }
 
-    *pCount = 0;
-    *ppProcessInfo = NULL;
-
-    return TRUE;
+    return WTSEnumerateProcessesExW(server, &level, WTS_ANY_SESSION, (WCHAR **)info, count);
 }
 
 /************************************************************

include/wtsapi32.h

@@ -171,6 +171,7 @@ DECL_WINELIB_TYPE_AW(PWTS_SERVER_INFO)
 
 #define WTS_CURRENT_SERVER_HANDLE ((HANDLE)NULL)
 #define WTS_CURRENT_SESSION (~0u)
+#define WTS_ANY_SESSION ((DWORD)-2)
 
 void WINAPI WTSCloseServer(HANDLE);
 BOOL WINAPI WTSConnectSessionA(ULONG, ULONG, PSTR, BOOL);