kernel32: Fix RegisterWaitForSingleObject for console handles.

Signed-off-by: Keno Fischer <keno@juliacomputing.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>

dlls/kernel32/sync.c

@@ -148,6 +148,23 @@ DWORD WINAPI WaitForMultipleObjects( DWORD count, const HANDLE *handles,
     return WaitForMultipleObjectsEx( count, handles, wait_all, timeout, FALSE );
 }
 
+static HANDLE normalize_handle_if_console(HANDLE handle)
+{
+    if ((handle == (HANDLE)STD_INPUT_HANDLE) ||
+        (handle == (HANDLE)STD_OUTPUT_HANDLE) ||
+        (handle == (HANDLE)STD_ERROR_HANDLE))
+        handle = GetStdHandle( HandleToULong(handle) );
+
+    /* yes, even screen buffer console handles are waitable, and are
+     * handled as a handle to the console itself !!
+     */
+    if (is_console_handle(handle))
+    {
+        if (VerifyConsoleIoHandle(handle))
+            handle = GetConsoleInputWaitHandle();
+    }
+    return handle;
+}
 
 /***********************************************************************
  *           WaitForMultipleObjectsEx   (KERNEL32.@)
@@ -167,23 +184,7 @@ DWORD WINAPI WaitForMultipleObjectsEx( DWORD count, const HANDLE *handles,
         return WAIT_FAILED;
     }
     for (i = 0; i < count; i++)
-    {
+        hloc[i] = normalize_handle_if_console(handles[i]);
-        if ((handles[i] == (HANDLE)STD_INPUT_HANDLE) ||
-            (handles[i] == (HANDLE)STD_OUTPUT_HANDLE) ||
-            (handles[i] == (HANDLE)STD_ERROR_HANDLE))
-            hloc[i] = GetStdHandle( HandleToULong(handles[i]) );
-        else
-            hloc[i] = handles[i];
-
-        /* yes, even screen buffer console handles are waitable, and are
-         * handled as a handle to the console itself !!
-         */
-        if (is_console_handle(hloc[i]))
-        {
-            if (VerifyConsoleIoHandle(hloc[i]))
-                hloc[i] = GetConsoleInputWaitHandle();
-        }
-    }
 
     status = NtWaitForMultipleObjects( count, hloc, !wait_all, alertable,
                                        get_nt_timeout( &time, timeout ) );
@@ -209,6 +210,7 @@ BOOL WINAPI RegisterWaitForSingleObject(PHANDLE phNewWaitObject, HANDLE hObject,
     TRACE("%p %p %p %p %d %d\n",
           phNewWaitObject,hObject,Callback,Context,dwMilliseconds,dwFlags);
 
+    hObject = normalize_handle_if_console(hObject);
     status = RtlRegisterWait( phNewWaitObject, hObject, Callback, Context, dwMilliseconds, dwFlags );
     if (status != STATUS_SUCCESS)
     {
@@ -231,6 +233,7 @@ HANDLE WINAPI RegisterWaitForSingleObjectEx( HANDLE hObject,
     TRACE("%p %p %p %d %d\n",
           hObject,Callback,Context,dwMilliseconds,dwFlags);
 
+    hObject = normalize_handle_if_console(hObject);
     status = RtlRegisterWait( &hNewWaitObject, hObject, Callback, Context, dwMilliseconds, dwFlags );
     if (status != STATUS_SUCCESS)
     {

dlls/kernel32/tests/console.c

@@ -908,6 +908,58 @@ static void testScreenBuffer(HANDLE hConOut)
     SetConsoleOutputCP(oldcp);
 }
 
+static void CALLBACK signaled_function(void *p, BOOLEAN timeout)
+{
+    HANDLE event = p;
+    SetEvent(event);
+    ok(!timeout, "wait shouldn't have timed out\n");
+}
+
+static void testWaitForConsoleInput(HANDLE input_handle)
+{
+    HANDLE wait_handle;
+    HANDLE complete_event;
+    INPUT_RECORD record;
+    DWORD events_written;
+    DWORD wait_ret;
+    BOOL ret;
+
+    complete_event = CreateEventW(NULL, FALSE, FALSE, NULL);
+
+    /* Test success case */
+    ret = RegisterWaitForSingleObject(&wait_handle, input_handle, signaled_function, complete_event, INFINITE, WT_EXECUTEONLYONCE);
+    ok(ret == TRUE, "Expected RegisterWaitForSingleObject to return TRUE, got %d\n", ret);
+    /* give worker thread a chance to start up */
+    Sleep(100);
+    record.EventType = KEY_EVENT;
+    record.Event.KeyEvent.bKeyDown = 1;
+    record.Event.KeyEvent.wRepeatCount = 1;
+    record.Event.KeyEvent.wVirtualKeyCode = VK_RETURN;
+    record.Event.KeyEvent.wVirtualScanCode = VK_RETURN;
+    record.Event.KeyEvent.uChar.UnicodeChar = '\r';
+    record.Event.KeyEvent.dwControlKeyState = 0;
+    ret = WriteConsoleInputW(input_handle, &record, 1, &events_written);
+    ok(ret == TRUE, "Expected WriteConsoleInputW to return TRUE, got %d\n", ret);
+    wait_ret = WaitForSingleObject(complete_event, INFINITE);
+    ok(wait_ret == WAIT_OBJECT_0, "Expected the handle to be signaled\n");
+    ret = UnregisterWait(wait_handle);
+    /* If the callback is still running, this fails with ERROR_IO_PENDING, but
+       that's ok and expected. */
+    ok(ret != 0 || GetLastError() == ERROR_IO_PENDING,
+        "UnregisterWait failed with error %d\n", GetLastError());
+
+    /* Test timeout case */
+    FlushConsoleInputBuffer(input_handle);
+    ret = RegisterWaitForSingleObject(&wait_handle, input_handle, signaled_function, complete_event, INFINITE, WT_EXECUTEONLYONCE);
+    wait_ret = WaitForSingleObject(complete_event, 100);
+    ok(wait_ret == WAIT_TIMEOUT, "Expected the wait to time out\n");
+    ret = UnregisterWait(wait_handle);
+    ok(ret, "UnregisterWait failed with error %d\n", GetLastError());
+
+    /* Clean up */
+    ok(CloseHandle(complete_event), "Failed to close event handle, last error %d\n", GetLastError());
+}
+
 static void test_GetSetConsoleInputExeName(void)
 {
     BOOL ret;
@@ -3093,6 +3145,8 @@ START_TEST(console)
     testScroll(hConOut, sbi.dwSize);
     /* will test sb creation / modification / codepage handling */
     testScreenBuffer(hConOut);
+    /* Test waiting for a console handle */
+    testWaitForConsoleInput(hConIn);
 
     /* clear duplicated console font table */
     CloseHandle(hConIn);