1 files changed, 41 insertions, 6 deletions

diff --git a/src/zenhttp/servers/httpsys.cpp b/src/zenhttp/servers/httpsys.cpp
index 2cad97725..c1b426bea 100644
--- a/src/zenhttp/servers/httpsys.cpp
+++ b/src/zenhttp/servers/httpsys.cpp
@@ -464,6 +464,8 @@ public:
 	inline int64_t	GetResponseBodySize() const { return m_TotalDataSize; }
 
 	void SetLocationHeader(std::string_view Location) { m_LocationHeader = Location; }
+	void SetContentTypeOverride(std::string Override) { m_ContentTypeOverride = std::move(Override); }
+	void SetContentRangeHeader(std::string V) { m_ContentRangeHeader = std::move(V); }
 
 private:
 	eastl::fixed_vector<HTTP_DATA_CHUNK, 16> m_HttpDataChunks;
@@ -473,6 +475,8 @@ private:
 	uint32_t								 m_RemainingChunkCount = 0;	 // Backlog for multi-call sends
 	bool									 m_IsInitialResponse   = true;
 	HttpContentType							 m_ContentType		   = HttpContentType::kBinary;
+	std::string								 m_ContentTypeOverride;
+	std::string								 m_ContentRangeHeader;
 	eastl::fixed_vector<IoBuffer, 16>		 m_DataBuffers;
 	std::string								 m_LocationHeader;
 
@@ -725,7 +729,8 @@ HttpMessageResponseRequest::IssueRequest(std::error_code& ErrorCode)
 
 		PHTTP_KNOWN_HEADER ContentTypeHeader = &HttpResponse.Headers.KnownHeaders[HttpHeaderContentType];
 
-		std::string_view ContentTypeString = MapContentTypeToString(m_ContentType);
+		std::string_view ContentTypeString =
+			m_ContentTypeOverride.empty() ? MapContentTypeToString(m_ContentType) : std::string_view(m_ContentTypeOverride);
 
 		ContentTypeHeader->pRawValue	  = ContentTypeString.data();
 		ContentTypeHeader->RawValueLength = (USHORT)ContentTypeString.size();
@@ -739,6 +744,15 @@ HttpMessageResponseRequest::IssueRequest(std::error_code& ErrorCode)
 			LocationHeader->RawValueLength	  = (USHORT)m_LocationHeader.size();
 		}
 
+		// Content-Range header (for 206 Partial Content single-range responses)
+
+		if (!m_ContentRangeHeader.empty())
+		{
+			PHTTP_KNOWN_HEADER ContentRangeHeader = &HttpResponse.Headers.KnownHeaders[HttpHeaderContentRange];
+			ContentRangeHeader->pRawValue		  = m_ContentRangeHeader.data();
+			ContentRangeHeader->RawValueLength	  = (USHORT)m_ContentRangeHeader.size();
+		}
+
 		std::string_view ReasonString = ReasonStringForHttpResultCode(m_ResponseCode);
 
 		HttpResponse.StatusCode	  = m_ResponseCode;
@@ -1258,7 +1272,7 @@ HttpSysServer::RegisterHttpUrls(int BasePort)
 					else if (InternalResult == ERROR_SHARING_VIOLATION || InternalResult == ERROR_ACCESS_DENIED)
 					{
 						// Port may be owned by another process's wildcard registration (access denied)
-						// or actively in use (sharing violation) — retry on a different port
+						// or actively in use (sharing violation) - retry on a different port
 						ShouldRetryNextPort = true;
 					}
 					else
@@ -1713,7 +1727,7 @@ HttpSysServer::OnRun(bool IsInteractive)
 
 		ShutdownRequested = m_ShutdownEvent.Wait(WaitTimeout);
 		UpdateLofreqTimerValue();
-	} while (!ShutdownRequested);
+	} while (!ShutdownRequested && !IsApplicationExitRequested());
 }
 
 void
@@ -2279,6 +2293,11 @@ HttpSysServerRequest::WriteResponse(HttpResponseCode ResponseCode)
 
 	HttpMessageResponseRequest* Response = new HttpMessageResponseRequest(m_HttpTx, (uint16_t)ResponseCode);
 
+	if (!m_ContentRangeHeader.empty())
+	{
+		Response->SetContentRangeHeader(std::move(m_ContentRangeHeader));
+	}
+
 	if (SuppressBody())
 	{
 		Response->SuppressResponseBody();
@@ -2307,6 +2326,15 @@ HttpSysServerRequest::WriteResponse(HttpResponseCode ResponseCode, HttpContentTy
 
 	HttpMessageResponseRequest* Response = new HttpMessageResponseRequest(m_HttpTx, (uint16_t)ResponseCode, ContentType, Blobs);
 
+	if (!m_ContentTypeOverride.empty())
+	{
+		Response->SetContentTypeOverride(std::move(m_ContentTypeOverride));
+	}
+	if (!m_ContentRangeHeader.empty())
+	{
+		Response->SetContentRangeHeader(std::move(m_ContentRangeHeader));
+	}
+
 	if (SuppressBody())
 	{
 		Response->SuppressResponseBody();
@@ -2595,7 +2623,14 @@ InitialRequestHandler::HandleCompletion(ULONG IoResult, ULONG_PTR NumberOfBytesT
 																					&Transaction().Server()));
 							Ref<WebSocketConnection> WsConnRef(WsConn.Get());
 
-							WsHandler->OnWebSocketOpen(std::move(WsConnRef));
+							ExtendableStringBuilder<128> UrlUtf8;
+							WideToUtf8({(wchar_t*)HttpReq->CookedUrl.pAbsPath,
+										gsl::narrow<size_t>(HttpReq->CookedUrl.AbsPathLength / sizeof(wchar_t))},
+									   UrlUtf8);
+							int				 PrefixLen = Service->UriPrefixLength();
+							std::string_view RelativeUri{UrlUtf8.ToView()};
+							RelativeUri.remove_prefix(std::min(PrefixLen, static_cast<int>(RelativeUri.size())));
+							WsHandler->OnWebSocketOpen(std::move(WsConnRef), RelativeUri);
 							WsConn->Start();
 
 							return nullptr;
@@ -2603,11 +2638,11 @@ InitialRequestHandler::HandleCompletion(ULONG IoResult, ULONG_PTR NumberOfBytesT
 
 						ZEN_WARN("WebSocket 101 send failed: {} ({:#x})", GetSystemErrorAsString(SendResult), SendResult);
 
-						// WebSocket upgrade failed — return nullptr since ServerRequest()
+						// WebSocket upgrade failed - return nullptr since ServerRequest()
 						// was never populated (no InvokeRequestHandler call)
 						return nullptr;
 					}
-					// Service doesn't support WebSocket or missing key — fall through to normal handling
+					// Service doesn't support WebSocket or missing key - fall through to normal handling
 				}
 			}