1 files changed, 229 insertions, 0 deletions
diff --git a/src/zenhttp/servers/wsframecodec.cpp b/src/zenhttp/servers/wsframecodec.cpp
new file mode 100644
index 000000000..a4c5e0f16
--- /dev/null
+++ b/src/zenhttp/servers/wsframecodec.cpp
@@ -0,0 +1,229 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "wsframecodec.h"
+
+#include <zencore/base64.h>
+#include <zencore/sha1.h>
+
+#include <cstring>
+#include <random>
+
+namespace zen {
+
+//////////////////////////////////////////////////////////////////////////
+//
+// Frame parsing
+//
+
+WsFrameParseResult
+WsFrameCodec::TryParseFrame(const uint8_t* Data, size_t Size)
+{
+	// Minimum frame: 2 bytes header (unmasked server frames) or 6 bytes (masked client frames)
+	if (Size < 2)
+	{
+		return {};
+	}
+
+	const bool	  Fin		 = (Data[0] & 0x80) != 0;
+	const uint8_t OpcodeRaw	 = Data[0] & 0x0F;
+	const bool	  Masked	 = (Data[1] & 0x80) != 0;
+	uint64_t	  PayloadLen = Data[1] & 0x7F;
+
+	size_t HeaderSize = 2;
+
+	if (PayloadLen == 126)
+	{
+		if (Size < 4)
+		{
+			return {};
+		}
+		PayloadLen = (uint64_t(Data[2]) << 8) | uint64_t(Data[3]);
+		HeaderSize = 4;
+	}
+	else if (PayloadLen == 127)
+	{
+		if (Size < 10)
+		{
+			return {};
+		}
+		PayloadLen = (uint64_t(Data[2]) << 56) | (uint64_t(Data[3]) << 48) | (uint64_t(Data[4]) << 40) | (uint64_t(Data[5]) << 32) |
+					 (uint64_t(Data[6]) << 24) | (uint64_t(Data[7]) << 16) | (uint64_t(Data[8]) << 8) | uint64_t(Data[9]);
+		HeaderSize = 10;
+	}
+
+	const size_t MaskSize	= Masked ? 4 : 0;
+	const size_t TotalFrame = HeaderSize + MaskSize + PayloadLen;
+
+	if (Size < TotalFrame)
+	{
+		return {};
+	}
+
+	const uint8_t* MaskKey	   = Masked ? (Data + HeaderSize) : nullptr;
+	const uint8_t* PayloadData = Data + HeaderSize + MaskSize;
+
+	WsFrameParseResult Result;
+	Result.IsValid		 = true;
+	Result.BytesConsumed = TotalFrame;
+	Result.Opcode		 = static_cast<WebSocketOpcode>(OpcodeRaw);
+	Result.Fin			 = Fin;
+
+	Result.Payload.resize(static_cast<size_t>(PayloadLen));
+	if (PayloadLen > 0)
+	{
+		std::memcpy(Result.Payload.data(), PayloadData, static_cast<size_t>(PayloadLen));
+
+		if (Masked)
+		{
+			for (size_t i = 0; i < Result.Payload.size(); ++i)
+			{
+				Result.Payload[i] ^= MaskKey[i & 3];
+			}
+		}
+	}
+
+	return Result;
+}
+
+//////////////////////////////////////////////////////////////////////////
+//
+// Frame building (server-to-client, no masking)
+//
+
+std::vector<uint8_t>
+WsFrameCodec::BuildFrame(WebSocketOpcode Opcode, std::span<const uint8_t> Payload)
+{
+	std::vector<uint8_t> Frame;
+
+	const size_t PayloadLen = Payload.size();
+
+	// FIN + opcode
+	Frame.push_back(0x80 | static_cast<uint8_t>(Opcode));
+
+	// Payload length (no mask bit for server frames)
+	if (PayloadLen < 126)
+	{
+		Frame.push_back(static_cast<uint8_t>(PayloadLen));
+	}
+	else if (PayloadLen <= 0xFFFF)
+	{
+		Frame.push_back(126);
+		Frame.push_back(static_cast<uint8_t>((PayloadLen >> 8) & 0xFF));
+		Frame.push_back(static_cast<uint8_t>(PayloadLen & 0xFF));
+	}
+	else
+	{
+		Frame.push_back(127);
+		for (int i = 7; i >= 0; --i)
+		{
+			Frame.push_back(static_cast<uint8_t>((PayloadLen >> (i * 8)) & 0xFF));
+		}
+	}
+
+	Frame.insert(Frame.end(), Payload.begin(), Payload.end());
+
+	return Frame;
+}
+
+std::vector<uint8_t>
+WsFrameCodec::BuildCloseFrame(uint16_t Code, std::string_view Reason)
+{
+	std::vector<uint8_t> Payload;
+	Payload.push_back(static_cast<uint8_t>((Code >> 8) & 0xFF));
+	Payload.push_back(static_cast<uint8_t>(Code & 0xFF));
+	Payload.insert(Payload.end(), Reason.begin(), Reason.end());
+
+	return BuildFrame(WebSocketOpcode::kClose, Payload);
+}
+
+//////////////////////////////////////////////////////////////////////////
+//
+// Frame building (client-to-server, with masking)
+//
+
+std::vector<uint8_t>
+WsFrameCodec::BuildMaskedFrame(WebSocketOpcode Opcode, std::span<const uint8_t> Payload)
+{
+	std::vector<uint8_t> Frame;
+
+	const size_t PayloadLen = Payload.size();
+
+	// FIN + opcode
+	Frame.push_back(0x80 | static_cast<uint8_t>(Opcode));
+
+	// Payload length with mask bit set
+	if (PayloadLen < 126)
+	{
+		Frame.push_back(0x80 | static_cast<uint8_t>(PayloadLen));
+	}
+	else if (PayloadLen <= 0xFFFF)
+	{
+		Frame.push_back(0x80 | 126);
+		Frame.push_back(static_cast<uint8_t>((PayloadLen >> 8) & 0xFF));
+		Frame.push_back(static_cast<uint8_t>(PayloadLen & 0xFF));
+	}
+	else
+	{
+		Frame.push_back(0x80 | 127);
+		for (int i = 7; i >= 0; --i)
+		{
+			Frame.push_back(static_cast<uint8_t>((PayloadLen >> (i * 8)) & 0xFF));
+		}
+	}
+
+	// Generate random 4-byte mask key
+	static thread_local std::mt19937 s_Rng(std::random_device{}());
+	uint32_t						 MaskValue = s_Rng();
+	uint8_t							 MaskKey[4];
+	std::memcpy(MaskKey, &MaskValue, 4);
+
+	Frame.insert(Frame.end(), MaskKey, MaskKey + 4);
+
+	// Masked payload
+	for (size_t i = 0; i < PayloadLen; ++i)
+	{
+		Frame.push_back(Payload[i] ^ MaskKey[i & 3]);
+	}
+
+	return Frame;
+}
+
+std::vector<uint8_t>
+WsFrameCodec::BuildMaskedCloseFrame(uint16_t Code, std::string_view Reason)
+{
+	std::vector<uint8_t> Payload;
+	Payload.push_back(static_cast<uint8_t>((Code >> 8) & 0xFF));
+	Payload.push_back(static_cast<uint8_t>(Code & 0xFF));
+	Payload.insert(Payload.end(), Reason.begin(), Reason.end());
+
+	return BuildMaskedFrame(WebSocketOpcode::kClose, Payload);
+}
+
+//////////////////////////////////////////////////////////////////////////
+//
+// Sec-WebSocket-Accept key computation (RFC 6455 section 4.2.2)
+//
+
+static constexpr std::string_view kWebSocketMagicGuid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
+
+std::string
+WsFrameCodec::ComputeAcceptKey(std::string_view ClientKey)
+{
+	// Concatenate client key with the magic GUID
+	std::string Combined;
+	Combined.reserve(ClientKey.size() + kWebSocketMagicGuid.size());
+	Combined.append(ClientKey);
+	Combined.append(kWebSocketMagicGuid);
+
+	// SHA1 hash
+	SHA1 Hash = SHA1::HashMemory(Combined.data(), Combined.size());
+
+	// Base64 encode the 20-byte hash
+	char	 Base64Buf[Base64::GetEncodedDataSize(20) + 1];
+	uint32_t EncodedLen	  = Base64::Encode(Hash.Hash, 20, Base64Buf);
+	Base64Buf[EncodedLen] = '\0';
+
+	return std::string(Base64Buf, EncodedLen);
+}
+
+}  // namespace zen