Maarten Lankhorst : user32: Fix enumeration for EnumWindowStations and EnumDesktops.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: fd69ababec3b26e76837ed0d3cd82509c619a38e
URL:    http://source.winehq.org/git/wine.git/?a=commit;h=fd69ababec3b26e76837ed0d3cd82509c619a38e

Author: Maarten Lankhorst <m.b.lankhorst at gmail.com>
Date:   Wed Mar  5 13:41:41 2008 -0800

user32: Fix enumeration for EnumWindowStations and EnumDesktops.

---

 dlls/user32/tests/winstation.c |  102 ++++++++++++++++++++++++++++++++++++++++
 dlls/user32/winstation.c       |   36 ++++++++++++--
 2 files changed, 132 insertions(+), 6 deletions(-)

diff --git a/dlls/user32/tests/winstation.c b/dlls/user32/tests/winstation.c
index 16a1c93..d69faa4 100644
--- a/dlls/user32/tests/winstation.c
+++ b/dlls/user32/tests/winstation.c
@@ -22,6 +22,7 @@
 #include "winbase.h"
 #include "wingdi.h"
 #include "winuser.h"
+#include "winnls.h"
 
 #define DESKTOP_ALL_ACCESS 0x01ff
 
@@ -248,6 +249,105 @@ static void test_handles(void)
     CloseHandle( hthread );
 }
 
+/* Enumeration tests */
+
+static BOOL CALLBACK window_station_callbackA(LPSTR winsta, LPARAM lp)
+{
+    trace("window_station_callbackA called with argument %s\n", winsta);
+    return lp;
+}
+
+static BOOL CALLBACK open_window_station_callbackA(LPSTR winsta, LPARAM lp)
+{
+    HWINSTA hwinsta;
+
+    trace("open_window_station_callbackA called with argument %s\n", winsta);
+    hwinsta = OpenWindowStationA(winsta, FALSE, WINSTA_ENUMERATE);
+    ok(hwinsta != NULL, "Could not open desktop %s!\n", winsta);
+    if (hwinsta)
+        CloseWindowStation(hwinsta);
+    return lp;
+}
+
+static void test_enumstations(void)
+{
+    BOOL ret;
+
+    if (0) /* Crashes instead */
+    {
+        SetLastError(0xbabefeed);
+        ret = EnumWindowStationsA(NULL, 0);
+        ok(!ret, "EnumWindowStationsA returned succesfully!\n");
+        ok(GetLastError() == ERROR_INVALID_PARAMETER, "LastError is set to %08x\n", GetLastError());
+    }
+
+    SetLastError(0xdeadbeef);
+    ret = EnumWindowStationsA(open_window_station_callbackA, 0x12345);
+    ok(ret == 0x12345, "EnumWindowStationsA returned %x\n", ret);
+    ok(GetLastError() == 0xdeadbeef, "LastError is set to %08x\n", GetLastError());
+
+    SetLastError(0xdeadbeef);
+    ret = EnumWindowStationsA(window_station_callbackA, 0);
+    ok(!ret, "EnumWindowStationsA returned %x\n", ret);
+    ok(GetLastError() == 0xdeadbeef, "LastError is set to %08x\n", GetLastError());
+}
+
+static BOOL CALLBACK desktop_callbackA(LPSTR desktop, LPARAM lp)
+{
+    trace("desktop_callbackA called with argument %s\n", desktop);
+    return lp;
+}
+
+static BOOL CALLBACK open_desktop_callbackA(LPSTR desktop, LPARAM lp)
+{
+    HDESK hdesk;
+    static int once;
+
+    trace("open_desktop_callbackA called with argument %s\n", desktop);
+    /* Only try to open one desktop */
+    if (once++)
+        return lp;
+
+    hdesk = OpenDesktopA(desktop, 0, FALSE, DESKTOP_ENUMERATE);
+    ok(hdesk != NULL, "Could not open desktop %s!\n", desktop);
+    if (hdesk)
+        CloseDesktop(hdesk);
+    return lp;
+}
+
+static void test_enumdesktops(void)
+{
+    BOOL ret;
+
+    if (0)  /* Crashes instead */
+    {
+        SetLastError(0xbabefeed);
+        ret = EnumDesktopsA(GetProcessWindowStation(), NULL, 0);
+        ok(!ret, "EnumDesktopsA returned succesfully!\n");
+        ok(GetLastError() == ERROR_INVALID_PARAMETER, "LastError is set to %08x\n", GetLastError());
+    }
+
+    SetLastError(0xdeadbeef);
+    ret = EnumDesktopsA(NULL, desktop_callbackA, 0x12345);
+    ok(ret == 0x12345, "EnumDesktopsA returned %x\n", ret);
+    ok(GetLastError() == 0xdeadbeef, "LastError is set to %08x\n", GetLastError());
+
+    SetLastError(0xdeadbeef);
+    ret = EnumDesktopsA(GetProcessWindowStation(), open_desktop_callbackA, 0x12345);
+    ok(ret == 0x12345, "EnumDesktopsA returned %x\n", ret);
+    ok(GetLastError() == 0xdeadbeef, "LastError is set to %08x\n", GetLastError());
+
+    SetLastError(0xdeadbeef);
+    ret = EnumDesktopsA(INVALID_HANDLE_VALUE, desktop_callbackA, 0x12345);
+    ok(!ret, "EnumDesktopsA returned %x\n", ret);
+    ok(GetLastError() == ERROR_INVALID_HANDLE, "LastError is set to %08x\n", GetLastError());
+
+    SetLastError(0xdeadbeef);
+    ret = EnumDesktopsA(GetProcessWindowStation(), desktop_callbackA, 0);
+    ok(!ret, "EnumDesktopsA returned %x\n", ret);
+    ok(GetLastError() == 0xdeadbeef, "LastError is set to %08x\n", GetLastError());
+}
+
 START_TEST(winstation)
 {
     /* Check whether this platform supports WindowStation calls */
@@ -260,5 +360,7 @@ START_TEST(winstation)
         return;
     }
 
+    test_enumstations();
+    test_enumdesktops();
     test_handles();
 }
diff --git a/dlls/user32/winstation.c b/dlls/user32/winstation.c
index 75cb286..906f2ba 100644
--- a/dlls/user32/winstation.c
+++ b/dlls/user32/winstation.c
@@ -18,6 +18,9 @@
  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
  */
 
+#include "ntstatus.h"
+#define WIN32_NO_STATUS
+
 #include <stdarg.h>
 #include "windef.h"
 #include "winbase.h"
@@ -211,18 +214,27 @@ BOOL WINAPI EnumWindowStationsW( WINSTAENUMPROCW func, LPARAM lparam )
     unsigned int index = 0;
     WCHAR name[MAX_PATH];
     BOOL ret = TRUE;
+    NTSTATUS status;
 
     while (ret)
     {
         SERVER_START_REQ( enum_winstation )
         {
             req->index = index;
-            wine_server_set_reply( req, name, sizeof(name) );
-            ret = !wine_server_call( req );
+            wine_server_set_reply( req, name, sizeof(name) - sizeof(WCHAR) );
+            status = wine_server_call( req );
+            name[wine_server_reply_size(reply)/sizeof(WCHAR)] = 0;
             index = reply->next;
         }
         SERVER_END_REQ;
-        if (ret) ret = func( name, lparam );
+        if (status == STATUS_NO_MORE_ENTRIES)
+            break;
+        if (status)
+        {
+            SetLastError( RtlNtStatusToDosError( status ) );
+            return FALSE;
+        }
+        ret = func( name, lparam );
     }
     return ret;
 }
@@ -401,6 +413,10 @@ BOOL WINAPI EnumDesktopsW( HWINSTA winsta, DESKTOPENUMPROCW func, LPARAM lparam
     unsigned int index = 0;
     WCHAR name[MAX_PATH];
     BOOL ret = TRUE;
+    NTSTATUS status;
+
+    if (!winsta)
+        winsta = GetProcessWindowStation();
 
     while (ret)
     {
@@ -408,12 +424,20 @@ BOOL WINAPI EnumDesktopsW( HWINSTA winsta, DESKTOPENUMPROCW func, LPARAM lparam
         {
             req->winstation = winsta;
             req->index      = index;
-            wine_server_set_reply( req, name, sizeof(name) );
-            ret = !wine_server_call( req );
+            wine_server_set_reply( req, name, sizeof(name) - sizeof(WCHAR) );
+            status = wine_server_call( req );
+            name[wine_server_reply_size(reply)/sizeof(WCHAR)] = 0;
             index = reply->next;
         }
         SERVER_END_REQ;
-        if (ret) ret = func( name, lparam );
+        if (status == STATUS_NO_MORE_ENTRIES)
+            break;
+        if (status)
+        {
+            SetLastError( RtlNtStatusToDosError( status ) );
+            return FALSE;
+        }
+        ret = func(name, lparam);
     }
     return ret;
 }