ws2_32: Only enable FD_WRITE on short sends.

dlls/ws2_32/socket.c

@@ -1250,7 +1250,12 @@ static NTSTATUS WS2_async_send(void* user, IO_STATUS_BLOCK* iosb, NTSTATUS statu
 
         if (result >= 0)
         {
+            int totalLength = 0;
+            int i;
             status = STATUS_SUCCESS;
+            for (i = 0; i < wsa->n_iovecs; i++)
+                totalLength += wsa->iovec[i].iov_len;
+            if (result < totalLength)
                 _enable_event( wsa->hSocket, FD_WRITE, 0, 0 );
         }
         else
@@ -2667,6 +2672,7 @@ INT WINAPI WSASendTo( SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount,
     unsigned int i, options;
     int n, fd, err;
     struct iovec iovec[WS_MSG_MAXIOVLEN];
+    int totalLength = 0;
     ULONG_PTR cvalue = (lpOverlapped && ((ULONG_PTR)lpOverlapped->hEvent & 1) == 0) ? (ULONG_PTR)lpOverlapped : 0;
 
     TRACE("socket %04lx, wsabuf %p, nbufs %d, flags %d, to %p, tolen %d, ovl %p, func %p\n",
@@ -2694,6 +2700,7 @@ INT WINAPI WSASendTo( SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount,
     {
         iovec[i].iov_base = lpBuffers[i].buf;
         iovec[i].iov_len  = lpBuffers[i].len;
+        totalLength += lpBuffers[i].len;
     }
 
     for (;;)
@@ -2819,6 +2826,7 @@ INT WINAPI WSASendTo( SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount,
     }
     else  /* non-blocking */
     {
+        if (n < totalLength)
             _enable_event(SOCKET2HANDLE(s), FD_WRITE, 0, 0);
         if (n == -1)
         {

dlls/ws2_32/tests/sock.c

@@ -1993,9 +1993,23 @@ static DWORD WINAPI drain_socket_thread(LPVOID arg)
 {
     char buffer[1024];
     SOCKET sock = *(SOCKET*)arg;
+    int ret;
 
-    while (recv(sock, buffer, sizeof(buffer), 0) > 0)
+    while ((ret = recv(sock, buffer, sizeof(buffer), 0)) != 0)
-        ;
+    {
+        if (ret < 0)
+        {
+            if (WSAGetLastError() == WSAEWOULDBLOCK)
+            {
+                fd_set readset;
+                FD_ZERO(&readset);
+                FD_SET(sock, &readset);
+                select(0, &readset, NULL, NULL, NULL);
+            }
+            else
+                break;
+        }
+    }
     return 0;
 }
 
@@ -2051,10 +2065,13 @@ static void test_write_events(void)
     SOCKET dst = INVALID_SOCKET;
     HANDLE hThread = NULL;
     HANDLE hEvent = INVALID_HANDLE_VALUE;
-    int len;
+    char *buffer = NULL;
+    int bufferSize = 1024*1024;
     u_long one = 1;
     int ret;
     DWORD id;
+    WSANETWORKEVENTS netEvents;
+    DWORD dwRet;
 
     if (tcp_socketpair(&src, &dst) != 0)
     {
@@ -2062,6 +2079,29 @@ static void test_write_events(void)
         return;
     }
 
+    /* On Windows it seems when a non-blocking socket sends to a
+       blocking socket on the same host, the send() is BLOCKING,
+       so make both sockets non-blocking */
+    ret = ioctlsocket(src, FIONBIO, &one);
+    if (ret)
+    {
+        ok(0, "ioctlsocket failed, error %d\n", WSAGetLastError());
+        goto end;
+    }
+    ret = ioctlsocket(dst, FIONBIO, &one);
+    if (ret)
+    {
+        ok(0, "ioctlsocket failed, error %d\n", WSAGetLastError());
+        goto end;
+    }
+
+    buffer = HeapAlloc(GetProcessHeap(), 0, bufferSize);
+    if (buffer == NULL)
+    {
+        ok(0, "could not allocate memory for test\n");
+        goto end;
+    }
+
     hThread = CreateThread(NULL, 0, drain_socket_thread, &dst, 0, &id);
     if (hThread == NULL)
     {
@@ -2076,13 +2116,6 @@ static void test_write_events(void)
         goto end;
     }
 
-    ret = ioctlsocket(src, FIONBIO, &one);
-    if (ret)
-    {
-        ok(0, "ioctlsocket failed, error %d\n", WSAGetLastError());
-        goto end;
-    }
-
     ret = WSAEventSelect(src, hEvent, FD_WRITE | FD_CLOSE);
     if (ret)
     {
@@ -2090,41 +2123,72 @@ static void test_write_events(void)
         goto end;
     }
 
-    for (len = 100; len > 0; --len)
+    /* FD_WRITE should be set initially, and allow us to send at least 1 byte */
-    {
+    dwRet = WaitForSingleObject(hEvent, 5000);
-         WSANETWORKEVENTS netEvents;
-         DWORD dwRet = WaitForSingleObject(hEvent, 5000);
     if (dwRet != WAIT_OBJECT_0)
     {
-             ok(0, "WaitForSingleObject failed, error %d\n", dwRet);
+        ok(0, "Initial WaitForSingleObject failed, error %d\n", dwRet);
         goto end;
     }
-
     ret = WSAEnumNetworkEvents(src, NULL, &netEvents);
     if (ret)
     {
         ok(0, "WSAEnumNetworkEvents failed, error %d\n", ret);
         goto end;
     }
-
     if (netEvents.lNetworkEvents & FD_WRITE)
     {
         ret = send(src, "a", 1, 0);
-             if (ret < 0 && WSAGetLastError() != WSAEWOULDBLOCK)
+        ok(ret == 1, "sending 1 byte failed, error %d\n", WSAGetLastError());
-             {
+        if (ret != 1)
-                 ok(0, "send failed, error %d\n", WSAGetLastError());
             goto end;
     }
+    else
+    {
+        ok(0, "FD_WRITE not among initial events\n");
+        goto end;
     }
 
-         if (netEvents.lNetworkEvents & FD_CLOSE)
+    /* Now FD_WRITE should not be set, because the socket send buffer isn't full yet */
+    dwRet = WaitForSingleObject(hEvent, 2000);
+    if (dwRet == WAIT_OBJECT_0)
     {
-             ok(0, "unexpected close\n");
+        ok(0, "WaitForSingleObject should have timed out, but succeeded!\n");
         goto end;
     }
+
+    /* Now if we send a tonne of data, the socket send buffer will only take some of it,
+       and we will get a short write, which will trigger another FD_WRITE event
+       as soon as data is sent and more space becomes available, but not any earlier. */
+    do
+    {
+        ret = send(src, buffer, bufferSize, 0);
+    } while (ret == bufferSize);
+    if (ret >= 0 || WSAGetLastError() == WSAEWOULDBLOCK)
+    {
+        dwRet = WaitForSingleObject(hEvent, 5000);
+        ok(dwRet == WAIT_OBJECT_0, "Waiting failed with %d\n", dwRet);
+        if (dwRet == WAIT_OBJECT_0)
+        {
+            ret = WSAEnumNetworkEvents(src, NULL, &netEvents);
+            ok(ret == 0, "WSAEnumNetworkEvents failed, error %d\n", ret);
+            if (ret == 0)
+                goto end;
+            ok(netEvents.lNetworkEvents & FD_WRITE,
+                "FD_WRITE event not set as expected, events are 0x%x\n", netEvents.lNetworkEvents);
+        }
+        else
+            goto end;
+    }
+    else
+    {
+        ok(0, "sending a lot of data failed with error %d\n", WSAGetLastError());
+        goto end;
     }
 
 end:
+    if (buffer != NULL)
+        HeapFree(GetProcessHeap(), 0, buffer);
     if (src != INVALID_SOCKET)
         closesocket(src);
     if (dst != INVALID_SOCKET)