TCPServer.cpp

Go to the documentation of this file.

#include "Network/TCPServer.h"
 
#include <arpa/inet.h>
#include <cerrno>
#include <cstring>
#include <fcntl.h>
#include <netinet/in.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <unistd.h>
 
namespace Ember {
namespace Monitor {
 
TCPServer::TCPServer() {}
 
TCPServer::~TCPServer() { Stop(); }
 
bool TCPServer::SetNonBlocking(int socket) {
    int flags = fcntl(socket, F_GETFL, 0);
    if (flags == -1) {
        return false;
    }
    return fcntl(socket, F_SETFL, flags | O_NONBLOCK) != -1;
}
 
bool TCPServer::Start(uint16_t port) {
    if (IsRunning()) {
        Stop();
    }
 
    m_serverSocket = socket(AF_INET, SOCK_STREAM, 0);
    if (m_serverSocket < 0) {
        return false;
    }
 
    int opt = 1;
    if (setsockopt(m_serverSocket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        close(m_serverSocket);
        m_serverSocket = -1;
        return false;
    }
 
    if (!SetNonBlocking(m_serverSocket)) {
        close(m_serverSocket);
        m_serverSocket = -1;
        return false;
    }
 
    struct sockaddr_in addr;
    std::memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(port);
 
    if (bind(m_serverSocket, reinterpret_cast<struct sockaddr *>(&addr), sizeof(addr)) < 0) {
        close(m_serverSocket);
        m_serverSocket = -1;
        return false;
    }
 
    if (listen(m_serverSocket, 1) < 0) {
        close(m_serverSocket);
        m_serverSocket = -1;
        return false;
    }
 
    m_port = port;
    return true;
}
 
void TCPServer::Stop() {
    DisconnectClient();
 
    if (m_serverSocket >= 0) {
        close(m_serverSocket);
        m_serverSocket = -1;
    }
    m_port = 0;
}
 
void TCPServer::DisconnectClient() {
    if (m_clientSocket >= 0) {
        close(m_clientSocket);
        m_clientSocket = -1;
        m_clientAddress.clear();
        m_clientName.clear();
        m_recvBuffer.clear();
 
        if (m_disconnectCallback) {
            m_disconnectCallback();
        }
    }
}
 
void TCPServer::AcceptConnection() {
    struct sockaddr_in clientAddr;
    socklen_t addrLen = sizeof(clientAddr);
 
    int newSocket = accept(m_serverSocket, reinterpret_cast<struct sockaddr *>(&clientAddr), &addrLen);
    if (newSocket < 0) {
        return;
    }
 
    if (IsClientConnected()) {
        DisconnectClient();
    }
 
    if (!SetNonBlocking(newSocket)) {
        close(newSocket);
        return;
    }
 
    m_clientSocket = newSocket;
 
    char addrStr[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &clientAddr.sin_addr, addrStr, INET_ADDRSTRLEN);
    m_clientAddress = std::string(addrStr) + ":" + std::to_string(ntohs(clientAddr.sin_port));
 
    if (m_connectCallback) {
        m_connectCallback(m_clientAddress);
    }
}
 
void TCPServer::ReadFromClient() {
    if (!IsClientConnected()) {
        return;
    }
 
    uint8_t tempBuffer[65536];
    for (;;) {
        ssize_t bytesRead = recv(m_clientSocket, tempBuffer, sizeof(tempBuffer), 0);
 
        if (bytesRead < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                break;
            }
            DisconnectClient();
            return;
        }
 
        if (bytesRead == 0) {
            DisconnectClient();
            return;
        }
 
        m_recvBuffer.insert(m_recvBuffer.end(), tempBuffer, tempBuffer + bytesRead);
    }
 
    while (m_recvBuffer.size() >= Network::FRAME_HEADER_SIZE) {
        uint32_t payloadLength = 0;
        Ember::Protocol::MessageType msgType;
 
        if (!Network::MessageFrame::ParseHeader(m_recvBuffer.data(), m_recvBuffer.size(), payloadLength, msgType)) {
            break;
        }
 
        size_t totalFrameSize = Network::FRAME_HEADER_SIZE + payloadLength;
        if (m_recvBuffer.size() < totalFrameSize) {
            break;
        }
 
        Network::MessageFrame frame = Network::MessageFrame::FromBytes(m_recvBuffer.data(), totalFrameSize);
 
        m_recvBuffer.erase(m_recvBuffer.begin(), m_recvBuffer.begin() + totalFrameSize);
 
        if (frame.IsValid() && m_messageCallback) {
            m_messageCallback(frame);
        }
    }
}
 
void TCPServer::Poll() {
    if (!IsRunning()) {
        return;
    }
 
    fd_set readFds;
    FD_ZERO(&readFds);
    FD_SET(m_serverSocket, &readFds);
 
    int maxFd = m_serverSocket;
 
    if (IsClientConnected()) {
        FD_SET(m_clientSocket, &readFds);
        if (m_clientSocket > maxFd) {
            maxFd = m_clientSocket;
        }
    }
 
    struct timeval timeout;
    timeout.tv_sec = 0;
    timeout.tv_usec = 0;
 
    int activity = select(maxFd + 1, &readFds, nullptr, nullptr, &timeout);
    if (activity < 0) {
        return;
    }
 
    if (FD_ISSET(m_serverSocket, &readFds)) {
        AcceptConnection();
    }
 
    if (IsClientConnected() && FD_ISSET(m_clientSocket, &readFds)) {
        ReadFromClient();
    }
}
 
bool TCPServer::Send(const Network::MessageFrame &frame) {
    if (!IsClientConnected()) {
        return false;
    }
 
    std::vector<uint8_t> data = frame.ToBytes();
    size_t totalSent = 0;
 
    while (totalSent < data.size()) {
        ssize_t sent = send(m_clientSocket, data.data() + totalSent, data.size() - totalSent, MSG_NOSIGNAL);
        if (sent < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                continue;
            }
            DisconnectClient();
            return false;
        }
        totalSent += static_cast<size_t>(sent);
    }
 
    return true;
}
 
} // namespace Monitor
} // namespace Ember