path: root/src/zenhttp/servers/wsasio.cpp

// Copyright Epic Games, Inc. All Rights Reserved.

#include "wsasio.h"
#include "asio_socket_traits.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;
}

//////////////////////////////////////////////////////////////////////////

template<typename SocketType>
WsAsioConnectionT<SocketType>::WsAsioConnectionT(std::unique_ptr<SocketType> Socket, IWebSocketHandler& Handler, HttpServer* Server)
: m_Socket(std::move(Socket))
, m_Handler(Handler)
, m_HttpServer(Server)
{
}

template<typename SocketType>
WsAsioConnectionT<SocketType>::~WsAsioConnectionT()
{
	m_IsOpen.store(false);
	if (m_HttpServer)
	{
		m_HttpServer->OnWebSocketConnectionClosed();
	}
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::Start()
{
	EnqueueRead();
}

template<typename SocketType>
bool
WsAsioConnectionT<SocketType>::IsOpen() const
{
	return m_IsOpen.load(std::memory_order_relaxed);
}

//////////////////////////////////////////////////////////////////////////
//
// Read loop
//

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::EnqueueRead()
{
	if (!m_IsOpen.load(std::memory_order_relaxed))
	{
		return;
	}

	Ref<WsAsioConnectionT> 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);
	});
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::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();
	}
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::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;
					SocketTraits<SocketType>::ShutdownBoth(*m_Socket, ShutdownEc);
					SocketTraits<SocketType>::Close(*m_Socket, ShutdownEc);
					return;
				}

			default:
				ZEN_LOG_WARN(WsLog(), "Unknown WebSocket opcode: {:#x}", static_cast<uint8_t>(Frame.Opcode));
				break;
		}
	}
}

//////////////////////////////////////////////////////////////////////////
//
// Write queue
//

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::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));
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::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));
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::Close(uint16_t Code, std::string_view Reason)
{
	DoClose(Code, Reason);
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::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);
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::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();
	}
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::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<WsAsioConnectionT> 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); });
}

template<typename SocketType>
void
WsAsioConnectionT<SocketType>::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();
}

//////////////////////////////////////////////////////////////////////////
// Explicit template instantiations

template class WsAsioConnectionT<asio::ip::tcp::socket>;

#if defined(ASIO_HAS_LOCAL_SOCKETS)
template class WsAsioConnectionT<asio::local::stream_protocol::socket>;
#endif

#if ZEN_USE_OPENSSL
template class WsAsioConnectionT<asio::ssl::stream<asio::ip::tcp::socket>>;
#endif

}  // namespace zen::asio_http