Merge pull request #6 from VerityIncorporated/main

Fix: Store and expose WebSocket headers for non-native clients

Author: ProjectSky

src/ws_server.cpp

@@ -59,6 +59,14 @@ void WebSocketServer::OnMessage(const std::string& message, std::shared_ptr<ix::
 		if (!pWebSocketClient->m_websocket_handle) return;
 		pWebSocketClient->m_keepConnecting = true;
 
+		{
+			std::lock_guard<std::mutex> lock(m_headersMutex);
+			auto it = m_connectionHeaders.find(connectionState->getId());
+			if (it != m_connectionHeaders.end()) {
+				pWebSocketClient->m_headers = it->second;
+			}
+		}
+
 		std::string remoteAddress = connectionState->getRemoteIp() + ":" + std::to_string(connectionState->getRemotePort());
 		pMessageForward->PushCell(m_webSocketServer_handle);
 		pMessageForward->PushCell(pWebSocketClient->m_websocket_handle);
@@ -79,6 +87,11 @@ void WebSocketServer::OnOpen(ix::WebSocketOpenInfo openInfo, std::shared_ptr<ix:
 		return;
 	}
 
+	{
+		std::lock_guard<std::mutex> lock(m_headersMutex);
+		m_connectionHeaders[connectionState->getId()] = openInfo.headers;
+	}
+
 	g_TaskQueue.Push([this, openInfo, connectionState]()
 	{
 		std::string remoteAddress = connectionState->getRemoteIp() + ":" + std::to_string(connectionState->getRemotePort());
@@ -95,6 +108,13 @@ void WebSocketServer::OnClose(ix::WebSocketCloseInfo closeInfo, std::shared_ptr<
 	{
 		return;
 	}
+	
+	std::string connectionId = connectionState->getId();
+
+	{
+		std::lock_guard<std::mutex> lock(m_headersMutex);
+		m_connectionHeaders.erase(connectionId);
+	}
 
 	g_TaskQueue.Push([this, closeInfo, connectionState]()
 	{
@@ -162,13 +182,13 @@ bool WebSocketServer::disconnectClient(const std::string& clientId) {
 }
 
 std::vector<std::string> WebSocketServer::getClientIds() {
-    std::vector<std::string> clientIds;
+	std::vector<std::string> clientIds;
 
-    auto clients = m_webSocketServer.getClients();
+	auto clients = m_webSocketServer.getClients();
 
-    for (const auto& client : clients) {
-        clientIds.push_back(client.second);
-    }
+	for (const auto& client : clients) {
+		clientIds.push_back(client.second);
+	}
 
-    return clientIds;
+	return clientIds;
 }

src/ws_server.h

@@ -1,6 +1,5 @@
 #include "extension.h"
 
-
 class WebSocketServer
 {
 public:
@@ -20,6 +19,9 @@ class WebSocketServer
 	ix::WebSocketServer m_webSocketServer;
 	Handle_t m_webSocketServer_handle = BAD_HANDLE;
 
+    std::mutex m_headersMutex;
+    std::unordered_map<std::string, ix::WebSocketHttpHeaders> m_connectionHeaders;
+
 	IChangeableForward *pMessageForward = nullptr;
 	IChangeableForward *pOpenForward = nullptr;
 	IChangeableForward *pCloseForward = nullptr;