1 files changed, 49 insertions, 0 deletions
diff --git a/src/zenhttp/servers/httpasio.cpp b/src/zenhttp/servers/httpasio.cpp
index 0c0238886..8c2dcd116 100644
--- a/src/zenhttp/servers/httpasio.cpp
+++ b/src/zenhttp/servers/httpasio.cpp
@@ -14,6 +14,8 @@
 #include <zenhttp/httpserver.h>
 
 #include "httpparser.h"
+#include "wsasio.h"
+#include "wsframecodec.h"
 
 #include <EASTL/fixed_vector.h>
 
@@ -1159,6 +1161,53 @@ HttpServerConnection::HandleRequest()
 {
 	ZEN_MEMSCOPE(GetHttpasioTag());
 
+	// WebSocket upgrade detection must happen before the keep-alive check below,
+	// because Upgrade requests have "Connection: Upgrade" which the HTTP parser
+	// treats as non-keep-alive, causing a premature shutdown of the receive side.
+	if (m_RequestData.IsWebSocketUpgrade())
+	{
+		if (HttpService* Service = m_Server.RouteRequest(m_RequestData.Url()))
+		{
+			IWebSocketHandler* WsHandler = dynamic_cast<IWebSocketHandler*>(Service);
+			if (WsHandler && !m_RequestData.SecWebSocketKey().empty())
+			{
+				std::string AcceptKey = WsFrameCodec::ComputeAcceptKey(m_RequestData.SecWebSocketKey());
+
+				auto ResponseStr = std::make_shared<std::string>();
+				ResponseStr->reserve(256);
+				ResponseStr->append(
+					"HTTP/1.1 101 Switching Protocols\r\n"
+					"Upgrade: websocket\r\n"
+					"Connection: Upgrade\r\n"
+					"Sec-WebSocket-Accept: ");
+				ResponseStr->append(AcceptKey);
+				ResponseStr->append("\r\n\r\n");
+
+				// Send the 101 response on the current socket, then hand the socket off
+				// to a WsAsioConnection for the WebSocket protocol.
+				asio::async_write(*m_Socket,
+								  asio::buffer(ResponseStr->data(), ResponseStr->size()),
+								  [Conn = AsSharedPtr(), WsHandler, OwnedResponse = ResponseStr](const asio::error_code& Ec, std::size_t) {
+									  if (Ec)
+									  {
+										  ZEN_WARN("WebSocket 101 send failed: {}", Ec.message());
+										  return;
+									  }
+
+									  Ref<WsAsioConnection>	   WsConn(new WsAsioConnection(std::move(Conn->m_Socket), *WsHandler));
+									  Ref<WebSocketConnection> WsConnRef(WsConn.Get());
+
+									  WsHandler->OnWebSocketOpen(std::move(WsConnRef));
+									  WsConn->Start();
+								  });
+
+				m_RequestState = RequestState::kDone;
+				return;
+			}
+		}
+		// Service doesn't support WebSocket or missing key — fall through to normal handling
+	}
+
 	if (!m_RequestData.IsKeepAlive())
 	{
 		m_RequestState = RequestState::kWritingFinal;