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

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);
+	if (lws_plat_init_lookup(context)) {
+		goto bail;
 	}
-#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;
-	}
-#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);

lib/libwebsockets.c

@@ -37,6 +37,29 @@ static const char * const log_level_names[] = {
 	"LATENCY",
 };
 
+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

lib/lws-plat-unix.c

@@ -271,6 +271,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)
 {
@@ -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);

lib/lws-plat-win.c

@@ -248,6 +248,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)
 {
@@ -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();
 }
 

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);