close should flush any truncated send before really closing

Signed-off-by: Andy Green <andy.green@linaro.org>

diff --git a/lib/extension.c b/lib/extension.c
index 19d68cb..176faf0 100644 (file)
--- a/lib/extension.c
+++ b/lib/extension.c
@@ -130,8 +130,10 @@ lws_issue_raw_ext_access(struct libwebsocket *wsi,
                if (eff_buf.token_len) {
                        n = lws_issue_raw(wsi, (unsigned char *)eff_buf.token,
                                                            eff_buf.token_len);
-                       if (n < 0)
+                       if (n < 0) {
+                               lwsl_info("closing from ext access\n");
                                return -1;
+                       }
 
                        /* always either sent it all or privately buffered */
                }

diff --git a/lib/libwebsockets.c b/lib/libwebsockets.c
index d8fa1e3..d2f06af 100644 (file)
--- a/lib/libwebsockets.c
+++ b/lib/libwebsockets.c
@@ -53,13 +53,30 @@ libwebsocket_close_and_free_session(struct libwebsocket_context *context,
 
        old_state = wsi->state;
 
-       if (old_state == WSI_STATE_DEAD_SOCKET)
+       switch (old_state) {
+       case WSI_STATE_DEAD_SOCKET:
                return;
 
        /* we tried the polite way... */
-       if (old_state == WSI_STATE_AWAITING_CLOSE_ACK)
+       case WSI_STATE_AWAITING_CLOSE_ACK:
                goto just_kill_connection;
 
+       case WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE:
+               if (wsi->truncated_send_len) {
+                       libwebsocket_callback_on_writable(context, wsi);
+                       return;
+               }
+               lwsl_info("wsi %p completed WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE\n", wsi);
+               goto just_kill_connection;
+       default:
+               if (wsi->truncated_send_len) {
+                       lwsl_info("wsi %p entering WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE\n", wsi);
+                       wsi->state = WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE;
+                       return;
+               }
+               break;
+       }
+
        wsi->u.ws.close_reason = reason;
 
        if (wsi->mode == LWS_CONNMODE_WS_CLIENT_WAITING_CONNECT ||

diff --git a/lib/output.c b/lib/output.c
index 6c060c9..b235954 100644 (file)
--- a/lib/output.c
+++ b/lib/output.c
@@ -99,6 +99,10 @@ int lws_issue_raw(struct libwebsocket *wsi, unsigned char *buf, size_t len)
        
        if (!len)
                return 0;
+       /* just ignore sends after we cleared the truncation buffer */
+       if (wsi->state == WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE &&
+                                               !wsi->truncated_send_len)
+               return len;
 
        if (wsi->truncated_send_len && (buf < wsi->truncated_send_malloc ||
                        buf > (wsi->truncated_send_malloc +
@@ -149,6 +153,9 @@ handle_truncated_send:
                        lwsl_info("***** %x partial send completed\n", wsi);
                        /* done with it, but don't free it */
                        n = real_len;
+                       if (wsi->state == WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE)
+                               lwsl_info("***** %x signalling to close now\n", wsi);
+                               return -1; /* retry closing now */
                }
                /* always callback on writeable */
                libwebsocket_callback_on_writable(
@@ -480,9 +487,12 @@ LWS_VISIBLE int libwebsockets_serve_http_file_fragment(
        while (!lws_send_pipe_choked(wsi)) {
 
                if (wsi->truncated_send_len) {
-                       lws_issue_raw(wsi, wsi->truncated_send_malloc +
+                       if (lws_issue_raw(wsi, wsi->truncated_send_malloc +
                                        wsi->truncated_send_offset,
-                                                      wsi->truncated_send_len);
+                                                      wsi->truncated_send_len) < 0) {
+                               lwsl_info("closing from libwebsockets_serve_http_file_fragment\n");
+                               return -1;
+                       }
                        continue;
                }
 

diff --git a/lib/private-libwebsockets.h b/lib/private-libwebsockets.h
index be6c2d9..94bd227 100644 (file)
--- a/lib/private-libwebsockets.h
+++ b/lib/private-libwebsockets.h
@@ -287,6 +287,7 @@ enum lws_connection_states {
        WSI_STATE_CLIENT_UNCONNECTED,
        WSI_STATE_RETURNED_CLOSE_ALREADY,
        WSI_STATE_AWAITING_CLOSE_ACK,
+       WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE,
 };
 
 enum lws_rx_parse_state {

diff --git a/lib/server.c b/lib/server.c
index ed2dd20..2d2ffb3 100644 (file)
--- a/lib/server.c
+++ b/lib/server.c
@@ -506,9 +506,12 @@ int lws_server_socket_service(struct libwebsocket_context *context,
 
                if (wsi->truncated_send_malloc) {
                        if (pollfd->revents & LWS_POLLOUT)
-                               lws_issue_raw(wsi, wsi->truncated_send_malloc +
+                               if (lws_issue_raw(wsi, wsi->truncated_send_malloc +
                                        wsi->truncated_send_offset,
-                                                       wsi->truncated_send_len);
+                                                       wsi->truncated_send_len) < 0) {
+                                       lwsl_info("closing from socket service\n");
+                                       return -1;
+                               }
                        /*
                         * we can't afford to allow input processing send
                         * something new, so spin around he event loop until
@@ -539,16 +542,20 @@ int lws_server_socket_service(struct libwebsocket_context *context,
                                break;
                        }
 
-                       /* hm this may want to send (via HTTP callback for example) */
+                       /* just ignore incoming if waiting for close */
+                       if (wsi->state != WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE) {
+                       
+                               /* hm this may want to send (via HTTP callback for example) */
 
-                       n = libwebsocket_read(context, wsi,
-                                               context->service_buffer, len);
-                       if (n < 0)
-                               /* we closed wsi */
-                               return 0;
+                               n = libwebsocket_read(context, wsi,
+                                                       context->service_buffer, len);
+                               if (n < 0)
+                                       /* we closed wsi */
+                                       return 0;
 
-                       /* hum he may have used up the writability above */
-                       break;
+                               /* hum he may have used up the writability above */
+                               break;
+                       }
                }
 
                /* this handles POLLOUT for http serving fragments */

diff --git a/lib/service.c b/lib/service.c
index 5edc2a3..85b6753 100644 (file)
--- a/lib/service.c
+++ b/lib/service.c
@@ -34,9 +34,12 @@ lws_handle_POLLOUT_event(struct libwebsocket_context *context,
        /* pending truncated sends have uber priority */
 
        if (wsi->truncated_send_len) {
-               lws_issue_raw(wsi, wsi->truncated_send_malloc +
+               if (lws_issue_raw(wsi, wsi->truncated_send_malloc +
                                wsi->truncated_send_offset,
-                                               wsi->truncated_send_len);
+                                               wsi->truncated_send_len) < 0) {
+                       lwsl_info("lws_handle_POLLOUT_event signalling to close\n");
+                       return -1;
+               }
                /* leave POLLOUT active either way */
                return 0;
        }
@@ -87,8 +90,10 @@ lws_handle_POLLOUT_event(struct libwebsocket_context *context,
                if (eff_buf.token_len) {
                        n = lws_issue_raw(wsi, (unsigned char *)eff_buf.token,
                                                             eff_buf.token_len);
-                       if (n < 0)
+                       if (n < 0) {
+                               lwsl_info("closing from POLLOUT spill\n");
                                return -1;
+                       }
                        /*
                         * Keep amount spilled small to minimize chance of this
                         */
@@ -340,6 +345,8 @@ libwebsocket_service_fd(struct libwebsocket_context *context,
        case LWS_CONNMODE_SERVER_LISTENER:
        case LWS_CONNMODE_SSL_ACK_PENDING:
                n = lws_server_socket_service(context, wsi, pollfd);
+               if (n < 0)
+                       goto close_and_handled;
                goto handled;
 
        case LWS_CONNMODE_WS_SERVING:
@@ -348,7 +355,8 @@ libwebsocket_service_fd(struct libwebsocket_context *context,
                /* the guy requested a callback when it was OK to write */
 
                if ((pollfd->revents & LWS_POLLOUT) &&
-                       wsi->state == WSI_STATE_ESTABLISHED &&
+                       (wsi->state == WSI_STATE_ESTABLISHED ||
+                               wsi->state == WSI_STATE_FLUSHING_STORED_SEND_BEFORE_CLOSE) &&
                           lws_handle_POLLOUT_event(context, wsi, pollfd)) {
                        lwsl_info("libwebsocket_service_fd: closing\n");
                        goto close_and_handled;