[PATCH v2 2/2] ws2_32: Make wait in WS2_recv_base() alertable.

[Paul Gofman]

From: Paul Gofman

Signed-off-by: Paul Gofman 
---
 dlls/ws2_32/socket.c     |  2 +-
 dlls/ws2_32/tests/sock.c | 75 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 70 insertions(+), 7 deletions(-)

diff --git a/dlls/ws2_32/socket.c b/dlls/ws2_32/socket.c
index 9be7190062f..9cc75aacf20 100644
--- a/dlls/ws2_32/socket.c
+++ b/dlls/ws2_32/socket.c
@@ -959,7 +959,7 @@ static int WS2_recv_base( SOCKET s, WSABUF *buffers, DWORD buffer_count, DWORD *
                                     IOCTL_AFD_WINE_RECVMSG, &params, sizeof(params), NULL, 0 );
     if (status == STATUS_PENDING && !overlapped)
     {
-        if (WaitForSingleObject( event, INFINITE ) == WAIT_FAILED)
+        if (wait_event_alertable( event ) == WAIT_FAILED)
             return -1;
         status = piosb->u.Status;
     }
diff --git a/dlls/ws2_32/tests/sock.c b/dlls/ws2_32/tests/sock.c
index 80d458b650d..af4226e6258 100644
--- a/dlls/ws2_32/tests/sock.c
+++ b/dlls/ws2_32/tests/sock.c
@@ -6747,21 +6747,58 @@ static void test_WSASendTo(void)
     ok(ret_addr.sin_port, "expected nonzero port\n");
 }

+struct recv_thread_apc_param
+{
+    SOCKET sock;
+    unsigned int apc_count;
+};
+
+static void WINAPI recv_thread_apc_func(ULONG_PTR param)
+{
+    struct recv_thread_apc_param *p = (struct recv_thread_apc_param *)param;
+    int ret;
+
+    ++p->apc_count;
+
+    ret = send(p->sock, "test", 4, 0);
+    ok(ret == 4, "got %d.\n", ret);
+}
+
+struct recv_thread_param
+{
+    SOCKET sock;
+    BOOL overlapped;
+};
+
 static DWORD WINAPI recv_thread(LPVOID arg)
 {
-    SOCKET sock = *(SOCKET *)arg;
+    struct recv_thread_param *p = arg;
+    SOCKET sock = p->sock;
     char buffer[32];
     WSABUF wsa;
     WSAOVERLAPPED ov;
     DWORD flags = 0;
+    DWORD len;
+    int ret;

wsa.buf = buffer;
     wsa.len = sizeof(buffer);
-    ov.hEvent = WSACreateEvent();
-    WSARecv(sock, &wsa, 1, NULL, &flags, &ov, NULL);
+    if (p->overlapped)
+    {
+        ov.hEvent = WSACreateEvent();
+        WSARecv(sock, &wsa, 1, NULL, &flags, &ov, NULL);

-    WaitForSingleObject(ov.hEvent, 1000);
-    WSACloseEvent(ov.hEvent);
+        WaitForSingleObject(ov.hEvent, 1000);
+        WSACloseEvent(ov.hEvent);
+    }
+    else
+    {
+        SetLastError(0xdeadbeef);
+        ret = WSARecv(sock, &wsa, 1, &len, &flags, NULL, NULL);
+        ok(!ret, "got ret %d.\n", ret);
+        ok(WSAGetLastError() == 0, "got error %d.\n", WSAGetLastError());
+        ok(len == 4, "got len %lu.\n", len);
+    }
     return 0;
 }

@@ -6775,11 +6812,14 @@ static void WINAPI io_completion(DWORD error, DWORD transferred, WSAOVERLAPPED *
 static void test_WSARecv(void)
 {
     SOCKET src, dest, server = INVALID_SOCKET;
+    struct recv_thread_apc_param apc_param;
+    struct recv_thread_param recv_param;
     char buf[20];
     WSABUF bufs[2];
     WSAOVERLAPPED ov;
     DWORD bytesReturned, flags, id;
     struct sockaddr_in addr;
+    unsigned int apc_count;
     int iret, len;
     DWORD dwret;
     BOOL bret;
@@ -6799,10 +6839,20 @@ static void test_WSARecv(void)
     ok(GetLastError() == ERROR_SUCCESS, "Expected 0, got %ld\n", GetLastError());
     SetLastError(0xdeadbeef);
     bytesReturned = 0xdeadbeef;
+
+    apc_count = 0;
+    dwret = QueueUserAPC(apc_func, GetCurrentThread(), (ULONG_PTR)&apc_count);
+    ok(dwret, "QueueUserAPC returned %lu\n", dwret);
+
     iret = WSARecv(dest, bufs, 1, &bytesReturned, &flags, NULL, NULL);
     ok(!iret, "Expected 0, got %d\n", iret);
     ok(bytesReturned == 2, "Expected 2, got %ld\n", bytesReturned);
     ok(GetLastError() == ERROR_SUCCESS, "Expected 0, got %ld\n", GetLastError());
+
+    ok(!apc_count, "got apc_count %u.\n", apc_count);
+    SleepEx(0, TRUE);
+    ok(apc_count == 1, "got apc_count %u.\n", apc_count);
+
     SetLastError(0xdeadbeef);
     bytesReturned = 0xdeadbeef;
     iret = WSARecv(dest, bufs, 1, &bytesReturned, &flags, NULL, NULL);
@@ -6895,10 +6945,23 @@ static void test_WSARecv(void)
     if (dest == INVALID_SOCKET) goto end;

send(src, "test message", sizeof("test message"), 0);
-    thread = CreateThread(NULL, 0, recv_thread, &dest, 0, &id);
+    recv_param.sock = dest;
+    recv_param.overlapped = TRUE;
+    thread = CreateThread(NULL, 0, recv_thread, &recv_param, 0, &id);
     WaitForSingleObject(thread, 3000);
     CloseHandle(thread);

+    recv_param.overlapped = FALSE;
+    thread = CreateThread(NULL, 0, recv_thread, &recv_param, 0, &id);
+    apc_param.apc_count = 0;
+    apc_param.sock = src;
+    dwret = QueueUserAPC(recv_thread_apc_func, thread, (ULONG_PTR)&apc_param);
+    ok(dwret, "QueueUserAPC returned %lu\n", dwret);
+    WaitForSingleObject(thread, 3000);
+    ok(apc_param.apc_count == 1, "got apc_count %u.\n", apc_param.apc_count);
+
+    CloseHandle(thread);
+
     memset(&ov, 0, sizeof(ov));
     ov.hEvent = event;
     ResetEvent(event);