Fix potential memory leak

- Got rid of ifdef _WIN32 stuff adn moved to plat_ files instead.
- Also, check all calls to lws_zalloc, was potential failure on WIN32
- Made context destory enable to destroy a half inited context as well. This way I got get rid of some of the error handling complexity in libwebsocket_create_context
- Added TODOs for some potential problems I see where things might be leaking and such

diff --git a/lib/context.c b/lib/context.c
index 02207dc..f1c8dd0 100644 (file)
--- a/lib/context.c
+++ b/lib/context.c
@@ -143,37 +143,15 @@ libwebsocket_create_context(struct lws_context_creation_info *info)
        if (context->fds == NULL) {
                lwsl_err("Unable to allocate fds array for %d connections\n",
                                                              context->max_fds);
-               lws_free(context);
-               return NULL;
+               goto bail;
        }
 
-#ifdef _WIN32
-       for (i = 0; i < FD_HASHTABLE_MODULUS; i++) {
-               context->fd_hashtable[i].wsi = lws_zalloc(sizeof(struct libwebsocket*) * context->max_fds);
-       }
-#else
-       context->lws_lookup = lws_zalloc(sizeof(struct libwebsocket *) * context->max_fds);
-       if (context->lws_lookup == NULL) {
-               lwsl_err(
-                 "Unable to allocate lws_lookup array for %d connections\n",
-                                                             context->max_fds);
-               lws_free(context->fds);
-               lws_free(context);
-               return NULL;
+       if (lws_plat_init_lookup(context)) {
+               goto bail;
        }
-#endif
 
        if (lws_plat_init_fd_tables(context)) {
-#ifdef _WIN32
-               for (i = 0; i < FD_HASHTABLE_MODULUS; i++) {
-                       lws_free(context->fd_hashtable[i].wsi);
-               }
-#else
-               lws_free(context->lws_lookup);
-#endif
-               lws_free(context->fds);
-               lws_free(context);
-               return NULL;
+               goto bail;
        }
 
        lws_context_init_extensions(info, context);
@@ -294,11 +272,16 @@ bail:
 LWS_VISIBLE void
 libwebsocket_context_destroy(struct libwebsocket_context *context)
 {
+       /* Note that this is used for freeing partially allocated structs as well
+        * so make sure you don't try to free something uninitialized */
        int n;
-       struct libwebsocket_protocols *protocol = context->protocols;
+       struct libwebsocket_protocols *protocol = NULL;
 
        lwsl_notice("%s\n", __func__);
 
+       if (!context)
+               return;
+
 #ifdef LWS_LATENCY
        if (context->worst_latency_info[0])
                lwsl_notice("Worst latency: %s\n", context->worst_latency_info);
@@ -318,38 +301,41 @@ libwebsocket_context_destroy(struct libwebsocket_context *context)
         * give all extensions a chance to clean up any per-context
         * allocations they might have made
         */
+       // TODO: I am not sure, but are we never supposed to be able to run a server
+       //       and client at the same time for a given context?
+       //       Otherwise both of these callbacks should always be called!
        if (context->listen_port != CONTEXT_PORT_NO_LISTEN) {
                if (lws_ext_callback_for_each_extension_type(context, NULL,
-                        LWS_EXT_CALLBACK_SERVER_CONTEXT_DESTRUCT, NULL, 0) < 0)
-                       return;
-       } else
+                               LWS_EXT_CALLBACK_SERVER_CONTEXT_DESTRUCT, NULL, 0) < 0) {
+                       lwsl_err("Got error from server extension callback on cleanup");
+               }
+       } else {
                if (lws_ext_callback_for_each_extension_type(context, NULL,
-                        LWS_EXT_CALLBACK_CLIENT_CONTEXT_DESTRUCT, NULL, 0) < 0)
-                       return;
+                               LWS_EXT_CALLBACK_CLIENT_CONTEXT_DESTRUCT, NULL, 0) < 0) {
+                       lwsl_err("Got error from client extension callback on cleanup");
+               }
+       }
 
        /*
         * inform all the protocols that they are done and will have no more
         * callbacks
         */
-
-       while (protocol->callback) {
-               protocol->callback(context, NULL, LWS_CALLBACK_PROTOCOL_DESTROY,
-                               NULL, NULL, 0);
-               protocol++;
+       protocol = context->protocols;
+       if (protocol) {
+               while (protocol->callback) {
+                       protocol->callback(context, NULL, LWS_CALLBACK_PROTOCOL_DESTROY,
+                                       NULL, NULL, 0);
+                       protocol++;
+               }
        }
 
        lws_plat_context_early_destroy(context);
 
        lws_ssl_context_destroy(context);
 
-       lws_free(context->fds);
-#ifdef _WIN32
-       for (n = 0; n < FD_HASHTABLE_MODULUS; n++) {
-               lws_free(context->fd_hashtable[n].wsi);
-       }
-#else
-       lws_free(context->lws_lookup);
-#endif
+       if (context->fds)
+               lws_free(context->fds);
+
        lws_plat_context_late_destroy(context);
 
        lws_free(context);

diff --git a/lib/libwebsockets.c b/lib/libwebsockets.c
index 25ff022..cb6f389 100644 (file)
--- a/lib/libwebsockets.c
+++ b/lib/libwebsockets.c
@@ -38,6 +38,29 @@ static const char * const log_level_names[] = {
 };
 
 void
+lws_free_wsi(struct libwebsocket *wsi)
+{
+       if (!wsi)
+               return;
+
+       /* Protocol user data may be allocated either internally by lws
+        * or by specified the user. Important we don't free external user data */
+       if (wsi->protocol && wsi->protocol->per_session_data_size
+               && wsi->user_space && !wsi->user_space_externally_allocated) {
+               lws_free(wsi->user_space);
+       }
+
+       lws_free2(wsi->rxflow_buffer);
+       lws_free2(wsi->truncated_send_malloc);
+
+       // TODO: Probably should handle the union structs in wsi->u here depending
+       //       on connection mode as well. Too spaghetti for me to follow however...
+
+       lws_free_header_table(wsi);
+       lws_free(wsi);
+}
+
+void
 libwebsocket_close_and_free_session(struct libwebsocket_context *context,
                         struct libwebsocket *wsi, enum lws_close_status reason)
 {
@@ -89,7 +112,6 @@ libwebsocket_close_and_free_session(struct libwebsocket_context *context,
                context->protocols[0].callback(context, wsi,
                        LWS_CALLBACK_CLIENT_CONNECTION_ERROR, wsi->user_space, NULL, 0);
 
-               lws_free_header_table(wsi);
                goto just_kill_connection;
        }
 
@@ -99,6 +121,7 @@ libwebsocket_close_and_free_session(struct libwebsocket_context *context,
 
        if (wsi->mode == LWS_CONNMODE_HTTP_SERVING_ACCEPTED) {
                if (wsi->u.http.fd != LWS_INVALID_FILE) {
+                       // TODO: If we're just closing with LWS_CLOSE_STATUS_NOSTATUS_CONTEXT_DESTROY this file descriptor might leak?
                        lwsl_debug("closing http file\n");
                        compatible_file_close(wsi->u.http.fd);
                        wsi->u.http.fd = LWS_INVALID_FILE;
@@ -218,10 +241,7 @@ just_kill_connection:
        wsi->state = WSI_STATE_DEAD_SOCKET;
 
        lws_free2(wsi->rxflow_buffer);
-
-       if (wsi->mode == LWS_CONNMODE_HTTP2_SERVING && wsi->u.hdr.ah) {
-               lws_free2(wsi->u.hdr.ah);
-       }
+       lws_free_header_table(wsi);
 
        if ((old_state == WSI_STATE_ESTABLISHED ||
             wsi->mode == LWS_CONNMODE_WS_SERVING ||
@@ -298,13 +318,7 @@ just_kill_connection:
        context->protocols[0].callback(context, wsi,
                        LWS_CALLBACK_WSI_DESTROY, wsi->user_space, NULL, 0);
 
-       if (wsi->protocol && wsi->protocol->per_session_data_size &&
-           wsi->user_space && !wsi->user_space_externally_allocated)
-               lws_free(wsi->user_space);
-
-       /* As a precaution, free the header table in case it lingered: */
-       lws_free_header_table(wsi);
-       lws_free(wsi);
+       lws_free_wsi(wsi);
 }
 
 LWS_VISIBLE int

diff --git a/lib/lws-plat-unix.c b/lib/lws-plat-unix.c
index a7f96a3..be642aa 100644 (file)
--- a/lib/lws-plat-unix.c
+++ b/lib/lws-plat-unix.c
@@ -272,6 +272,20 @@ lws_plat_drop_app_privileges(struct lws_context_creation_info *info)
 }
 
 LWS_VISIBLE int
+lws_plat_init_lookup(struct libwebsocket_context *context)
+{
+       context->lws_lookup = lws_zalloc(sizeof(struct libwebsocket *) * context->max_fds);
+       if (context->lws_lookup == NULL) {
+               lwsl_err(
+                 "Unable to allocate lws_lookup array for %d connections\n",
+                                                             context->max_fds);
+               return 1;
+       }
+
+       return 0;
+}
+
+LWS_VISIBLE int
 lws_plat_init_fd_tables(struct libwebsocket_context *context)
 {
        context->fd_random = open(SYSTEM_RANDOM_FILEPATH, O_RDONLY);
@@ -328,6 +342,9 @@ lws_plat_context_early_destroy(struct libwebsocket_context *context)
 LWS_VISIBLE void
 lws_plat_context_late_destroy(struct libwebsocket_context *context)
 {
+       if (context->lws_lookup)
+               lws_free(context->lws_lookup);
+
        close(context->dummy_pipe_fds[0]);
        close(context->dummy_pipe_fds[1]);
        close(context->fd_random);

diff --git a/lib/lws-plat-win.c b/lib/lws-plat-win.c
index 863c716..122ec84 100644 (file)
--- a/lib/lws-plat-win.c
+++ b/lib/lws-plat-win.c
@@ -249,6 +249,22 @@ lws_plat_drop_app_privileges(struct lws_context_creation_info *info)
 }
 
 LWS_VISIBLE int
+lws_plat_init_lookup(struct libwebsocket_context *context)
+{
+       int i;
+
+       for (i = 0; i < FD_HASHTABLE_MODULUS; i++) {
+               context->fd_hashtable[i].wsi = lws_zalloc(sizeof(struct libwebsocket*) * context->max_fds);
+
+               if (!context->fd_hashtable[i].wsi) {
+                       return -1;
+               }
+       }
+
+       return 0;
+}
+
+LWS_VISIBLE int
 lws_plat_init_fd_tables(struct libwebsocket_context *context)
 {
        context->events = lws_malloc(sizeof(WSAEVENT) * (context->max_fds + 1));
@@ -300,6 +316,13 @@ lws_plat_context_early_destroy(struct libwebsocket_context *context)
 LWS_VISIBLE void
 lws_plat_context_late_destroy(struct libwebsocket_context *context)
 {
+       int n;
+
+       for (n = 0; n < FD_HASHTABLE_MODULUS; n++) {
+               if (context->fd_hashtable[n].wsi)
+                       lws_free(context->fd_hashtable[n].wsi);
+       }
+
        WSACleanup();
 }
 

diff --git a/lib/private-libwebsockets.h b/lib/private-libwebsockets.h
index 8f6d043..85e750b 100644 (file)
--- a/lib/private-libwebsockets.h
+++ b/lib/private-libwebsockets.h
@@ -1245,6 +1245,8 @@ lws_poll_listen_fd(struct libwebsocket_pollfd *fd);
 LWS_EXTERN int
 lws_plat_service(struct libwebsocket_context *context, int timeout_ms);
 LWS_EXTERN int
+lws_plat_init_lookup(struct libwebsocket_context *context);
+LWS_EXTERN int
 lws_plat_init_fd_tables(struct libwebsocket_context *context);
 LWS_EXTERN void
 lws_plat_drop_app_privileges(struct lws_context_creation_info *info);