server.cpp

#include <arpa/inet.h>
#include <errno.h>
#include <netinet/in.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/timerfd.h>
#include <sys/types.h>
#include <unistd.h>
#include <netinet/tcp.h>
#include <math.h>
#include <vector>
#include <string>

#include "common.h"
#include "helpers.h"

#define MAX_CONNECTIONS 5

std::vector<tcp_client> clients;
std::vector<struct pollfd> poll_fds(3);

void tcp_connection(int tcp_socket) {
    struct sockaddr_in cli_addr;
    socklen_t cli_len = sizeof(cli_addr);
    memset(&cli_addr, 0, cli_len);

    // accept the new connection
    int newsockfd = accept(tcp_socket, (struct sockaddr *)&cli_addr, &cli_len);
    DIE(newsockfd < 0, "ERROR on accept");

    // disable Nagle's algorithm for the tcp socket
    const int enable = 1;
    DIE(setsockopt(newsockfd, IPPROTO_TCP, TCP_NODELAY, &enable, sizeof(int)) < 0, "setsockopt");

    // add the new socket to the poll_fds array
    struct pollfd new_poll_fd;
    new_poll_fd.fd = newsockfd;
    new_poll_fd.events = POLLIN;
    poll_fds.push_back(new_poll_fd);

    // get the id client
    protocol recv_id;
    memset(&recv_id, 0, sizeof(protocol));
    int rc = recv_all(newsockfd, &recv_id.length, sizeof(int));
    DIE(rc < 0, "ERROR reading from socket");
    rc = recv_all(newsockfd, &recv_id.payload, recv_id.length);
    DIE(rc < 0, "ERROR reading from socket");

    //check if the client id already exists
    for (std::vector<pollfd>::size_type i = 0; i < clients.size(); i++) {
        if (strncmp(clients[i].id_client, recv_id.payload, recv_id.length) == 0) {
            if (clients[i].sockfd != -1) {
                printf("Client %s already connected.\n", recv_id.payload);
                close(newsockfd);
                return;
            }
            clients[i].sockfd = newsockfd;
            printf("New client %s connected from %s:%d\n", recv_id.payload,
                                                           inet_ntoa(cli_addr.sin_addr),
                                                           ntohs(cli_addr.sin_port));
            return;
        }
    }

    tcp_client new_client = {};
    new_client.sockfd = newsockfd;
    strncpy(new_client.id_client, recv_id.payload, recv_id.length);
    new_client.addr = cli_addr;
    clients.push_back(new_client);

    printf("New client %s connected from %s:%d\n", recv_id.payload,
                                                   inet_ntoa(cli_addr.sin_addr),
                                                   ntohs(cli_addr.sin_port));
}

void udp_message(int udp_socket) {
    struct sockaddr_in cli_addr;
    socklen_t cli_len = sizeof(cli_addr);
    memset(&cli_addr, 0, cli_len);

    chat_packet packet;
    memset(&packet, 0, sizeof(chat_packet));

    char buffer[1551];
    memset(buffer, 0, 1551);

    // read the entire packet
    int rc = recvfrom(udp_socket, &buffer, sizeof(buffer), 0, (struct sockaddr *)&cli_addr, &cli_len);
    DIE(rc < 0, "ERROR reading from socket");
    buffer[rc] = '\0';

    // extract the topic
    memcpy(packet.topic, buffer, 50);

    // extract the data_type
    packet.data_type = buffer[50];

    char *ptr = buffer;
    ptr += 51;

    if (packet.data_type == 0) {
        // extract the signed byte
        int8_t sign_byte = ptr[0];

        // extract the uint32_t value from the next 4 bytes (network byte order)
        uint32_t value;
        memcpy(&value, ptr + 1, sizeof(uint32_t));

        // convert from network byte order to host byte order
        value = ntohl(value);

        // combine the sign and value to get the final integer
        int32_t final_value = (sign_byte == 0) ? (int32_t)value : -(int32_t)value;

        memcpy(packet.payload, &final_value, sizeof(int32_t));
    } else if (packet.data_type == 1) {
        // extract the uint16_t value from the payload
        uint16_t value;
        memcpy(&value, ptr, sizeof(uint16_t));

        // convert from network byte order to host byte order
        value = ntohs(value);

        // divide the value by 100 to get the final short real value
        float final_value = value / 100.0;

        memcpy(packet.payload, &final_value, sizeof(float));
    } else if (packet.data_type == 2) {
        // extract the sign byte
        int8_t sign_byte = ptr[0];

        // extract the uint32_t value from the next 4 bytes (network byte order)
        uint32_t integer_part;
        memcpy(&integer_part, ptr + 1, sizeof(uint32_t));

        // convert from network byte order to host byte order
        integer_part = ntohl(integer_part);

        // extract the uint8_t value (negative power of 10)
        uint8_t power_of_10 = ptr[5];

        // calculate the float value
        double float_value = integer_part * pow(10, -power_of_10);

        // apply the sign if negative
        sign_byte == 1 ? (float_value *= -1) : (float_value *= 1);

        memcpy(packet.payload, &float_value, sizeof(double));
    } else if (packet.data_type == 3) {
        strcpy(packet.payload, ptr);
    }

    protocol send_protocol;
    memset(&send_protocol, 0, sizeof(protocol));
    send_protocol.length = sizeof(chat_packet);
    memcpy(send_protocol.payload, &packet, send_protocol.length);

    // iterate through all the clients and send the message to those who are following the topic
    for (std::vector<pollfd>::size_type i = 0; i < clients.size(); i++) {
        for (std::vector<pollfd>::size_type j = 0; j < clients[i].followed_topics.size(); j++) {
            if (strncmp(clients[i].followed_topics[j].c_str(), packet.topic, 50) == 0) {
                // check if the client is disconnected, if so, skip it
                if (clients[i].sockfd == -1) {
                    continue;
                }
                rc = send_all(clients[i].sockfd, &send_protocol, send_protocol.length + sizeof(int));
                DIE(rc < 0, "ERROR writing to socket");
            }
        }
    }
}

void exit_command(int tcp_client, int udp_client) {
    close(tcp_client);
    close(udp_client);
    for (std::vector<pollfd>::size_type i = 0; i < clients.size(); i++) {
        close(clients[i].sockfd);
    }
}

int add_followed_topic(tcp_client *client, const char *topic) {
    // check if the topic is already followed
    for (std::vector<pollfd>::size_type i = 0; i < client->followed_topics.size(); i++) {
        if (client->followed_topics[i] == topic) {
            return -1;
        }
    }

    // add the new topic
    client->followed_topics.push_back(topic);
    return 1;
}

int remove_followed_topic(tcp_client *client, const char *topic) {
    for (std::vector<pollfd>::size_type i = 0; i < client->followed_topics.size(); i++) {
        if (client->followed_topics[i] == topic) {
            client->followed_topics.erase(client->followed_topics.begin() + i);
            return 1;
        }
    }
    return -1;
}

void tcp_receive(int sockfd, std::vector<pollfd>::size_type *i) {
    // receive the packet from the client
    protocol recv_protocol;
    memset(&recv_protocol, 0, sizeof(protocol));
    int rc = recv_all(sockfd, &recv_protocol.length, sizeof(int));
    DIE(rc < 0, "ERROR reading from socket");
    rc = recv_all(sockfd, &recv_protocol.payload, recv_protocol.length);
    DIE(rc < 0, "ERROR reading from socket");

    chat_packet packet;
    memset(&packet, 0, sizeof(chat_packet));
    memcpy(&packet, recv_protocol.payload, recv_protocol.length);
    if (strncmp(packet.payload, "subscribe", 9) == 0) {
        for (std::vector<pollfd>::size_type j = 0; j < clients.size(); j++) {
            if (clients[j].sockfd == sockfd) {
                int rc = add_followed_topic(&clients[j], packet.topic);
                DIE(rc == -1, "Already subscribed to topic");
                break;
            }
        }
    } else if (strncmp(packet.payload, "unsubscribe", 11) == 0) {
        for (std::vector<pollfd>::size_type j = 0; j < clients.size(); j++) {
            if (clients[j].sockfd == sockfd) {
                int rc = remove_followed_topic(&clients[j], packet.topic);
                DIE(rc == -1, "Not subscribed to topic");
                break;
            }
        }
    } else if (strncmp(packet.payload, "exit", 4) == 0) {
        for (std::vector<pollfd>::size_type j = 0; j < clients.size(); j++) {
            if (clients[j].sockfd == sockfd) {
                close(sockfd);
                // set the client's socket to -1 in order to mark it as disconnected
                clients[j].sockfd = -1;
                // remove the client from the poll_fds array
                poll_fds.erase(poll_fds.begin() + *i);
                (*i)--;
                printf("Client %s disconnected.\n", packet.topic);
                break;
            }
        }
    }
}

void run_chat_multi_server(int tcp_socket, int udp_socket) {
    // add the standard input, tcp socket, and udp socket to the poll_fds array
    poll_fds[0].fd = STDIN_FILENO;
    poll_fds[0].events = POLLIN;
    poll_fds[1].fd = tcp_socket;
    poll_fds[1].events = POLLIN;
    poll_fds[2].fd = udp_socket;
    poll_fds[2].events = POLLIN;

    while (1) {
        int rc = poll(poll_fds.data(), poll_fds.size(), -1);
        DIE(rc < 0, "poll");

        for (std::vector<pollfd>::size_type i = 0; i < poll_fds.size(); i++) {
            if (poll_fds[i].fd == STDIN_FILENO && poll_fds[i].revents & POLLIN) {
                char buffer[10];
                fgets(buffer, 10, stdin);
                if (strncmp(buffer, "exit", 4) == 0) {
                    exit_command(tcp_socket, udp_socket);
                    return;
                }
            } else if (poll_fds[i].fd == tcp_socket && poll_fds[i].revents & POLLIN) {
                tcp_connection(tcp_socket);
                break;
            } else if (poll_fds[i].fd == udp_socket && poll_fds[i].revents & POLLIN) {
                udp_message(udp_socket);
            } else {
                if (poll_fds[i].revents & POLLIN) {
                    tcp_receive(poll_fds[i].fd, &i);
                }
            }
        }
    }
}

int main(int argc, char *argv[]) {
    setvbuf(stdout, NULL, _IONBF, BUFSIZ);
    DIE(argc != 2, "Usage: ./server <port>\n");

    int tcp_socket, udp_socket;

    // create the tcp socket
    DIE((tcp_socket = socket(AF_INET, SOCK_STREAM, 0)) == -1, "socket");

    // create the udp socket
    DIE((udp_socket = socket(AF_INET, SOCK_DGRAM, 0)) == -1, "socket");

    uint16_t port = atoi(argv[1]);
    int rc = sscanf(argv[1], "%hu", &port);
    DIE(rc != 1, "Given port is invalid");

    // enable SO_REUSEADDR for both sockets
    const int enable = 1;
    DIE(setsockopt(tcp_socket, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)) < 0, "setsockopt");
    DIE(setsockopt(udp_socket, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)) < 0, "setsockopt");

    // disable Nagle's algorithm for the tcp socket
    DIE(setsockopt(tcp_socket, IPPROTO_TCP, TCP_NODELAY, &enable, sizeof(int)) < 0, "setsockopt");

    // complete the server address structure
    struct sockaddr_in serv_addr;
    socklen_t socket_len = sizeof(struct sockaddr_in);
    memset(&serv_addr, 0, socket_len);
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_port = htons(port);
    serv_addr.sin_addr.s_addr = INADDR_ANY;

    // bind the tcp socket to the specified port
    DIE(bind(tcp_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) == -1, "bind");

    // bind the udp socket to the specified port
    DIE(bind(udp_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) == -1, "bind");

    // start listening for incoming tcp connections
    DIE(listen(tcp_socket, MAX_CONNECTIONS) == -1, "listen");

    run_chat_multi_server(tcp_socket, udp_socket);

    // close the sockets
    close(tcp_socket);
    close(udp_socket);

    return 0;
}