[PATCH v3 2/2] wtsapi32: Improve WTSQuerySessionInformationW stub.

[Gijs Vermeulen]

Signed-off-by: Gijs Vermeulen 
---
 dlls/wtsapi32/tests/wtsapi.c | 19 +++++++++++++++++++
 dlls/wtsapi32/wtsapi32.c     | 11 +++++++++++
 2 files changed, 30 insertions(+)

diff --git a/dlls/wtsapi32/tests/wtsapi.c b/dlls/wtsapi32/tests/wtsapi.c
index f45598f548..c9312cd97c 100644
--- a/dlls/wtsapi32/tests/wtsapi.c
+++ b/dlls/wtsapi32/tests/wtsapi.c
@@ -97,6 +97,25 @@ static void test_WTSQuerySessionInformation(void)
     char *buf2;
     DWORD count;

+    SetLastError(0xdeadbeef);
+    count = 0;
+    ret = WTSQuerySessionInformationW(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 = WTSQuerySessionInformationW(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 = WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, &buf1, NULL);
+    ok(!ret, "got %u\n", GetLastError());
+    ok(GetLastError() == ERROR_INVALID_USER_BUFFER, "got %u\n", GetLastError());
+
     count = 0;
     buf1 = NULL;
     ret = WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, &buf1, &count);
diff --git a/dlls/wtsapi32/wtsapi32.c b/dlls/wtsapi32/wtsapi32.c
index 93f8e96fd0..6e20b75128 100644
--- a/dlls/wtsapi32/wtsapi32.c
+++ b/dlls/wtsapi32/wtsapi32.c
@@ -341,6 +341,12 @@ BOOL WINAPI WTSQuerySessionInformationW(
     FIXME("Stub %p 0x%08x %d %p %p\n", hServer, SessionId, WTSInfoClass,
         Buffer, BytesReturned);

+    if (!Buffer || !BytesReturned)
+    {
+        SetLastError(ERROR_INVALID_USER_BUFFER);
+        return FALSE;
+    }
+
     if (WTSInfoClass == WTSUserName)
     {
         WCHAR *username;
@@ -354,6 +360,11 @@ BOOL WINAPI WTSQuerySessionInformationW(
         *BytesReturned = count * sizeof(WCHAR);
         return TRUE;
     }
+    else
+    {
+        *Buffer = NULL;
+        *BytesReturned = 0;
+    }
     return FALSE;
 }