Implemented and fixed many issues

Source/Applications/StreamHub/WSServer.cpp

@@ -8,6 +8,7 @@
 #include "SHA1.h"
 #include "Base64.h"
 #include "AdvancedErrorManagement.h"
+#include "Select.h"
 #include "Sleep.h"
 #include "Threads.h"
 #include "TimeoutType.h"
@@ -57,8 +58,10 @@ static const char *FindSubstr(const char *s, const char *pattern) {
 
 WSServer::WSServer()
     : numClients(0u),
+      liveReadThreads(0u),
       callback(static_cast<WSCommandCallback *>(0)),
       running(false),
+      numAllowedOrigins(0u),
       acceptTid(MARTe::InvalidThreadIdentifier) {
 
     for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
@@ -66,6 +69,20 @@ WSServer::WSServer()
         clients[i].active = false;
         clients[i].readTid = MARTe::InvalidThreadIdentifier;
     }
+    for (uint32 i = 0u; i < WS_MAX_ORIGINS; i++) {
+        allowedOrigins[i][0] = '\0';
+    }
+}
+
+bool WSServer::AddAllowedOrigin(const char *origin) {
+    if ((origin == static_cast<const char *>(0)) || (origin[0] == '\0')) {
+        return false;
+    }
+    if (numAllowedOrigins >= WS_MAX_ORIGINS) { return false; }
+    if (strlen(origin) >= WS_MAX_ORIGIN_LEN) { return false; }
+    strcpy(allowedOrigins[numAllowedOrigins], origin);
+    numAllowedOrigins++;
+    return true;
 }
 
 WSServer::~WSServer() {
@@ -104,9 +121,9 @@ bool WSServer::Start(uint16 port, WSCommandCallback *cb) {
 bool WSServer::Stop() {
     if (!running) { return true; }
     running = false;
-    Sleep::MSec(200u);
 
-    /* Close all client connections — their read threads will exit on error */
+    /* Close all client connections — their read threads wake out of select()
+     * and unwind through FreeSlot. */
     (void) clientsMutex.FastLock();
     for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
         if (clients[i].active && (clients[i].sock != static_cast<BasicTCPSocket *>(0))) {
@@ -114,10 +131,23 @@ bool WSServer::Stop() {
         }
     }
     clientsMutex.FastUnLock();
-    Sleep::MSec(200u);
 
+    /* The accept loop polls WaitConnection with a 500 ms timeout, so it is out
+     * of the listener by now. */
+    Sleep::MSec(600u);
     tcpListener.Close();
-    Sleep::MSec(100u);
+
+    /* Wait for the read threads: they hold pointers to the sockets freed
+     * below.  Bounded — leaking a socket at exit beats deleting one that a
+     * wedged thread is still reading from. */
+    static const uint32 kReadJoinMs = 3000u;
+    for (uint32 waited = 0u; waited < kReadJoinMs; waited += 20u) {
+        (void) clientsMutex.FastLock();
+        const uint32 live = liveReadThreads;
+        clientsMutex.FastUnLock();
+        if (live == 0u) { break; }
+        Sleep::MSec(20u);
+    }
 
     /* Free any remaining slots */
     (void) clientsMutex.FastLock();
@@ -170,6 +200,10 @@ void WSServer::AcceptLoop() {
         }
 
         /* Start per-client read thread */
+        (void) clientsMutex.FastLock();
+        liveReadThreads++;
+        clientsMutex.FastUnLock();
+
         ClientThreadArg *arg = new ClientThreadArg();
         arg->srv  = this;
         arg->slot = slot;
@@ -200,12 +234,28 @@ bool WSServer::UpgradeHTTP(BasicTCPSocket *sock) {
     }
 
     /* 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. */
+     * If an Origin header is present it must either be on the configured
+     * allowlist or 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++; }
+
+        /* Full origin value "scheme://host[:port]", for the allowlist. */
+        char originFull[WS_MAX_ORIGIN_LEN];
+        uint32 ofLen = 0u;
+        while ((originHdr[ofLen] != '\r') && (originHdr[ofLen] != '\n') &&
+               (originHdr[ofLen] != '\0') && (ofLen < (WS_MAX_ORIGIN_LEN - 1u))) {
+            originFull[ofLen] = originHdr[ofLen];
+            ofLen++;
+        }
+        originFull[ofLen] = '\0';
+        bool allowed = false;
+        for (uint32 i = 0u; (i < numAllowedOrigins) && !allowed; i++) {
+            if (strcmp(originFull, allowedOrigins[i]) == 0) { allowed = true; }
+        }
+
         /* Extract the host part of Origin: "scheme://host[:port]" */
         char originHost[256];
         uint32 ohLen = 0u;
@@ -221,7 +271,7 @@ bool WSServer::UpgradeHTTP(BasicTCPSocket *sock) {
 
         /* Extract Host header value */
         const char *hostHdr = FindSubstr(hdrBuf, "Host:");
-        if (hostHdr != static_cast<const char *>(0)) {
+        if (!allowed && (hostHdr != static_cast<const char *>(0))) {
             hostHdr += 5; /* skip "Host:" */
             while (*hostHdr == ' ') { hostHdr++; }
             char hostVal[256];
@@ -299,23 +349,30 @@ void WSServer::ClientReadLoop(uint32 slotIdx) {
     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;
+
+        /* Wait for readability before reading.  BasicTCPSocket::Read reports a
+         * timeout and a closed peer identically (false, zero bytes), so polling
+         * it on its own cannot end the loop: once the client goes away recv
+         * returns immediately and forever, and the thread spins at 100% CPU
+         * until it starves the rest of the hub.  select() tells the two apart —
+         * readable followed by no data is end of stream.  A wait consumes the
+         * handle set, hence a fresh Select each pass. */
+        MARTe::Select sel;
+        if (!sel.AddReadHandle(*sock)) { break; }
+        const MARTe::int32 ready = sel.WaitUntil(TimeoutType(500u));
+        if (ready == 0) { continue; }   /* idle client — recheck running */
+        if (ready < 0) { break; }       /* socket closed or errored */
+
+        /* Readable: this returns at once, and only fails at end of stream. */
+        if (!sock->Read(reinterpret_cast<char *>(buf + filled), want,
+                        TimeoutType(500u))) {
+            break;
         }
         filled += want;
 
@@ -383,6 +440,10 @@ client_done:
         callback->OnWSClientDisconnected();
     }
     FreeSlot(slotIdx);
+
+    (void) clientsMutex.FastLock();
+    if (liveReadThreads > 0u) { liveReadThreads--; }
+    clientsMutex.FastUnLock();
 }
 
 /*---------------------------------------------------------------------------*/