From 116890da122da856c49dc66643bda5cd8360339b Mon Sep 17 00:00:00 2001
From: Gijs Vermeulen <gijsvrm@gmail.com>
Date: Thu, 28 May 2020 15:59:27 +0200
Subject: [PATCH] wtsapi32: Improve WTSQuerySessionInformationW stub.

Signed-off-by: Gijs Vermeulen <gijsvrm@gmail.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>
---
 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 f45598f5484..c9312cd97c5 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 0ba05abf2c7..ef3e0d10a0f 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;
 }