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

[Zhiyi Zhang]

On 5/28/20 6:03 PM, Gijs Vermeulen wrote:
> Wine-Bug: https://bugs.winehq.org/show_bug.cgi?id=49178
> Signed-off-by: Gijs Vermeulen <gijsvrm at gmail.com>
> ---
>  dlls/wtsapi32/tests/wtsapi.c | 44 +++++++++++++++++++++++++++++-------
>  dlls/wtsapi32/wtsapi32.c     | 40 ++++++++++++++++++++++++--------
>  2 files changed, 67 insertions(+), 17 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..e7629dff5a 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,17 +288,38 @@ 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;
> +
> +    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))
> +    {
bufferW should be freed using WTSFreeMemory() instead of calling heap_free() directly.

I suggest return early on failures to avoid nested indentations.
> +        *count = WideCharToMultiByte(CP_ACP, 0, bufferW, -1, NULL, 0, NULL, NULL);
> +        if (*count)
> +        {
> +            if (!(*buffer = heap_alloc(*count)))
> +            {
> +                heap_free(bufferW);
> +                return FALSE;
> +            }
> +
> +            if ((*count = WideCharToMultiByte(CP_ACP, 0, bufferW, -1, *buffer, *count, NULL, NULL)))
> +            {
> +                heap_free(bufferW);
> +                return TRUE;
> +            }
> +            heap_free(*buffer);
> +        }
> +        heap_free(bufferW);
> +    }
>      return FALSE;
>  }
>