1 files changed, 35 insertions, 7 deletions

diff --git a/src/zenhttp/servers/httpasio.cpp b/src/zenhttp/servers/httpasio.cpp
index 7972777b8..b624c3a29 100644
--- a/src/zenhttp/servers/httpasio.cpp
+++ b/src/zenhttp/servers/httpasio.cpp
@@ -625,6 +625,8 @@ public:
 
 	void SetAllowZeroCopyFileSend(bool Allow) { m_AllowZeroCopyFileSend = Allow; }
 	void SetKeepAlive(bool KeepAlive) { m_IsKeepAlive = KeepAlive; }
+	void SetContentTypeOverride(std::string Override) { m_ContentTypeOverride = std::move(Override); }
+	void SetContentRangeHeader(std::string V) { m_ContentRangeHeader = std::move(V); }
 
 	/**
 	 * Initialize the response for sending a payload made up of multiple blobs
@@ -768,10 +770,18 @@ public:
 	{
 		ZEN_MEMSCOPE(GetHttpasioTag());
 
+		std::string_view ContentTypeStr =
+			m_ContentTypeOverride.empty() ? MapContentTypeToString(m_ContentType) : std::string_view(m_ContentTypeOverride);
+
 		m_Headers << "HTTP/1.1 " << ResponseCode() << " " << ReasonStringForHttpResultCode(ResponseCode()) << "\r\n"
-				  << "Content-Type: " << MapContentTypeToString(m_ContentType) << "\r\n"
+				  << "Content-Type: " << ContentTypeStr << "\r\n"
 				  << "Content-Length: " << ContentLength() << "\r\n"sv;
 
+		if (!m_ContentRangeHeader.empty())
+		{
+			m_Headers << "Content-Range: " << m_ContentRangeHeader << "\r\n"sv;
+		}
+
 		if (!m_IsKeepAlive)
 		{
 			m_Headers << "Connection: close\r\n"sv;
@@ -898,7 +908,9 @@ private:
 	bool							 m_AllowZeroCopyFileSend = true;
 	State							 m_State				 = State::kUninitialized;
 	HttpContentType					 m_ContentType			 = HttpContentType::kBinary;
-	uint64_t						 m_ContentLength		 = 0;
+	std::string						 m_ContentTypeOverride;
+	std::string						 m_ContentRangeHeader;
+	uint64_t						 m_ContentLength = 0;
 	eastl::fixed_vector<IoBuffer, 8> m_DataBuffers;	 // This is here to keep the IoBuffer buffers/handles alive
 	ExtendableStringBuilder<160>	 m_Headers;
 
@@ -1275,7 +1287,9 @@ HttpServerConnectionT<SocketType>::HandleRequest()
 					asio::buffer(ResponseStr->data(), ResponseStr->size()),
 					asio::bind_executor(
 						m_Strand,
-						[Conn = AsSharedPtr(), WsHandler, OwnedResponse = ResponseStr](const asio::error_code& Ec, std::size_t) {
+						[Conn = AsSharedPtr(), WsHandler, OwnedResponse = ResponseStr, PrefixLen = Service->UriPrefixLength()](
+							const asio::error_code& Ec,
+							std::size_t) {
 							if (Ec)
 							{
 								ZEN_WARN("WebSocket 101 send failed: {}", Ec.message());
@@ -1287,7 +1301,9 @@ HttpServerConnectionT<SocketType>::HandleRequest()
 							Ref<WsConnType> WsConn(new WsConnType(std::move(Conn->m_Socket), *WsHandler, Conn->m_Server.m_HttpServer));
 							Ref<WebSocketConnection> WsConnRef(WsConn.Get());
 
-							WsHandler->OnWebSocketOpen(std::move(WsConnRef));
+							std::string_view FullUrl	 = Conn->m_RequestData.Url();
+							std::string_view RelativeUri = FullUrl.substr(std::min(PrefixLen, static_cast<int>(FullUrl.size())));
+							WsHandler->OnWebSocketOpen(std::move(WsConnRef), RelativeUri);
 							WsConn->Start();
 						}));
 
@@ -1295,7 +1311,7 @@ HttpServerConnectionT<SocketType>::HandleRequest()
 				return;
 			}
 		}
-		// 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
 	}
 
 	if (!m_RequestData.IsKeepAlive())
@@ -2127,6 +2143,10 @@ HttpAsioServerRequest::WriteResponse(HttpResponseCode ResponseCode)
 	m_Response.reset(new HttpResponse(HttpContentType::kBinary, m_RequestNumber));
 	m_Response->SetAllowZeroCopyFileSend(m_AllowZeroCopyFileSend);
 	m_Response->SetKeepAlive(m_Request.IsKeepAlive());
+	if (!m_ContentRangeHeader.empty())
+	{
+		m_Response->SetContentRangeHeader(std::move(m_ContentRangeHeader));
+	}
 	std::array<IoBuffer, 0> Empty;
 
 	m_Response->InitializeForPayload((uint16_t)ResponseCode, Empty);
@@ -2142,6 +2162,14 @@ HttpAsioServerRequest::WriteResponse(HttpResponseCode ResponseCode, HttpContentT
 	m_Response.reset(new HttpResponse(ContentType, m_RequestNumber));
 	m_Response->SetAllowZeroCopyFileSend(m_AllowZeroCopyFileSend);
 	m_Response->SetKeepAlive(m_Request.IsKeepAlive());
+	if (!m_ContentTypeOverride.empty())
+	{
+		m_Response->SetContentTypeOverride(std::move(m_ContentTypeOverride));
+	}
+	if (!m_ContentRangeHeader.empty())
+	{
+		m_Response->SetContentRangeHeader(std::move(m_ContentRangeHeader));
+	}
 	m_Response->InitializeForPayload((uint16_t)ResponseCode, Blobs);
 }
 
@@ -2590,7 +2618,7 @@ HttpAsioServer::OnRun(bool IsInteractive)
 		}
 
 		ShutdownRequested = m_ShutdownEvent.Wait(WaitTimeout);
-	} while (!ShutdownRequested);
+	} while (!ShutdownRequested && !IsApplicationExitRequested());
 #else
 	if (IsInteractive)
 	{
@@ -2600,7 +2628,7 @@ HttpAsioServer::OnRun(bool IsInteractive)
 	do
 	{
 		ShutdownRequested = m_ShutdownEvent.Wait(WaitTimeout);
-	} while (!ShutdownRequested);
+	} while (!ShutdownRequested && !IsApplicationExitRequested());
 #endif
 }