1 files changed, 311 insertions, 0 deletions
diff --git a/src/zenhttp/servers/wsasio.cpp b/src/zenhttp/servers/wsasio.cpp
new file mode 100644
index 000000000..b2543277a
--- /dev/null
+++ b/src/zenhttp/servers/wsasio.cpp
@@ -0,0 +1,311 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "wsasio.h"
+#include "wsframecodec.h"
+
+#include <zencore/logging.h>
+#include <zenhttp/httpserver.h>
+
+namespace zen::asio_http {
+
+static LoggerRef
+WsLog()
+{
+	static LoggerRef g_Logger = logging::Get("ws");
+	return g_Logger;
+}
+
+//////////////////////////////////////////////////////////////////////////
+
+WsAsioConnection::WsAsioConnection(std::unique_ptr<asio::ip::tcp::socket> Socket, IWebSocketHandler& Handler, HttpServer* Server)
+: m_Socket(std::move(Socket))
+, m_Handler(Handler)
+, m_HttpServer(Server)
+{
+}
+
+WsAsioConnection::~WsAsioConnection()
+{
+	m_IsOpen.store(false);
+	if (m_HttpServer)
+	{
+		m_HttpServer->OnWebSocketConnectionClosed();
+	}
+}
+
+void
+WsAsioConnection::Start()
+{
+	EnqueueRead();
+}
+
+bool
+WsAsioConnection::IsOpen() const
+{
+	return m_IsOpen.load(std::memory_order_relaxed);
+}
+
+//////////////////////////////////////////////////////////////////////////
+//
+// Read loop
+//
+
+void
+WsAsioConnection::EnqueueRead()
+{
+	if (!m_IsOpen.load(std::memory_order_relaxed))
+	{
+		return;
+	}
+
+	Ref<WsAsioConnection> Self(this);
+
+	asio::async_read(*m_Socket, m_ReadBuffer, asio::transfer_at_least(1), [Self](const asio::error_code& Ec, std::size_t ByteCount) {
+		Self->OnDataReceived(Ec, ByteCount);
+	});
+}
+
+void
+WsAsioConnection::OnDataReceived(const asio::error_code& Ec, [[maybe_unused]] std::size_t ByteCount)
+{
+	if (Ec)
+	{
+		if (Ec != asio::error::eof && Ec != asio::error::operation_aborted)
+		{
+			ZEN_LOG_DEBUG(WsLog(), "WebSocket read error: {}", Ec.message());
+		}
+
+		if (m_IsOpen.exchange(false))
+		{
+			m_Handler.OnWebSocketClose(*this, 1006, "connection lost");
+		}
+		return;
+	}
+
+	ProcessReceivedData();
+
+	if (m_IsOpen.load(std::memory_order_relaxed))
+	{
+		EnqueueRead();
+	}
+}
+
+void
+WsAsioConnection::ProcessReceivedData()
+{
+	while (m_ReadBuffer.size() > 0)
+	{
+		const auto& InputBuffer = m_ReadBuffer.data();
+		const auto* Data		= static_cast<const uint8_t*>(InputBuffer.data());
+		const auto	Size		= InputBuffer.size();
+
+		WsFrameParseResult Frame = WsFrameCodec::TryParseFrame(Data, Size);
+		if (!Frame.IsValid)
+		{
+			break;	// not enough data yet
+		}
+
+		m_ReadBuffer.consume(Frame.BytesConsumed);
+
+		if (m_HttpServer)
+		{
+			m_HttpServer->OnWebSocketFrameReceived(Frame.BytesConsumed);
+		}
+
+		switch (Frame.Opcode)
+		{
+			case WebSocketOpcode::kText:
+			case WebSocketOpcode::kBinary:
+				{
+					WebSocketMessage Msg;
+					Msg.Opcode	= Frame.Opcode;
+					Msg.Payload = IoBuffer(IoBuffer::Clone, Frame.Payload.data(), Frame.Payload.size());
+					m_Handler.OnWebSocketMessage(*this, Msg);
+					break;
+				}
+
+			case WebSocketOpcode::kPing:
+				{
+					// Auto-respond with pong carrying the same payload
+					std::vector<uint8_t> PongFrame = WsFrameCodec::BuildFrame(WebSocketOpcode::kPong, Frame.Payload);
+					EnqueueWrite(std::move(PongFrame));
+					break;
+				}
+
+			case WebSocketOpcode::kPong:
+				// Unsolicited pong — ignore per RFC 6455
+				break;
+
+			case WebSocketOpcode::kClose:
+				{
+					uint16_t		 Code = 1000;
+					std::string_view Reason;
+
+					if (Frame.Payload.size() >= 2)
+					{
+						Code = (uint16_t(Frame.Payload[0]) << 8) | uint16_t(Frame.Payload[1]);
+						if (Frame.Payload.size() > 2)
+						{
+							Reason = std::string_view(reinterpret_cast<const char*>(Frame.Payload.data() + 2), Frame.Payload.size() - 2);
+						}
+					}
+
+					// Echo close frame back if we haven't sent one yet
+					if (!m_CloseSent.exchange(true))
+					{
+						std::vector<uint8_t> CloseFrame = WsFrameCodec::BuildCloseFrame(Code);
+						EnqueueWrite(std::move(CloseFrame));
+					}
+
+					m_IsOpen.store(false);
+					m_Handler.OnWebSocketClose(*this, Code, Reason);
+
+					// Shut down the socket
+					std::error_code ShutdownEc;
+					m_Socket->shutdown(asio::socket_base::shutdown_both, ShutdownEc);
+					m_Socket->close(ShutdownEc);
+					return;
+				}
+
+			default:
+				ZEN_LOG_WARN(WsLog(), "Unknown WebSocket opcode: {:#x}", static_cast<uint8_t>(Frame.Opcode));
+				break;
+		}
+	}
+}
+
+//////////////////////////////////////////////////////////////////////////
+//
+// Write queue
+//
+
+void
+WsAsioConnection::SendText(std::string_view Text)
+{
+	if (!m_IsOpen.load(std::memory_order_relaxed))
+	{
+		return;
+	}
+
+	std::span<const uint8_t> Payload(reinterpret_cast<const uint8_t*>(Text.data()), Text.size());
+	std::vector<uint8_t>	 Frame = WsFrameCodec::BuildFrame(WebSocketOpcode::kText, Payload);
+	EnqueueWrite(std::move(Frame));
+}
+
+void
+WsAsioConnection::SendBinary(std::span<const uint8_t> Data)
+{
+	if (!m_IsOpen.load(std::memory_order_relaxed))
+	{
+		return;
+	}
+
+	std::vector<uint8_t> Frame = WsFrameCodec::BuildFrame(WebSocketOpcode::kBinary, Data);
+	EnqueueWrite(std::move(Frame));
+}
+
+void
+WsAsioConnection::Close(uint16_t Code, std::string_view Reason)
+{
+	DoClose(Code, Reason);
+}
+
+void
+WsAsioConnection::DoClose(uint16_t Code, std::string_view Reason)
+{
+	if (!m_IsOpen.exchange(false))
+	{
+		return;
+	}
+
+	if (!m_CloseSent.exchange(true))
+	{
+		std::vector<uint8_t> CloseFrame = WsFrameCodec::BuildCloseFrame(Code, Reason);
+		EnqueueWrite(std::move(CloseFrame));
+	}
+
+	m_Handler.OnWebSocketClose(*this, Code, Reason);
+}
+
+void
+WsAsioConnection::EnqueueWrite(std::vector<uint8_t> Frame)
+{
+	if (m_HttpServer)
+	{
+		m_HttpServer->OnWebSocketFrameSent(Frame.size());
+	}
+
+	bool ShouldFlush = false;
+
+	m_WriteLock.WithExclusiveLock([&] {
+		m_WriteQueue.push_back(std::move(Frame));
+		if (!m_IsWriting)
+		{
+			m_IsWriting = true;
+			ShouldFlush = true;
+		}
+	});
+
+	if (ShouldFlush)
+	{
+		FlushWriteQueue();
+	}
+}
+
+void
+WsAsioConnection::FlushWriteQueue()
+{
+	std::vector<uint8_t> Frame;
+
+	m_WriteLock.WithExclusiveLock([&] {
+		if (m_WriteQueue.empty())
+		{
+			m_IsWriting = false;
+			return;
+		}
+		Frame = std::move(m_WriteQueue.front());
+		m_WriteQueue.pop_front();
+	});
+
+	if (Frame.empty())
+	{
+		return;
+	}
+
+	Ref<WsAsioConnection> Self(this);
+
+	// Move Frame into a shared_ptr so we can create the buffer and capture ownership
+	// in the same async_write call without evaluation order issues.
+	auto OwnedFrame = std::make_shared<std::vector<uint8_t>>(std::move(Frame));
+
+	asio::async_write(*m_Socket,
+					  asio::buffer(OwnedFrame->data(), OwnedFrame->size()),
+					  [Self, OwnedFrame](const asio::error_code& Ec, std::size_t ByteCount) { Self->OnWriteComplete(Ec, ByteCount); });
+}
+
+void
+WsAsioConnection::OnWriteComplete(const asio::error_code& Ec, [[maybe_unused]] std::size_t ByteCount)
+{
+	if (Ec)
+	{
+		if (Ec != asio::error::operation_aborted)
+		{
+			ZEN_LOG_DEBUG(WsLog(), "WebSocket write error: {}", Ec.message());
+		}
+
+		m_WriteLock.WithExclusiveLock([&] {
+			m_IsWriting = false;
+			m_WriteQueue.clear();
+		});
+
+		if (m_IsOpen.exchange(false))
+		{
+			m_Handler.OnWebSocketClose(*this, 1006, "write error");
+		}
+		return;
+	}
+
+	FlushWriteQueue();
+}
+
+}  // namespace zen::asio_http