From c2ecb4acb52c9f9b2e77af86a38a9b36ee7d51e6 Mon Sep 17 00:00:00 2001
From: Dmitry Ilyin <dima@doty.ru>
Date: Wed, 12 Oct 2022 14:13:44 +0300
Subject: [PATCH] Add locks for server WS, fixes #1357

---
 include/event2/ws.h     |  2 +-
 sample/ws-chat-server.c |  2 +-
 test/regress_ws.c       |  2 +-
 ws.c                    | 27 +++++++++++++++++++++++----
 4 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/include/event2/ws.h b/include/event2/ws.h
index 0816a1266e..11fc2e5880 100644
--- a/include/event2/ws.h
+++ b/include/event2/ws.h
@@ -28,7 +28,7 @@ typedef void (*ws_on_close_cb)(struct evws_connection *, void *);
  */
 EVENT2_EXPORT_SYMBOL
 struct evws_connection *evws_new_session(
-	struct evhttp_request *req, ws_on_msg_cb, void *arg);
+	struct evhttp_request *req, ws_on_msg_cb, void *arg, int options);
 
 /** Sends data over WebSocket connection */
 EVENT2_EXPORT_SYMBOL
diff --git a/sample/ws-chat-server.c b/sample/ws-chat-server.c
index 8761df3185..731d1281e7 100644
--- a/sample/ws-chat-server.c
+++ b/sample/ws-chat-server.c
@@ -152,7 +152,7 @@ on_ws(struct evhttp_request *req, void *arg)
 	socklen_t len;
 
 	client = calloc(sizeof(*client), 1);
-	client->evws = evws_new_session(req, on_msg_cb, client);
+	client->evws = evws_new_session(req, on_msg_cb, client, 0);
 	fd = bufferevent_getfd(evws_connection_get_bufferevent(client->evws));
 
 	len = sizeof(addr);
diff --git a/test/regress_ws.c b/test/regress_ws.c
index 06c3f5f066..29850be8ca 100644
--- a/test/regress_ws.c
+++ b/test/regress_ws.c
@@ -115,7 +115,7 @@ http_on_ws_cb(struct evhttp_request *req, void *arg)
 	struct evws_connection *evws;
 	const char *hello = "Server: hello";
 
-	evws = evws_new_session(req, on_ws_msg_cb, (void *)0xDEADBEEF);
+	evws = evws_new_session(req, on_ws_msg_cb, (void *)0xDEADBEEF, 0);
 	if (!evws)
 		return;
 	test_ok++;
diff --git a/ws.c b/ws.c
index 8c9314daa2..c8a806cefc 100644
--- a/ws.c
+++ b/ws.c
@@ -12,6 +12,7 @@
 #include "event2/bufferevent.h"
 #include "sys/queue.h"
 #include "http-internal.h"
+#include "bufferevent-internal.h"
 
 #include <assert.h>
 #include <string.h>
@@ -269,16 +270,18 @@ ws_evhttp_read_cb(struct bufferevent *bufev, void *arg)
 	int msg_len, in_len, header_sz;
 	struct evbuffer *input = bufferevent_get_input(evws->bufev);
 
+	bufferevent_incref_and_lock_(evws->bufev);
+	bufferevent_disable(evws->bufev, EV_WRITE);
 	while ((in_len = evbuffer_get_length(input))) {
 		unsigned char *data = evbuffer_pullup(input, in_len);
 		if (data == NULL) {
-			return;
+			goto bailout;
 		}
 
 		type = get_ws_frame(data, in_len, &payload, &msg_len);
 		if (type == INCOMPLETE_DATA) {
 			/* incomplete data received, wait for next chunk */
-			return;
+			goto bailout;
 		}
 		header_sz = payload - data;
 		evbuffer_drain(input, header_sz);
@@ -324,6 +327,9 @@ ws_evhttp_read_cb(struct bufferevent *bufev, void *arg)
 		}
 		evbuffer_drain(input, msg_len);
 	}
+
+bailout:
+	bufferevent_decref_and_unlock_(evws->bufev);
 }
 
 static void
@@ -336,7 +342,8 @@ ws_evhttp_error_cb(struct bufferevent *bufev, short what, void *arg)
 }
 
 struct evws_connection *
-evws_new_session(struct evhttp_request *req, ws_on_msg_cb cb, void *arg)
+evws_new_session(
+	struct evhttp_request *req, ws_on_msg_cb cb, void *arg, int options)
 {
 	struct evws_connection *evws = NULL;
 	struct evkeyvalq *in_hdrs;
@@ -380,6 +387,12 @@ evws_new_session(struct evhttp_request *req, ws_on_msg_cb cb, void *arg)
 	evws->http_server = evcon->http_server;
 
 	evws->bufev = evhttp_start_ws_(req);
+
+	if (options & BEV_OPT_THREADSAFE) {
+		if (bufferevent_enable_locking_(evws->bufev, NULL) < 0)
+			goto error;
+	}
+
 	bufferevent_setcb(
 		evws->bufev, ws_evhttp_read_cb, NULL, ws_evhttp_error_cb, evws);
 
@@ -389,6 +402,9 @@ evws_new_session(struct evhttp_request *req, ws_on_msg_cb cb, void *arg)
 	return evws;
 
 error:
+	if (evws)
+		evws_connection_free(evws);
+
 	evhttp_send_reply(req, HTTP_BADREQUEST, NULL, NULL);
 	return NULL;
 }
@@ -419,9 +435,12 @@ make_ws_frame(struct evbuffer *output, enum WebSocketFrameType frame_type,
 void
 evws_send(struct evws_connection *evws, const char *packet_str, size_t str_len)
 {
-	struct evbuffer *output = bufferevent_get_output(evws->bufev);
+	struct evbuffer *output;
 
+	bufferevent_lock(evws->bufev);
+	output = bufferevent_get_output(evws->bufev);
 	make_ws_frame(output, TEXT_FRAME, (unsigned char *)packet_str, str_len);
+	bufferevent_unlock(evws->bufev);
 }
 
 void