Files
MARTe-Integrated-Components/Source/Applications/StreamHub/WSServer.cpp
T
2026-07-01 16:39:34 +02:00

499 lines
18 KiB
C++

/**
* @file WSServer.cpp
* @brief Multi-client WebSocket server implementation.
*/
#include "WSServer.h"
#include "WSFrame.h"
#include "SHA1.h"
#include "Base64.h"
#include "AdvancedErrorManagement.h"
#include "Sleep.h"
#include "Threads.h"
#include "TimeoutType.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
namespace StreamHub {
using MARTe::TimeoutType;
using MARTe::Sleep;
/*---------------------------------------------------------------------------*/
/* Thread trampolines */
/*---------------------------------------------------------------------------*/
struct AcceptThreadArg { WSServer *srv; };
struct ClientThreadArg { WSServer *srv; uint32 slot; };
static void AcceptThreadEntry(const void *arg) {
WSServer *srv = reinterpret_cast<AcceptThreadArg *>(
const_cast<void *>(arg))->srv;
delete reinterpret_cast<const AcceptThreadArg *>(arg);
srv->AcceptLoop();
}
static void ClientReadEntry(const void *arg) {
const ClientThreadArg *a = reinterpret_cast<const ClientThreadArg *>(arg);
WSServer *srv = a->srv;
uint32 slot = a->slot;
delete a;
srv->ClientReadLoop(slot);
}
/*---------------------------------------------------------------------------*/
/* Tiny JSON helpers (no library, only what we need) */
/*---------------------------------------------------------------------------*/
/** Locate the first occurrence of pattern in s, return pointer or NULL. */
static const char *FindSubstr(const char *s, const char *pattern) {
return strstr(s, pattern);
}
/*---------------------------------------------------------------------------*/
/* WSServer */
/*---------------------------------------------------------------------------*/
WSServer::WSServer()
: numClients(0u),
callback(static_cast<WSCommandCallback *>(0)),
running(false),
acceptTid(MARTe::InvalidThreadIdentifier) {
for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
clients[i].sock = static_cast<BasicTCPSocket *>(0);
clients[i].active = false;
clients[i].readTid = MARTe::InvalidThreadIdentifier;
}
}
WSServer::~WSServer() {
(void) Stop();
}
bool WSServer::Start(uint16 port, WSCommandCallback *cb) {
if (running) { return true; }
callback = cb;
running = true;
if (!tcpListener.Open()) {
REPORT_ERROR_STATIC(MARTe::ErrorManagement::FatalError,
"WSServer: Failed to open TCP listener socket.");
running = false;
return false;
}
if (!tcpListener.Listen(port, static_cast<MARTe::int32>(WS_MAX_CLIENTS))) {
REPORT_ERROR_STATIC(MARTe::ErrorManagement::FatalError,
"WSServer: Failed to listen on port %u.", static_cast<uint32>(port));
running = false;
return false;
}
AcceptThreadArg *arg = new AcceptThreadArg();
arg->srv = this;
acceptTid = MARTe::Threads::BeginThread(
reinterpret_cast<MARTe::ThreadFunctionType>(AcceptThreadEntry),
arg, MARTe::THREADS_DEFAULT_STACKSIZE, "WSAccept");
REPORT_ERROR_STATIC(MARTe::ErrorManagement::Information,
"WSServer: Listening on port %u.", static_cast<uint32>(port));
return true;
}
bool WSServer::Stop() {
if (!running) { return true; }
running = false;
Sleep::MSec(200u);
/* Close all client connections — their read threads will exit on error */
(void) clientsMutex.FastLock();
for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
if (clients[i].active && (clients[i].sock != static_cast<BasicTCPSocket *>(0))) {
clients[i].sock->Close();
}
}
clientsMutex.FastUnLock();
Sleep::MSec(200u);
tcpListener.Close();
Sleep::MSec(100u);
/* Free any remaining slots */
(void) clientsMutex.FastLock();
for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
if (clients[i].sock != static_cast<BasicTCPSocket *>(0)) {
delete clients[i].sock;
clients[i].sock = static_cast<BasicTCPSocket *>(0);
clients[i].active = false;
}
}
numClients = 0u;
clientsMutex.FastUnLock();
return true;
}
/*---------------------------------------------------------------------------*/
/* Accept loop */
/*---------------------------------------------------------------------------*/
void WSServer::AcceptLoop() {
while (running) {
BasicTCPSocket *client = tcpListener.WaitConnection(
TimeoutType(500u), static_cast<BasicTCPSocket *>(0));
if (client == static_cast<BasicTCPSocket *>(0)) {
/* Timeout — loop again to check running flag */
continue;
}
/* HTTP upgrade */
if (!UpgradeHTTP(client)) {
client->Close();
delete client;
continue;
}
uint32 slot = AllocSlot(client);
if (slot == WS_MAX_CLIENTS) {
/* No room */
static const char *full = "HTTP/1.1 503 Service Unavailable\r\n\r\n";
uint32 l = static_cast<uint32>(strlen(full));
(void) client->Write(full, l);
client->Close();
delete client;
continue;
}
/* Notify StreamHub */
if (callback != static_cast<WSCommandCallback *>(0)) {
callback->OnWSClientConnected();
}
/* Start per-client read thread */
ClientThreadArg *arg = new ClientThreadArg();
arg->srv = this;
arg->slot = slot;
clients[slot].readTid = MARTe::Threads::BeginThread(
reinterpret_cast<MARTe::ThreadFunctionType>(ClientReadEntry),
arg, MARTe::THREADS_DEFAULT_STACKSIZE, "WSClient");
}
}
/*---------------------------------------------------------------------------*/
/* HTTP upgrade */
/*---------------------------------------------------------------------------*/
bool WSServer::UpgradeHTTP(BasicTCPSocket *sock) {
/* Read raw HTTP request until \r\n\r\n */
static const uint32 kMaxHdr = 4096u;
char hdrBuf[4096];
uint32 totalRead = 0u;
while (totalRead < kMaxHdr - 1u) {
uint32 rem = kMaxHdr - 1u - totalRead;
bool ok = sock->Read(hdrBuf + totalRead, rem,
TimeoutType(3000u));
if (!ok || rem == 0u) { return false; }
totalRead += rem;
hdrBuf[totalRead] = '\0';
if (strstr(hdrBuf, "\r\n\r\n") != static_cast<char *>(0)) { break; }
}
/* Origin validation (CSWSH / CSRF defence, RFC 6455 §10.2).
* If an Origin header is present, its host must match the Host header
* (same-origin). Non-browser clients (no Origin) are allowed. */
const char *originHdr = FindSubstr(hdrBuf, "Origin:");
if (originHdr != static_cast<const char *>(0)) {
originHdr += 7; /* skip "Origin:" */
while (*originHdr == ' ') { originHdr++; }
/* Extract the host part of Origin: "scheme://host[:port]" */
char originHost[256];
uint32 ohLen = 0u;
const char *op = originHdr;
/* Skip scheme:// */
const char *schemeEnd = strstr(op, "://");
if (schemeEnd != static_cast<const char *>(0)) { op = schemeEnd + 3; }
while (*op != '\r' && *op != '\n' && *op != '\0' &&
*op != '/' && ohLen < 255u) {
originHost[ohLen++] = *op++;
}
originHost[ohLen] = '\0';
/* Extract Host header value */
const char *hostHdr = FindSubstr(hdrBuf, "Host:");
if (hostHdr != static_cast<const char *>(0)) {
hostHdr += 5; /* skip "Host:" */
while (*hostHdr == ' ') { hostHdr++; }
char hostVal[256];
uint32 hvLen = 0u;
while (*hostHdr != '\r' && *hostHdr != '\n' &&
*hostHdr != '\0' && hvLen < 255u) {
hostVal[hvLen++] = *hostHdr++;
}
hostVal[hvLen] = '\0';
if (strcmp(originHost, hostVal) != 0) {
/* Cross-origin — reject the upgrade */
const char *forbidden =
"HTTP/1.1 403 Forbidden\r\n"
"Content-Type: text/plain\r\n"
"Connection: close\r\n"
"\r\nOrigin not allowed\r\n";
uint32 forbLen = static_cast<uint32>(strlen(forbidden));
(void) sock->Write(forbidden, forbLen);
return false;
}
}
}
/* Find Sec-WebSocket-Key */
const char *keyHdr = FindSubstr(hdrBuf, "Sec-WebSocket-Key:");
if (keyHdr == static_cast<const char *>(0)) { return false; }
keyHdr += 18; /* skip "Sec-WebSocket-Key:" */
while (*keyHdr == ' ') { keyHdr++; }
char key[64];
uint32 keyLen = 0u;
while (keyHdr[keyLen] != '\r' && keyHdr[keyLen] != '\n' &&
keyHdr[keyLen] != '\0' && keyLen < 63u) {
key[keyLen] = keyHdr[keyLen];
keyLen++;
}
key[keyLen] = '\0';
/* Compute Accept = Base64(SHA1(key + GUID)) */
static const char *kGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
char concat[128];
snprintf(concat, sizeof(concat), "%s%s", key, kGUID);
uint8 digest[20];
SHA1(reinterpret_cast<const uint8 *>(concat),
static_cast<uint32>(strlen(concat)), digest);
char accept[64];
Base64Encode(digest, 20u, accept);
/* Send 101 response */
char resp[512];
snprintf(resp, sizeof(resp),
"HTTP/1.1 101 Switching Protocols\r\n"
"Upgrade: websocket\r\n"
"Connection: Upgrade\r\n"
"Sec-WebSocket-Accept: %s\r\n"
"\r\n", accept);
uint32 respLen = static_cast<uint32>(strlen(resp));
return sock->Write(resp, respLen);
}
/*---------------------------------------------------------------------------*/
/* Client read loop */
/*---------------------------------------------------------------------------*/
void WSServer::ClientReadLoop(uint32 slotIdx) {
WSClientSlot &slot = clients[slotIdx];
BasicTCPSocket *sock = slot.sock;
/* Receive buffer: WS_MAX_RECV_PAYLOAD + max header (14: 2 + 8 ext-length +
* 4 mask) + 1 spare byte for in-place NUL-termination of the payload. */
static const uint32 kRecvBuf = WS_MAX_RECV_PAYLOAD + 14u + 1u;
uint8 *buf = new uint8[kRecvBuf];
uint32 filled = 0u;
while (running && slot.active) {
/* Read more bytes (with short timeout so we can check running) */
uint32 want = kRecvBuf - filled;
if (want == 0u) {
/* Buffer full — discard old frame (shouldn't happen with reasonable clients) */
filled = 0u;
continue;
}
bool ok = sock->Read(reinterpret_cast<char *>(buf + filled), want,
TimeoutType(500u));
if (!ok) {
/* Timeout or error — check running and retry */
if (!running) { break; }
if (want == kRecvBuf) {
/* Zero bytes read — connection likely closed */
break;
}
continue;
}
filled += want;
/* Parse as many complete frames as possible */
uint32 consumed = 0u;
while (consumed < filled) {
const uint8 *frameStart = buf + consumed;
uint32 available = filled - consumed;
WSFrameHeader hdr;
if (!WSParseHeader(frameStart, available, hdr)) {
break; /* incomplete header — need more data */
}
if (hdr.payloadLen > static_cast<uint64>(WS_MAX_RECV_PAYLOAD)) {
/* Payload too large — close connection */
goto client_done;
}
uint32 plen = static_cast<uint32>(hdr.payloadLen);
uint32 totalFrameBytes = hdr.headerSize + plen;
if (available < totalFrameBytes) {
break; /* incomplete payload — need more data */
}
/* Unmask payload */
uint8 *payload = const_cast<uint8 *>(frameStart) + hdr.headerSize;
if (hdr.masked) {
WSUnmask(payload, plen, hdr.maskKey);
}
/* Handle opcode */
if (hdr.opcode == WS_OPCODE_CLOSE) {
goto client_done;
} else if (hdr.opcode == WS_OPCODE_PING) {
/* Send pong */
(void) slot.writeMutex.FastLock();
SendFrame(slot, WS_OPCODE_PONG, payload, plen);
slot.writeMutex.FastUnLock();
} else if (hdr.opcode == WS_OPCODE_TEXT) {
if (callback != static_cast<WSCommandCallback *>(0)) {
/* NUL-terminate in place for the callback, but restore the
* byte afterwards: when several client frames coalesce in
* one TCP read, payload[plen] is the first header byte of
* the NEXT frame. */
uint8 savedByte = payload[plen];
payload[plen] = '\0'; /* safe: buf has extra byte */
callback->OnWSCommand(reinterpret_cast<const char *>(payload), plen, slotIdx);
payload[plen] = savedByte;
}
}
/* Binary frames from client are ignored (server only sends binary) */
consumed += totalFrameBytes;
}
/* Compact buffer: move unconsumed bytes to front */
if (consumed > 0u && consumed < filled) {
memmove(buf, buf + consumed, filled - consumed);
}
filled = (consumed <= filled) ? (filled - consumed) : 0u;
}
client_done:
delete[] buf;
if (callback != static_cast<WSCommandCallback *>(0)) {
callback->OnWSClientDisconnected();
}
FreeSlot(slotIdx);
}
/*---------------------------------------------------------------------------*/
/* Broadcast helpers */
/*---------------------------------------------------------------------------*/
void WSServer::BroadcastText(const char *json, uint32 len) {
for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
if (!clients[i].active) { continue; }
(void) clients[i].writeMutex.FastLock();
if (clients[i].active) {
(void) SendFrame(clients[i], WS_OPCODE_TEXT,
reinterpret_cast<const uint8 *>(json), len);
}
clients[i].writeMutex.FastUnLock();
}
}
void WSServer::BroadcastBinary(const uint8 *buf, uint32 len) {
for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
if (!clients[i].active) { continue; }
(void) clients[i].writeMutex.FastLock();
if (clients[i].active) {
(void) SendFrame(clients[i], WS_OPCODE_BINARY, buf, len);
}
clients[i].writeMutex.FastUnLock();
}
}
void WSServer::SendText(uint32 slotIdx, const char *json, uint32 len) {
if (slotIdx >= WS_MAX_CLIENTS) { return; }
WSClientSlot &slot = clients[slotIdx];
if (!slot.active) { return; }
(void) slot.writeMutex.FastLock();
if (slot.active) {
(void) SendFrame(slot, WS_OPCODE_TEXT,
reinterpret_cast<const uint8 *>(json), len);
}
slot.writeMutex.FastUnLock();
}
uint32 WSServer::ClientCount() const {
(void) clientsMutex.FastLock();
uint32 c = numClients;
clientsMutex.FastUnLock();
return c;
}
/*---------------------------------------------------------------------------*/
/* Internal helpers */
/*---------------------------------------------------------------------------*/
bool WSServer::SendFrame(WSClientSlot &slot, uint8 opcode,
const uint8 *payload, uint32 payloadLen) {
if (!slot.active || (slot.sock == static_cast<BasicTCPSocket *>(0))) {
return false;
}
uint8 hdrBuf[10];
uint32 hdrLen = WSEncodeHeader(hdrBuf, opcode,
static_cast<uint64>(payloadLen));
/* Write header */
uint32 sz = hdrLen;
if (!slot.sock->Write(reinterpret_cast<const char *>(hdrBuf), sz)) {
return false;
}
/* Write payload */
sz = payloadLen;
if (payloadLen > 0u) {
if (!slot.sock->Write(reinterpret_cast<const char *>(payload), sz)) {
return false;
}
}
return true;
}
uint32 WSServer::AllocSlot(BasicTCPSocket *sock) {
(void) clientsMutex.FastLock();
uint32 idx = WS_MAX_CLIENTS;
for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
if (!clients[i].active) {
clients[i].sock = sock;
clients[i].active = true;
clients[i].readTid = MARTe::InvalidThreadIdentifier;
idx = i;
numClients++;
break;
}
}
clientsMutex.FastUnLock();
return idx;
}
void WSServer::FreeSlot(uint32 idx) {
if (idx >= WS_MAX_CLIENTS) { return; }
/* HI-5: acquire writeMutex before modifying active/sock to prevent
* use-after-free when BroadcastText/BroadcastBinary are iterating. */
(void) clients[idx].writeMutex.FastLock();
(void) clientsMutex.FastLock();
if (clients[idx].active) {
clients[idx].active = false;
if (clients[idx].sock != static_cast<BasicTCPSocket *>(0)) {
clients[idx].sock->Close();
delete clients[idx].sock;
clients[idx].sock = static_cast<BasicTCPSocket *>(0);
}
if (numClients > 0u) { numClients--; }
}
clientsMutex.FastUnLock();
clients[idx].writeMutex.FastUnLock();
}
} /* namespace StreamHub */