#include "pch.h"
#include "System/Net/WebSockets/WebSocket.h"
#include "System/Net/WebSockets/WebSocketException.h"
#include "System/ArgumentException.h"
#include <openssl/sha.h>
#include <openssl/evp.h>
#include <array>
#include <vector>
#include <string>
#include <cstdint>
namespace DotNetDupe {
namespace System {
namespace Net {
namespace WebSockets {
static void BuildWebSocketFrameHeader(uint8_t opcode, size_t len, std::vector<uint8_t>& frame) {
frame.push_back(opcode);
if (len <= 125) {
frame.push_back(static_cast<uint8_t>(len));
} else if (len <= 65535) {
frame.push_back(126);
frame.push_back(static_cast<uint8_t>((len >> 8) & 0xFF));
frame.push_back(static_cast<uint8_t>(len & 0xFF));
} else {
frame.push_back(127);
for (int i = 7; i >= 0; --i) frame.push_back(static_cast<uint8_t>((len >> (i * 8)) & 0xFF));
}
}
class WebSocket::Impl : public Object {
public:
SmartPointer<Sockets::NetworkStream> m_pStream;
WebSocketState m_eState;
Threading::CriticalSection m_csLock;
Impl(SmartPointer<Sockets::NetworkStream> pStream)
: m_pStream(pStream), m_eState(WebSocketState::Open) {}
void CloseState() {
Threading::Lock<Threading::CriticalSection> lock(m_csLock);
m_eState = WebSocketState::Closed;
}
bool ReadExtended16(uint64_t& payloadLen) {
uint8_t extLen[2] = { 0 };
int bytesRead = m_pStream->Read(reinterpret_cast<char*>(extLen), 0, 2);
if (bytesRead <= 0) {
throw WebSocketException(WebSocketError::ConnectionClosedPrematurely, "Unexpected EOF reading 16-bit extended length.");
}
payloadLen = (static_cast<uint64_t>(extLen[0]) << 8) | extLen[1];
return true;
}
bool ReadExtended64(uint64_t& payloadLen) {
uint8_t extLen[8] = { 0 };
int bytesRead = m_pStream->Read(reinterpret_cast<char*>(extLen), 0, 8);
if (bytesRead <= 0) {
throw WebSocketException(WebSocketError::ConnectionClosedPrematurely, "Unexpected EOF reading 64-bit extended length.");
}
payloadLen = 0;
for (int i = 0; i < 8; ++i) payloadLen = (payloadLen << 8) | extLen[i];
return true;
}
bool ReadFrameHeader(uint8_t& opcode, bool& masked, uint64_t& payloadLen) {
uint8_t header[2] = { 0 };
int bytesRead = 0;
try {
bytesRead = m_pStream->Read(reinterpret_cast<char*>(header), 0, 2);
} catch (const Exception& ex) {
throw WebSocketException(WebSocketError::NativeError, ex);
}
if (bytesRead <= 0) return false;
opcode = header[0] & 0x0F;
masked = (header[1] & 0x80) != 0;
payloadLen = header[1] & 0x7F;
return true;
}
bool ReadExtendedLength(uint64_t& payloadLen) {
if (payloadLen == 126) return ReadExtended16(payloadLen);
if (payloadLen == 127) return ReadExtended64(payloadLen);
return true;
}
void ReadPayloadBytes(uint64_t payloadLen, std::vector<uint8_t>& payload) {
payload.resize(payloadLen, 0);
uint64_t totalRead = 0;
while (totalRead < payloadLen) {
int toRead = static_cast<int>(payloadLen - totalRead);
int bytesRead = m_pStream->Read(reinterpret_cast<char*>(payload.data() + totalRead), 0, toRead);
if (bytesRead <= 0) {
throw WebSocketException(WebSocketError::ConnectionClosedPrematurely, "Connection lost while reading payload bytes.");
}
totalRead += bytesRead;
}
}
void UnmaskPayload(const std::array<uint8_t, 4>& maskKey, uint64_t payloadLen, std::vector<uint8_t>& payload) {
for (uint64_t i = 0; i < payloadLen; ++i) {
payload[i] ^= maskKey[i % 4];
}
}
bool ReadMaskKeyAndPayload(bool masked, uint64_t payloadLen, std::vector<uint8_t>& payload) {
std::array<uint8_t, 4> maskKey = { 0 };
if (masked && m_pStream->Read(reinterpret_cast<char*>(maskKey.data()), 0, 4) <= 0) {
throw WebSocketException(WebSocketError::ConnectionClosedPrematurely, "Connection lost while reading mask key.");
}
ReadPayloadBytes(payloadLen, payload);
if (masked) UnmaskPayload(maskKey, payloadLen, payload);
return true;
}
bool WriteFrame(uint8_t opcode, const uint8_t* pData, size_t len) {
Threading::Lock<Threading::CriticalSection> lock(m_csLock);
if (m_pStream.IsNull() || m_eState != WebSocketState::Open) {
throw WebSocketException(WebSocketError::InvalidState, "WebSocket is not connected or already closed.");
}
std::vector<uint8_t> frame;
BuildWebSocketFrameHeader(opcode, len, frame);
if (pData != nullptr && len > 0) frame.insert(frame.end(), pData, pData + len);
try {
m_pStream->Write(reinterpret_cast<const char*>(frame.data()), 0, static_cast<int>(frame.size()));
} catch (const Exception& ex) {
throw WebSocketException(WebSocketError::NativeError, ex);
}
return true;
}
};
WebSocket::WebSocket(SmartPointer<Sockets::NetworkStream> pStream)
: m_pImpl(SmartPointer<Impl>::NewShared(pStream)) {
}
WebSocket::~WebSocket() {
}
WebSocketState WebSocket::GetState() const { return m_pImpl->m_eState; }
void WebSocket::SetState(WebSocketState state) {
m_pImpl->m_eState = state;
}
String WebSocket::ComputeSecWebSocketAccept(const String& secWebSocketKey) {
if (secWebSocketKey.IsEmpty()) {
throw ArgumentException("secWebSocketKey cannot be empty.");
}
std::string key = secWebSocketKey.GetRawString();
std::string magic = key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
unsigned char hash[SHA_DIGEST_LENGTH];
SHA1(reinterpret_cast<const unsigned char*>(magic.c_str()), magic.length(), hash);
char encoded[128] = { 0 };
EVP_EncodeBlock(reinterpret_cast<unsigned char*>(encoded), hash, SHA_DIGEST_LENGTH);
return String(encoded);
}
bool WebSocket::SendAsync(const String& message) {
std::string text = message.GetRawString();
return m_pImpl->WriteFrame(0x81, reinterpret_cast<const uint8_t*>(text.data()), text.length());
}
bool WebSocket::SendBytes(const Array<uint8_t>& data) {
return m_pImpl->WriteFrame(0x82, data.GetData(), data.GetLength());
}
bool WebSocket::ReceiveText(String& outMessage) {
{
Threading::Lock<Threading::CriticalSection> lock(m_pImpl->m_csLock);
if (m_pImpl->m_pStream.IsNull() || m_pImpl->m_eState != WebSocketState::Open) {
throw WebSocketException(WebSocketError::InvalidState, "WebSocket is not connected or already closed.");
}
}
uint8_t opcode = 0; bool masked = false; uint64_t payloadLen = 0;
if (!m_pImpl->ReadFrameHeader(opcode, masked, payloadLen) || opcode == 0x08) {
m_pImpl->CloseState();
return false;
}
std::vector<uint8_t> payload;
m_pImpl->ReadExtendedLength(payloadLen);
m_pImpl->ReadMaskKeyAndPayload(masked, payloadLen, payload);
outMessage = String(std::string(payload.begin(), payload.end()).c_str());
return true;
}
void WebSocket::Close() {
Threading::Lock<Threading::CriticalSection> lock(m_pImpl->m_csLock);
if (m_pImpl->m_eState == WebSocketState::Open) {
uint8_t closeFrame[2] = { 0x88, 0x00 };
if (!m_pImpl->m_pStream.IsNull()) {
try {
m_pImpl->m_pStream->Write(reinterpret_cast<const char*>(closeFrame), 0, 2);
} catch (...) {
// Suppress stream write errors during close handshake
(void)0;
}
}
m_pImpl->m_eState = WebSocketState::Closed;
}
}
}
}
}
}