From ffaac93ba6e85f71eaea3693779cda7afb77223b Mon Sep 17 00:00:00 2001
From: tt2468 <tt2468@gmail.com>
Date: Fri, 30 Apr 2021 08:56:08 -0700
Subject: [PATCH] WebSocketSession: Add mutex for operations which change
 session parameters

---
 src/WebSocketProtocol.cpp | 3 ++-
 src/WebSocketServer.cpp   | 3 +++
 src/WebSocketSession.h    | 2 ++
 3 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/WebSocketProtocol.cpp b/src/WebSocketProtocol.cpp
index 37eebeb5..ebd223b5 100644
--- a/src/WebSocketProtocol.cpp
+++ b/src/WebSocketProtocol.cpp
@@ -49,6 +49,7 @@ WebSocketProtocol::ProcessResult WebSocketProtocol::ProcessMessage(SessionPtr se
 	} else if (messageType == "RequestBatch") {
 		;
 	} else if (messageType == "Identify") {
+		std::unique_lock<std::mutex> sessionLock(session->OperationMutex);
 		if (session->IsIdentified()) {
 			if (!session->IgnoreInvalidMessages()) {
 				ret.closeCode = WebSocketServer::WebSocketCloseCode::AlreadyIdentified;
@@ -87,7 +88,7 @@ WebSocketProtocol::ProcessResult WebSocketProtocol::ProcessMessage(SessionPtr se
 		ret.result["negotiatedRpcVersion"] = session->RpcVersion();
 		return ret;
 	} else if (messageType == "Reidentify") {
-		;
+		std::unique_lock<std::mutex> sessionLock(session->OperationMutex);
 	} else {
 		if (!session->IgnoreInvalidMessages()) {
 			ret.closeCode = WebSocketServer::WebSocketCloseCode::UnknownMessageType;
diff --git a/src/WebSocketServer.cpp b/src/WebSocketServer.cpp
index ab7fd353..8307f60a 100644
--- a/src/WebSocketServer.cpp
+++ b/src/WebSocketServer.cpp
@@ -253,6 +253,7 @@ void WebSocketServer::onOpen(websocketpp::connection_hdl hdl)
 	// Build new session
 	std::unique_lock<std::mutex> lock(_sessionMutex);
 	SessionPtr session = _sessions[hdl] = std::make_shared<WebSocketSession>();
+	std::unique_lock<std::mutex> sessionLock(session->OperationMutex);
 	lock.unlock();
 
 	// Configure session details
@@ -284,6 +285,8 @@ void WebSocketServer::onOpen(websocketpp::connection_hdl hdl)
 		helloMessage["authentication"]["salt"] = AuthenticationSalt;
 	}
 
+	sessionLock.unlock();
+
 	// Send object to client
 	websocketpp::lib::error_code errorCode;
 	auto sessionEncoding = session->Encoding();
diff --git a/src/WebSocketSession.h b/src/WebSocketSession.h
index c280a253..a9c58dd4 100644
--- a/src/WebSocketSession.h
+++ b/src/WebSocketSession.h
@@ -42,6 +42,8 @@ class WebSocketSession
 		uint64_t EventSubscriptions();
 		void SetEventSubscriptions(uint64_t subscriptions);
 
+		std::mutex OperationMutex;
+
 	private:
 		std::mutex _remoteAddressMutex;
 		std::string _remoteAddress;