MARTe-Integrated-Components/Client/streamhub/WSClient.cpp

/**
 * @file WSClient.cpp
 * @brief RFC 6455 WebSocket client implementation.
 */

#include "WSClient.h"
#include "WSFrame_client.h"

#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <cstring>
#include <cstdio>
#include <cstdlib>
#include <ctime>
#include <chrono>
#include <random>

namespace StreamHubClient {

/* ── Helpers ─────────────────────────────────────────────────────────────── */

static std::string base64Key() {
    /* HI-7: use /dev/urandom (CSPRNG) instead of srand(time)/rand() */
    uint8_t raw[16];
    int fd = open("/dev/urandom", O_RDONLY);
    if (fd < 0 || read(fd, raw, sizeof(raw)) != static_cast<ssize_t>(sizeof(raw))) {
        /* Fallback: std::random_device (still better than srand/rand) */
        std::random_device rd;
        for (size_t i = 0; i < sizeof(raw); i += sizeof(unsigned)) {
            unsigned val = rd();
            for (size_t j = 0; j < sizeof(unsigned) && i + j < sizeof(raw); j++) {
                raw[i + j] = static_cast<uint8_t>(val >> (j * 8));
            }
        }
    }
    if (fd >= 0) { close(fd); }
    char out[32];
    WS_Base64Encode(raw, 16, out);
    return std::string(out);
}

/* ── Construction ────────────────────────────────────────────────────────── */

WSClient::WSClient() = default;

WSClient::~WSClient() {
    disconnect();
}

/* ── Public API ──────────────────────────────────────────────────────────── */

void WSClient::connect(const std::string& host, uint16_t port) {
    /* If already running toward the same target just signal a reconnect */
    host_ = host;
    port_ = port;
    reconnect_.store(true);

    if (!recvThread_.joinable()) {
        stopThread_.store(false);
        recvThread_ = std::thread(&WSClient::recvLoop, this);
    }
}

void WSClient::disconnect() {
    stopThread_.store(true);
    connected_.store(false);
    if (sock_ >= 0) {
        ::shutdown(sock_, SHUT_RDWR);
        ::close(sock_);
        sock_ = -1;
    }
    if (recvThread_.joinable()) {
        recvThread_.join();
    }
}

void WSClient::reconnect(const std::string& host, uint16_t port) {
    host_ = host;
    port_ = port;
    /* Closing the socket causes the recv thread to get n<=0 and loop back
     * to doConnect() with the new host/port. No thread join needed. */
    connected_.store(false);
    if (sock_ >= 0) {
        ::shutdown(sock_, SHUT_RDWR);
        ::close(sock_);
        sock_ = -1;
    }
    /* Ensure thread is running (first call before connect()) */
    if (!recvThread_.joinable()) {
        stopThread_.store(false);
        recvThread_ = std::thread(&WSClient::recvLoop, this);
    }
}

bool WSClient::isConnected() const {
    return connected_.load();
}

void WSClient::poll(const std::function<void(const WSMessage&)>& handler) {
    std::queue<WSMessage> local;
    {
        std::lock_guard<std::mutex> lk(queueMutex_);
        std::swap(local, recvQueue_);
    }
    while (!local.empty()) {
        handler(local.front());
        local.pop();
    }
}

void WSClient::sendText(const std::string& json) {
    if (!connected_.load()) { return; }

    uint8_t hdrBuf[10];
    uint32_t hdrLen = WSEncodeHeader(hdrBuf, WS_OPCODE_TEXT,
                                     static_cast<uint64_t>(json.size()));
    std::lock_guard<std::mutex> lk(sendMutex_);
    if (!sendRaw(hdrBuf, hdrLen)) { connected_.store(false); return; }
    if (!json.empty()) {
        if (!sendRaw(reinterpret_cast<const uint8_t*>(json.data()), json.size())) {
            connected_.store(false);
        }
    }
}

void WSClient::sendBinary(const uint8_t* data, size_t len) {
    if (!connected_.load()) { return; }

    uint8_t hdrBuf[10];
    uint32_t hdrLen = WSEncodeHeader(hdrBuf, WS_OPCODE_BINARY,
                                     static_cast<uint64_t>(len));
    std::lock_guard<std::mutex> lk(sendMutex_);
    if (!sendRaw(hdrBuf, hdrLen)) { connected_.store(false); return; }
    if (len > 0 && data != nullptr) {
        if (!sendRaw(data, len)) {
            connected_.store(false);
        }
    }
}

/* ── Receive loop (background thread) ────────────────────────────────────── */

void WSClient::recvLoop() {
    /* Per-connection receive buffer: declared here so it can be cleared on
     * each reconnect, preventing stale WS frames from a prior session from
     * being parsed as frames of the new connection. */
    std::vector<uint8_t> recvBuf;
    recvBuf.reserve(1u << 20); /* 1 MiB initial reservation */

    while (!stopThread_.load()) {
        if (!connected_.load()) {
            recvBuf.clear(); /* discard any partial data from previous connection */
            if (!doConnect() || !doHandshake()) {
                connected_.store(false);
                if (sock_ >= 0) { ::close(sock_); sock_ = -1; }
                /* Wait 3 s before retry */
                for (int i = 0; i < 30 && !stopThread_.load(); i++) {
                    std::this_thread::sleep_for(std::chrono::milliseconds(100));
                }
                continue;
            }
            connected_.store(true);
            reconnect_.store(false);
            fprintf(stderr, "WSClient: connected to %s:%u\n", host_.c_str(), port_);
        }

        /* Read and accumulate bytes */

        uint8_t chunk[4096];
        ssize_t n = ::recv(sock_, chunk, sizeof(chunk), 0);
        if (n <= 0) {
            if (n == 0 || (errno != EAGAIN && errno != EWOULDBLOCK)) {
                fprintf(stderr, "WSClient: recv error / connection closed\n");
                connected_.store(false);
                ::close(sock_); sock_ = -1;
            }
            continue;
        }
        recvBuf.insert(recvBuf.end(), chunk, chunk + n);

        /* Parse as many complete frames as possible */
        size_t consumed = 0;
        while (consumed < recvBuf.size()) {
            const uint8_t* frameStart = recvBuf.data() + consumed;
            size_t available = recvBuf.size() - consumed;

            WSFrameHeader hdr;
            if (!WSParseHeader(frameStart, static_cast<uint32_t>(available), hdr)) {
                break; /* Need more data */
            }
            if (hdr.payloadLen > 16u * 1024u * 1024u) {
                /* Frame too large — disconnect */
                fprintf(stderr, "WSClient: frame too large (%llu bytes)\n",
                        static_cast<unsigned long long>(hdr.payloadLen));
                connected_.store(false);
                ::close(sock_); sock_ = -1;
                recvBuf.clear();
                goto next_iter;
            }
            {
                size_t total = hdr.headerSize + static_cast<size_t>(hdr.payloadLen);
                if (available < total) { break; }

                uint8_t* payload = const_cast<uint8_t*>(frameStart) + hdr.headerSize;
                uint32_t plen    = static_cast<uint32_t>(hdr.payloadLen);
                if (hdr.masked) { WSUnmask(payload, plen, hdr.maskKey); }

                if (hdr.opcode == WS_OPCODE_CLOSE) {
                    connected_.store(false);
                    ::close(sock_); sock_ = -1;
                    recvBuf.clear();
                    goto next_iter;
                } else if (hdr.opcode == WS_OPCODE_PING) {
                    /* Send pong */
                    uint8_t pongHdr[10];
                    uint32_t pongHdrLen = WSEncodeHeader(pongHdr, WS_OPCODE_PONG, plen);
                    std::lock_guard<std::mutex> lk(sendMutex_);
                    sendRaw(pongHdr, pongHdrLen);
                    if (plen > 0) { sendRaw(payload, plen); }
                } else if (hdr.opcode == WS_OPCODE_TEXT ||
                           hdr.opcode == WS_OPCODE_BINARY) {
                    WSMessage msg;
                    msg.isBinary = (hdr.opcode == WS_OPCODE_BINARY);
                    msg.data.assign(payload, payload + plen);
                    std::lock_guard<std::mutex> lk(queueMutex_);
                    recvQueue_.push(std::move(msg));
                }
                consumed += total;
            }
        }

        /* Compact buffer */
        if (consumed > 0) {
            recvBuf.erase(recvBuf.begin(), recvBuf.begin() + consumed);
        }
        next_iter:;
    }

    if (sock_ >= 0) { ::close(sock_); sock_ = -1; }
    connected_.store(false);
}

/* ── TCP connect ─────────────────────────────────────────────────────────── */

bool WSClient::doConnect() {
    struct addrinfo hints = {};
    hints.ai_family   = AF_INET;
    hints.ai_socktype = SOCK_STREAM;

    char portStr[16];
    snprintf(portStr, sizeof(portStr), "%u", static_cast<unsigned>(port_));

    struct addrinfo* res = nullptr;
    if (getaddrinfo(host_.c_str(), portStr, &hints, &res) != 0 || res == nullptr) {
        fprintf(stderr, "WSClient: getaddrinfo failed for %s\n", host_.c_str());
        return false;
    }

    sock_ = ::socket(res->ai_family, res->ai_socktype, res->ai_protocol);
    if (sock_ < 0) { freeaddrinfo(res); return false; }

    if (::connect(sock_, res->ai_addr, res->ai_addrlen) < 0) {
        fprintf(stderr, "WSClient: connect failed: %s\n", strerror(errno));
        ::close(sock_); sock_ = -1;
        freeaddrinfo(res);
        return false;
    }
    freeaddrinfo(res);
    return true;
}

/* ── HTTP/1.1 WebSocket upgrade ──────────────────────────────────────────── */

bool WSClient::doHandshake() {
    std::string key = base64Key();

    /* Send HTTP upgrade request */
    char req[512];
    snprintf(req, sizeof(req),
        "GET / HTTP/1.1\r\n"
        "Host: %s:%u\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Key: %s\r\n"
        "Sec-WebSocket-Version: 13\r\n"
        "\r\n",
        host_.c_str(), static_cast<unsigned>(port_), key.c_str());

    if (!sendRaw(reinterpret_cast<const uint8_t*>(req), strlen(req))) {
        return false;
    }

    /* Read HTTP response until \r\n\r\n */
    char resp[2048];
    size_t totalRead = 0;
    while (totalRead < sizeof(resp) - 1) {
        ssize_t n = ::recv(sock_, resp + totalRead, 1, 0);
        if (n <= 0) { return false; }
        totalRead++;
        resp[totalRead] = '\0';
        if (totalRead >= 4 &&
            memcmp(resp + totalRead - 4, "\r\n\r\n", 4) == 0) {
            break;
        }
    }

    /* Check for 101 */
    if (strstr(resp, "101") == nullptr) {
        fprintf(stderr, "WSClient: handshake failed:\n%s\n", resp);
        return false;
    }

    /* Verify Sec-WebSocket-Accept (optional but good practice) */
    static const char* kGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
    std::string concat = key + kGUID;
    uint8_t digest[20];
    WS_SHA1(reinterpret_cast<const uint8_t*>(concat.data()),
            static_cast<uint32_t>(concat.size()), digest);
    char expectedAccept[32];
    WS_Base64Encode(digest, 20, expectedAccept);

    if (strstr(resp, expectedAccept) == nullptr) {
        fprintf(stderr, "WSClient: Sec-WebSocket-Accept mismatch\n");
        return false;
    }

    return true;
}

/* ── Low-level I/O ───────────────────────────────────────────────────────── */

bool WSClient::recvExact(uint8_t* buf, size_t len) {
    size_t received = 0;
    while (received < len) {
        ssize_t n = ::recv(sock_, buf + received, len - received, MSG_WAITALL);
        if (n <= 0) { return false; }
        received += static_cast<size_t>(n);
    }
    return true;
}

bool WSClient::sendRaw(const uint8_t* buf, size_t len) {
    size_t sent = 0;
    while (sent < len) {
        ssize_t n = ::send(sock_, buf + sent, len - sent, MSG_NOSIGNAL);
        if (n <= 0) { return false; }
        sent += static_cast<size_t>(n);
    }
    return true;
}

} /* namespace StreamHubClient */