1 files changed, 127 insertions, 53 deletions

diff --git a/src/zenhttp/servers/httpsys.cpp b/src/zenhttp/servers/httpsys.cpp
index e93ae4853..23d57af57 100644
--- a/src/zenhttp/servers/httpsys.cpp
+++ b/src/zenhttp/servers/httpsys.cpp
@@ -156,6 +156,10 @@ private:
 
 #if ZEN_WITH_HTTPSYS
 
+#	include "httpsys_iocontext.h"
+#	include "wshttpsys.h"
+#	include "wsframecodec.h"
+
 #	include <conio.h>
 #	include <mstcpip.h>
 #	pragma comment(lib, "httpapi.lib")
@@ -380,7 +384,7 @@ public:
 
 	PTP_IO				  Iocp();
 	HANDLE				  RequestQueueHandle();
-	inline OVERLAPPED*	  Overlapped() { return &m_HttpOverlapped; }
+	inline OVERLAPPED*	  Overlapped() { return &m_IoContext.Overlapped; }
 	inline HttpSysServer& Server() { return m_HttpServer; }
 	inline HTTP_REQUEST*  HttpRequest() { return m_InitialHttpHandler.HttpRequest(); }
 
@@ -397,8 +401,8 @@ public:
 	};
 
 private:
-	OVERLAPPED	   m_HttpOverlapped{};
-	HttpSysServer& m_HttpServer;
+	HttpSysIoContext m_IoContext{};
+	HttpSysServer&	 m_HttpServer;
 
 	// Tracks which handler is due to handle the next I/O completion event
 	HttpSysRequestHandler*				m_CompletionHandler = nullptr;
@@ -1555,7 +1559,23 @@ HttpSysTransaction::IoCompletionCallback(PTP_CALLBACK_INSTANCE Instance,
 	// than one thread at any given moment. This means we need to be careful about what
 	// happens in here
 
-	HttpSysTransaction* Transaction = CONTAINING_RECORD(pOverlapped, HttpSysTransaction, m_HttpOverlapped);
+	HttpSysIoContext* IoContext = CONTAINING_RECORD(pOverlapped, HttpSysIoContext, Overlapped);
+
+	switch (IoContext->ContextType)
+	{
+		case HttpSysIoContext::Type::kWebSocketRead:
+			static_cast<WsHttpSysConnection*>(IoContext->Owner)->OnReadCompletion(IoResult, NumberOfBytesTransferred);
+			return;
+
+		case HttpSysIoContext::Type::kWebSocketWrite:
+			static_cast<WsHttpSysConnection*>(IoContext->Owner)->OnWriteCompletion(IoResult, NumberOfBytesTransferred);
+			return;
+
+		case HttpSysIoContext::Type::kTransaction:
+			break;
+	}
+
+	HttpSysTransaction* Transaction = CONTAINING_RECORD(IoContext, HttpSysTransaction, m_IoContext);
 
 	if (Transaction->HandleCompletion(IoResult, NumberOfBytesTransferred) == HttpSysTransaction::Status::kDone)
 	{
@@ -2111,64 +2131,118 @@ InitialRequestHandler::HandleCompletion(ULONG IoResult, ULONG_PTR NumberOfBytesT
 	{
 		HTTP_REQUEST* HttpReq = HttpRequest();
 
-#	if 0
-		for (int i = 0; i < HttpReq->RequestInfoCount; ++i)
+		if (HttpService* Service = reinterpret_cast<HttpService*>(HttpReq->UrlContext))
 		{
-			auto& ReqInfo = HttpReq->pRequestInfo[i];
-
-			switch (ReqInfo.InfoType)
+			// WebSocket upgrade detection
+			if (m_IsInitialRequest)
 			{
-				case HttpRequestInfoTypeRequestTiming:
+				const HTTP_KNOWN_HEADER& UpgradeHeader = HttpReq->Headers.KnownHeaders[HttpHeaderUpgrade];
+				if (UpgradeHeader.RawValueLength > 0 &&
+					StrCaseCompare(UpgradeHeader.pRawValue, "websocket", UpgradeHeader.RawValueLength) == 0)
+				{
+					if (IWebSocketHandler* WsHandler = dynamic_cast<IWebSocketHandler*>(Service))
 					{
-						const HTTP_REQUEST_TIMING_INFO* TimingInfo = reinterpret_cast<HTTP_REQUEST_TIMING_INFO*>(ReqInfo.pInfo);
+						// Extract Sec-WebSocket-Key from the unknown headers
+						// (http.sys has no known-header slot for it)
+						std::string_view SecWebSocketKey;
+						for (USHORT i = 0; i < HttpReq->Headers.UnknownHeaderCount; ++i)
+						{
+							const HTTP_UNKNOWN_HEADER& Hdr = HttpReq->Headers.pUnknownHeaders[i];
+							if (Hdr.NameLength == 17 && _strnicmp(Hdr.pName, "Sec-WebSocket-Key", 17) == 0)
+							{
+								SecWebSocketKey = std::string_view(Hdr.pRawValue, Hdr.RawValueLength);
+								break;
+							}
+						}
 
-						ZEN_INFO("");
-					}
-					break;
-				case HttpRequestInfoTypeAuth:
-					ZEN_INFO("");
-					break;
-				case HttpRequestInfoTypeChannelBind:
-					ZEN_INFO("");
-					break;
-				case HttpRequestInfoTypeSslProtocol:
-					ZEN_INFO("");
-					break;
-				case HttpRequestInfoTypeSslTokenBindingDraft:
-					ZEN_INFO("");
-					break;
-				case HttpRequestInfoTypeSslTokenBinding:
-					ZEN_INFO("");
-					break;
-				case HttpRequestInfoTypeTcpInfoV0:
-					{
-						const TCP_INFO_v0* TcpInfo = reinterpret_cast<const TCP_INFO_v0*>(ReqInfo.pInfo);
+						if (SecWebSocketKey.empty())
+						{
+							ZEN_WARN("WebSocket upgrade missing Sec-WebSocket-Key header");
+							return nullptr;
+						}
 
-						ZEN_INFO("");
-					}
-					break;
-				case HttpRequestInfoTypeRequestSizing:
-					{
-						const HTTP_REQUEST_SIZING_INFO* SizingInfo = reinterpret_cast<const HTTP_REQUEST_SIZING_INFO*>(ReqInfo.pInfo);
-						ZEN_INFO("");
-					}
-					break;
-				case HttpRequestInfoTypeQuicStats:
-					ZEN_INFO("");
-					break;
-				case HttpRequestInfoTypeTcpInfoV1:
-					{
-						const TCP_INFO_v1* TcpInfo = reinterpret_cast<const TCP_INFO_v1*>(ReqInfo.pInfo);
+						const std::string AcceptKey = WsFrameCodec::ComputeAcceptKey(SecWebSocketKey);
+
+						HANDLE			RequestQueueHandle = Transaction().RequestQueueHandle();
+						HTTP_REQUEST_ID RequestId		   = HttpReq->RequestId;
+
+						// Build the 101 Switching Protocols response
+						HTTP_RESPONSE Response = {};
+						Response.StatusCode	   = 101;
+						Response.pReason	   = "Switching Protocols";
+						Response.ReasonLength  = (USHORT)strlen(Response.pReason);
+
+						Response.Headers.KnownHeaders[HttpHeaderUpgrade].pRawValue		= "websocket";
+						Response.Headers.KnownHeaders[HttpHeaderUpgrade].RawValueLength = 9;
+
+						eastl::fixed_vector<HTTP_UNKNOWN_HEADER, 8> UnknownHeaders;
+
+						// IMPORTANT: Due to some quirk in HttpSendHttpResponse, this cannot use KnownHeaders
+						// despite there being an entry for it there (HttpHeaderConnection). If you try to do
+						// that you get an ERROR_INVALID_PARAMETERS error from HttpSendHttpResponse below
+
+						UnknownHeaders.push_back({.NameLength = 10, .RawValueLength = 7, .pName = "Connection", .pRawValue = "Upgrade"});
+
+						UnknownHeaders.push_back({.NameLength	  = 20,
+												  .RawValueLength = (USHORT)AcceptKey.size(),
+												  .pName		  = "Sec-WebSocket-Accept",
+												  .pRawValue	  = AcceptKey.c_str()});
+
+						Response.Headers.UnknownHeaderCount = (USHORT)UnknownHeaders.size();
+						Response.Headers.pUnknownHeaders	= UnknownHeaders.data();
+
+						const ULONG Flags = HTTP_SEND_RESPONSE_FLAG_OPAQUE | HTTP_SEND_RESPONSE_FLAG_MORE_DATA;
+
+						// Use an OVERLAPPED with an event so we can wait synchronously.
+						// The request queue is IOCP-associated, so passing NULL for pOverlapped
+						// may return ERROR_IO_PENDING.  Setting the low-order bit of hEvent
+						// prevents IOCP delivery and lets us wait on the event directly.
+						OVERLAPPED SendOverlapped = {};
+						HANDLE	   SendEvent	  = CreateEventW(nullptr, TRUE, FALSE, nullptr);
+						SendOverlapped.hEvent	  = (HANDLE)((uintptr_t)SendEvent | 1);
+
+						ULONG SendResult = HttpSendHttpResponse(RequestQueueHandle,
+																RequestId,
+																Flags,
+																&Response,
+																nullptr,  // CachePolicy
+																nullptr,  // BytesSent
+																nullptr,  // Reserved1
+																0,		  // Reserved2
+																&SendOverlapped,
+																nullptr	 // LogData
+						);
+
+						if (SendResult == ERROR_IO_PENDING)
+						{
+							WaitForSingleObject(SendEvent, INFINITE);
+							SendResult = (SendOverlapped.Internal == 0) ? NO_ERROR : ERROR_IO_INCOMPLETE;
+						}
+
+						CloseHandle(SendEvent);
+
+						if (SendResult == NO_ERROR)
+						{
+							Ref<WsHttpSysConnection> WsConn(
+								new WsHttpSysConnection(RequestQueueHandle, RequestId, *WsHandler, Transaction().Iocp()));
+							Ref<WebSocketConnection> WsConnRef(WsConn.Get());
+
+							WsHandler->OnWebSocketOpen(std::move(WsConnRef));
+							WsConn->Start();
+
+							return nullptr;
+						}
 
-						ZEN_INFO("");
+						ZEN_WARN("WebSocket 101 send failed: {}", SendResult);
+
+						// WebSocket upgrade failed — return nullptr since ServerRequest()
+						// was never populated (no InvokeRequestHandler call)
+						return nullptr;
 					}
-					break;
+					// Service doesn't support WebSocket or missing key — fall through to normal handling
+				}
 			}
-		}
-#	endif
 
-		if (HttpService* Service = reinterpret_cast<HttpService*>(HttpReq->UrlContext))
-		{
 			if (m_IsInitialRequest)
 			{
 				m_ContentLength					  = GetContentLength(HttpReq);