[PATCH 3/3] wtsapi32: Handle WTSConnectState class in WTSQuerySessionInformationW().

[Paul Gofman]

From: Paul Gofman

---
 dlls/wtsapi32/tests/wtsapi.c | 15 +++++++++++++++
 dlls/wtsapi32/wtsapi32.c     | 24 +++++++++++++-----------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/dlls/wtsapi32/tests/wtsapi.c b/dlls/wtsapi32/tests/wtsapi.c
index cac5da94e77..d45d5776f39 100644
--- a/dlls/wtsapi32/tests/wtsapi.c
+++ b/dlls/wtsapi32/tests/wtsapi.c
@@ -192,10 +192,25 @@ static void test_WTSQuerySessionInformation(void)
 {
     WCHAR *buf1, usernameW[UNLEN + 1], computernameW[MAX_COMPUTERNAME_LENGTH + 1];
     char *buf2, username[UNLEN + 1], computername[MAX_COMPUTERNAME_LENGTH + 1];
+    WTS_CONNECTSTATE_CLASS *state;
     DWORD count, tempsize;
     USHORT *protocol;
     BOOL ret;

+    count = 0;
+    ret = WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSConnectState, (WCHAR **)&state, &count);
+    ok(ret, "got error %lu\n", GetLastError());
+    ok(count == sizeof(*state), "got %lu\n", count);
+    ok(*state == WTSActive, "got %d.\n", *state);
+    WTSFreeMemory(state);
+
+    count = 0;
+    ret = WTSQuerySessionInformationA(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSConnectState, (char **)&state, &count);
+    ok(ret, "got error %lu\n", GetLastError());
+    ok(count == sizeof(*state), "got %lu\n", count);
+    ok(*state == WTSActive, "got %d.\n", *state);
+    WTSFreeMemory(state);
+
     SetLastError(0xdeadbeef);
     count = 0;
     ret = WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSUserName, NULL, &count);
diff --git a/dlls/wtsapi32/wtsapi32.c b/dlls/wtsapi32/wtsapi32.c
index 6eaea8d3c6f..f4072e7090b 100644
--- a/dlls/wtsapi32/wtsapi32.c
+++ b/dlls/wtsapi32/wtsapi32.c
@@ -468,17 +468,8 @@ BOOL WINAPI WTSQuerySessionInformationA(HANDLE server, DWORD session_id, WTS_INF
         return FALSE;
     }

-    if (class == WTSClientProtocolType)
-    {
-        USHORT *protocol;
-
-        if (!(protocol = heap_alloc(sizeof(*protocol)))) return FALSE;
-        FIXME("returning 0 protocol type\n");
-        *protocol = 0;
-        *buffer = (char *)protocol;
-        *count = sizeof(*protocol);
-        return TRUE;
-    }
+    if (class == WTSClientProtocolType || class == WTSConnectState)
+        return WTSQuerySessionInformationW(server, session_id, class, (WCHAR **)buffer, count);

if (!WTSQuerySessionInformationW(server, session_id, class, &bufferW, count))
         return FALSE;
@@ -520,6 +511,17 @@ BOOL WINAPI WTSQuerySessionInformationW(HANDLE server, DWORD session_id, WTS_INF
         return FALSE;
     }

+    if (class == WTSConnectState)
+    {
+        WTS_CONNECTSTATE_CLASS *state;
+
+        if (!(state = heap_alloc(sizeof(*state)))) return FALSE;
+        *state = WTSActive;
+        *buffer = (WCHAR *)state;
+        *count = sizeof(*state);
+        return TRUE;
+    }
+
     if (class == WTSClientProtocolType)
     {
         USHORT *protocol;