recv() may not return all data

While performing handshake recv() is called only once.
It may return only part of the data and handshake
will fail. This patch modifies libwebsocket_service_fd()
to ensure that there is not data left in the socket.

Signed-off-by: Pavel Borzenkov <pavel.borzenkov@auriga.com>

diff --git a/lib/client-handshake.c b/lib/client-handshake.c
index de306a6..9b502cf 100644 (file)
--- a/lib/client-handshake.c
+++ b/lib/client-handshake.c
@@ -34,6 +34,7 @@ libwebsocket_client_connect(struct libwebsocket_context *context,
 {
        struct hostent *server_hostent;
        struct sockaddr_in server_addr;
+       struct timeval tv;
        char pkt[512];
        struct pollfd pfd;
        struct libwebsocket *wsi;
@@ -168,6 +169,12 @@ libwebsocket_client_connect(struct libwebsocket_context *context,
        /* Disable Nagle */
        setsockopt(wsi->sock, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(opt));
 
+       /* Set receiving timeout */
+       tv.tv_sec = 0;
+       tv.tv_usec = 100 * 1000;
+       setsockopt(wsi->sock, SOL_SOCKET, SO_RCVTIMEO, (char *)&tv,
+                       sizeof(tv));
+
        if (connect(wsi->sock, (struct sockaddr *)&server_addr,
                                              sizeof(struct sockaddr)) == -1)  {
                fprintf(stderr, "Connect failed\n");

diff --git a/lib/libwebsockets.c b/lib/libwebsockets.c
index 05f7a40..4d7b3d3 100644 (file)
--- a/lib/libwebsockets.c
+++ b/lib/libwebsockets.c
@@ -614,7 +614,19 @@ user_service:
        return 0;
 }
 
+static ssize_t
+libwebsocket_recv(int sock, void *buf, size_t len, int flags)
+{
+       int n, ret = 0;
+
+       while ((len > 0) && (n = recv(sock, (char *)buf + ret, len,
+                       flags)) > 0) {
+               ret += n;
+               len -= n;
+       }
 
+       return ret;
+}
 
 /**
  * libwebsocket_service_fd() - Service polled socket with something waiting
@@ -1008,7 +1020,7 @@ libwebsocket_service_fd(struct libwebsocket_context *context,
                        return 1;
                }
 
-               n = recv(wsi->sock, pkt, sizeof pkt, 0);
+               n = libwebsocket_recv(wsi->sock, pkt, sizeof pkt, 0);
                if (n < 0) {
                        libwebsocket_close_and_free_session(context, wsi,
                                                     LWS_CLOSE_STATUS_NOSTATUS);
@@ -1352,7 +1364,7 @@ issue_hdr:
                        len = SSL_read(wsi->ssl, pkt, sizeof pkt);
                else
        #endif
-                       len = recv(wsi->sock, pkt, sizeof pkt, 0);
+               len = libwebsocket_recv(wsi->sock, pkt, sizeof pkt, 0);
 
                if (len < 0) {
                        fprintf(stderr,
@@ -1737,7 +1749,8 @@ bail2:
                else
 #endif
                        eff_buf.token_len =
-                                          recv(pollfd->fd, buf, sizeof buf, 0);
+                                       libwebsocket_recv(pollfd->fd, buf,
+                                       sizeof buf, 0);
 
                if (eff_buf.token_len < 0) {
                        fprintf(stderr, "Socket read returned %d\n",