From 2d51d6168874cdb2b72ee4f56f414d9a9a9d92e5 Mon Sep 17 00:00:00 2001
From: unknown <diego.nehab@gmail.com>
Date: Tue, 28 May 2013 17:27:06 +0800
Subject: [PATCH] Fix "final" bug in pton and TCP connreset handling

---
 src/inet.c    |  9 ++++-----
 src/wsocket.c | 37 +++++++++++++++++++++++--------------
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/src/inet.c b/src/inet.c
index 1f55d2a..1c44464 100644
--- a/src/inet.c
+++ b/src/inet.c
@@ -559,12 +559,11 @@ const char *inet_ntop(int af, const void *src, char *dst, socklen_t cnt)
 int inet_pton(int af, const char *src, void *dst) 
 {
     struct addrinfo hints, *res;
+    int ret = 1;
     memset(&hints, 0, sizeof(struct addrinfo));
     hints.ai_family = af;
     hints.ai_flags = AI_NUMERICHOST;
-    if (getaddrinfo(src, NULL, &hints, &res) != 0) {
-        return -1;
-    }
+    if (getaddrinfo(src, NULL, &hints, &res) != 0) return -1;
     if (af == AF_INET) {
         struct sockaddr_in *in = (struct sockaddr_in *) res->ai_addr;
         memcpy(dst, &in->sin_addr, sizeof(in->sin_addr));
@@ -572,10 +571,10 @@ int inet_pton(int af, const char *src, void *dst)
         struct sockaddr_in6 *in = (struct sockaddr_in6 *) res->ai_addr;
         memcpy(dst, &in->sin6_addr, sizeof(in->sin6_addr));
     } else {
-        return -1;
+        ret = -1;
     }
     freeaddrinfo(res); 
-    return 0;
+    return ret;
 }
 
 #endif
diff --git a/src/wsocket.c b/src/wsocket.c
index d34724b..b4a4384 100644
--- a/src/wsocket.c
+++ b/src/wsocket.c
@@ -238,8 +238,10 @@ int socket_sendto(p_socket ps, const char *data, size_t count, size_t *sent,
 /*-------------------------------------------------------------------------*\
 * Receive with timeout
 \*-------------------------------------------------------------------------*/
-int socket_recv(p_socket ps, char *data, size_t count, size_t *got, p_timeout tm) {
-    int err;
+int socket_recv(p_socket ps, char *data, size_t count, size_t *got, 
+        p_timeout tm) 
+{
+    int err, prev = IO_DONE;
     *got = 0;
     if (*ps == SOCKET_INVALID) return IO_CLOSED;
     for ( ;; ) {
@@ -250,11 +252,14 @@ int socket_recv(p_socket ps, char *data, size_t count, size_t *got, p_timeout tm
         }
         if (taken == 0) return IO_CLOSED;
         err = WSAGetLastError();
-        /* On Windows, and on UDP, a connreset simply means the
-         * previous send failed. On TCP, it means our socket
-         * is now useless, so the error must pass. I am
-         * hoping waitfd will still get the error. */
-        if (err != WSAEWOULDBLOCK && err != WSAECONNRESET) return err;
+        /* On UDP, a connreset simply means the previous send failed. 
+         * So we try again. 
+         * On TCP, it means our socket is now useless, so the error passes. 
+         * (We will loop again, exiting because the same error will happen) */
+        if (err != WSAEWOULDBLOCK) {
+            if (err != WSAECONNRESET || prev == WSAECONNRESET) return err;
+            prev = err;
+        }
         if ((err = socket_waitfd(ps, WAITFD_R, tm)) != IO_DONE) return err;
     }
 }
@@ -263,8 +268,9 @@ int socket_recv(p_socket ps, char *data, size_t count, size_t *got, p_timeout tm
 * Recvfrom with timeout
 \*-------------------------------------------------------------------------*/
 int socket_recvfrom(p_socket ps, char *data, size_t count, size_t *got, 
-        SA *addr, socklen_t *len, p_timeout tm) {
-    int err;
+        SA *addr, socklen_t *len, p_timeout tm) 
+{
+    int err, prev = IO_DONE;
     *got = 0;
     if (*ps == SOCKET_INVALID) return IO_CLOSED;
     for ( ;; ) {
@@ -275,11 +281,14 @@ int socket_recvfrom(p_socket ps, char *data, size_t count, size_t *got,
         }
         if (taken == 0) return IO_CLOSED;
         err = WSAGetLastError();
-        /* On Windows, and on UDP, a connreset simply means the
-         * previous send failed. On TCP, it means our socket
-         * is now useless, so the error must pass. I am
-         * hoping waitfd will still get the error. */
-        if (err != WSAEWOULDBLOCK && err != WSAECONNRESET) return err;
+        /* On UDP, a connreset simply means the previous send failed. 
+         * So we try again. 
+         * On TCP, it means our socket is now useless, so the error passes.
+         * (We will loop again, exiting because the same error will happen) */
+        if (err != WSAEWOULDBLOCK) {
+            if (err != WSAECONNRESET || prev == WSAECONNRESET) return err;
+            prev = err;
+        }
         if ((err = socket_waitfd(ps, WAITFD_R, tm)) != IO_DONE) return err;
     }
 }