[PATCH v3 1/2] wtsapi32: Implement WTSQuerySessionInformationA.

[Gijs Vermeulen]

Wine-Bug: https://bugs.winehq.org/show_bug.cgi?id=49178
Signed-off-by: Gijs Vermeulen 
---

v3: Use WTSFreeMemory and return early on error, thanks Zhiyi

dlls/wtsapi32/tests/wtsapi.c | 44 +++++++++++++++++++++++++++-------
 dlls/wtsapi32/wtsapi32.c     | 46 ++++++++++++++++++++++++++++--------
 2 files changed, 72 insertions(+), 18 deletions(-)

diff --git a/dlls/wtsapi32/tests/wtsapi.c b/dlls/wtsapi32/tests/wtsapi.c
index 67f56bbd7f..f45598f548 100644
--- a/dlls/wtsapi32/tests/wtsapi.c
+++ b/dlls/wtsapi32/tests/wtsapi.c
@@ -90,19 +90,47 @@ static void test_WTSEnumerateProcessesW(void)
     WTSFreeMemory(info);
 }

-static void test_WTSQuerySessionInformationW(void)
+static void test_WTSQuerySessionInformation(void)
 {
     BOOL ret;
-    WCHAR *buf;
+    WCHAR *buf1;
+    char *buf2;
     DWORD count;

count = 0;
-    buf = NULL;
-    ret = WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, &buf, &count);
+    buf1 = NULL;
+    ret = WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, &buf1, &count);
     ok(ret, "got %u\n", GetLastError());
-    ok(buf != NULL, "buf not set\n");
-    ok(count == (lstrlenW(buf) + 1) * sizeof(WCHAR), "got %u\n", count);
-    WTSFreeMemory(buf);
+    ok(buf1 != NULL, "buf not set\n");
+    ok(count == (lstrlenW(buf1) + 1) * sizeof(WCHAR), "expected %u, got %u\n", (lstrlenW(buf1) + 1) * sizeof(WCHAR), count);
+    WTSFreeMemory(buf1);
+
+    SetLastError(0xdeadbeef);
+    count = 0;
+    ret = WTSQuerySessionInformationA(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, NULL, &count);
+    ok(!ret, "got %u\n", GetLastError());
+    ok(count == 0, "got %u\n", count);
+    ok(GetLastError() == ERROR_INVALID_USER_BUFFER, "got %u\n", GetLastError());
+
+    SetLastError(0xdeadbeef);
+    count = 1;
+    ret = WTSQuerySessionInformationA(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, NULL, &count);
+    ok(!ret, "got %u\n", GetLastError());
+    ok(count == 1, "got %u\n", count);
+    ok(GetLastError() == ERROR_INVALID_USER_BUFFER, "got %u\n", GetLastError());
+
+    SetLastError(0xdeadbeef);
+    ret = WTSQuerySessionInformationA(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, &buf2, NULL);
+    ok(!ret, "got %u\n", GetLastError());
+    ok(GetLastError() == ERROR_INVALID_USER_BUFFER, "got %u\n", GetLastError());
+
+    count = 0;
+    buf2 = NULL;
+    ret = WTSQuerySessionInformationA(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, &buf2, &count);
+    ok(ret, "got %u\n", GetLastError());
+    ok(buf2 != NULL, "buf not set\n");
+    ok(count == lstrlenA(buf2) + 1, "expected %u, got %u\n", lstrlenA(buf2) + 1, count);
+    WTSFreeMemory(buf2);
 }

static void test_WTSQueryUserToken(void)
@@ -118,6 +146,6 @@ static void test_WTSQueryUserToken(void)
 START_TEST (wtsapi)
 {
     test_WTSEnumerateProcessesW();
-    test_WTSQuerySessionInformationW();
+    test_WTSQuerySessionInformation();
     test_WTSQueryUserToken();
 }
diff --git a/dlls/wtsapi32/wtsapi32.c b/dlls/wtsapi32/wtsapi32.c
index 026e7f4369..93f8e96fd0 100644
--- a/dlls/wtsapi32/wtsapi32.c
+++ b/dlls/wtsapi32/wtsapi32.c
@@ -19,6 +19,7 @@
 #include <stdlib.h>
 #include "windef.h"
 #include "winbase.h"
+#include "winnls.h"
 #include "wtsapi32.h"
 #include "wine/debug.h"
 #include "wine/heap.h"
@@ -287,18 +288,43 @@ HANDLE WINAPI WTSOpenServerW(LPWSTR pServerName)
 /************************************************************
  *                WTSQuerySessionInformationA  (WTSAPI32.@)
  */
-BOOL WINAPI WTSQuerySessionInformationA(
-    HANDLE hServer,
-    DWORD SessionId,
-    WTS_INFO_CLASS WTSInfoClass,
-    LPSTR* Buffer,
-    DWORD* BytesReturned)
+BOOL WINAPI WTSQuerySessionInformationA(HANDLE server, DWORD session_id, WTS_INFO_CLASS class, char **buffer, DWORD *count)
 {
-    /* FIXME: Forward request to winsta.dll::WinStationQueryInformationA */
-    FIXME("Stub %p 0x%08x %d %p %p\n", hServer, SessionId, WTSInfoClass,
-        Buffer, BytesReturned);
+    WCHAR *bufferW = NULL;

-    return FALSE;
+    TRACE("%p 0x%08x %d %p %p\n", server, session_id, class, buffer, count);
+
+    if (!buffer || !count)
+    {
+        SetLastError(ERROR_INVALID_USER_BUFFER);
+        return FALSE;
+    }
+
+    if (!WTSQuerySessionInformationW(server, session_id, class, &bufferW, count))
+        return FALSE;
+    
+    *count = WideCharToMultiByte(CP_ACP, 0, bufferW, -1, NULL, 0, NULL, NULL);
+    if (!*count)
+    {
+        WTSFreeMemory(bufferW);
+        return FALSE;
+    }
+    
+    if (!(*buffer = heap_alloc(*count)))
+    {
+        WTSFreeMemory(bufferW);
+        return FALSE;
+    }
+
+    if (!(*count = WideCharToMultiByte(CP_ACP, 0, bufferW, -1, *buffer, *count, NULL, NULL)))
+    {
+        WTSFreeMemory(bufferW);
+        heap_free(*buffer);
+        return FALSE;
+    }
+
+    WTSFreeMemory(bufferW);
+    return TRUE;
 }

/************************************************************