server.cpp

#include "utils.h"
#include "pattern.h"

#define INITIAL_CONNECTIONS 3

bool checkIfIdExist(std::unordered_map<std::string, int> connected_clients, char id[11]) {
    return connected_clients.find(std::string(id)) != connected_clients.end();
}

void close_server(std::vector<struct pollfd> &fds) {
    for (uint32_t i = 0; i < fds.size(); i++) {
        if (fds[i].fd != STDIN_FILENO) {
            close(fds[i].fd);
        }
    }
}

int findId(std::vector<std::string> vector, std::string str) {
    for (size_t i = 0; i < vector.size(); i++) {
        if (vector[i] == str) {
            return i;
        }
    }
    return -1;
}

void run_server(int tcp_fd, int udp_fd) {
    std::vector<struct pollfd> fds(INITIAL_CONNECTIONS, {0, 0, 0});
    fds[0].fd = udp_fd;
    fds[0].events = POLLIN;

    fds[1].fd = tcp_fd;
    fds[1].events = POLLIN;

    fds[2].fd = STDIN_FILENO;
    fds[2].events = POLLIN;

    std::unordered_map<std::string, int> connected_clients;
    std::vector<std::string> clients;
    std::unordered_map<std::string, std::unordered_set<std::string>> client_topics;

    while (1) {
        int ready = poll(fds.data(), fds.size(), 2);
        if (ready < 0) {
            perror("Error at poll");
            close_server(fds);
            break;
        }

        for (uint32_t i = 0; i < fds.size(); i++) {
            if (fds[i].fd == STDIN_FILENO
            && fds[i].revents & POLLIN) {
                char buffer[MAX_SIZE_FROM_READ];
                int bytes = read(fds[i].fd, buffer, sizeof(buffer));
                if (bytes < 0) {
                    perror("Error when reading from stdin");
                    close_server(fds);
                    exit(1);
                }
                buffer[strcspn(buffer, "\n")] = '\0';
                if (strcmp(buffer, "exit") == 0) {
                    for (auto &pair : connected_clients) {
                        struct tcp_packet_to_client message;
                        message.type = EXIT_CODE;
                        send_message(pair.second, &message, sizeof(message));
                    }
                    close_server(fds);
                    return;
                } else {
                    WRITE_ERROR("That command isn't supported on the server\n");
                }
            } else if (fds[i].fd == udp_fd
            && fds[i].revents & POLLIN) {
                char buffer[MAX_SIZE_TOPIC];
                struct sockaddr_in client;
                socklen_t length = sizeof(struct sockaddr_in);

                int bytes_received = recvfrom(fds[i].fd,
                                    buffer, sizeof(buffer),
                                    0, (struct sockaddr *)&client,
                                    &length);
                if (bytes_received < 0) {
                    close_server(fds);
                    perror("Error when receiving an UDP packet\n");
                    exit(1);
                } else {
                    struct udp_packet *packet = (struct udp_packet *)buffer;
                    struct tcp_packet_to_client message;
                    message.ip_udp = client.sin_addr.s_addr;
                    message.port_udp = client.sin_port;
                    message.type = TOPIC;
                    memcpy(&message.packet, packet, sizeof(struct udp_packet));
                    std::string topic(packet->topic);
                    for (const auto& pair : client_topics) {
                        std::unordered_set<std::string> senders;
                        if (pair.first.find("*") == std::string::npos
                        && pair.first.find("+") == std::string::npos) {
                            if (pair.first == topic) {
                                auto it = pair.second;
                                for (auto client_id : it) {
                                    if (senders.find(client_id) == senders.end()) {
                                        int fd = connected_clients[client_id];
                                        if (fd == 0) {
                                            continue;
                                        }
                                        send_message(fd, &message, sizeof(message));
                                        senders.insert(client_id);
                                    }
                                }
                            }
                        } else {
                            std::string tmp(pair.first);
                            std::regex regex_string(get_regex_string(tmp));
                            if (std::regex_match(topic, regex_string)) {
                                auto it = pair.second;
                                for (auto client_id : it) {
                                    if (senders.find(client_id) == senders.end()) {
                                        int fd = connected_clients[client_id];
                                        if (fd == 0) {
                                            continue;
                                        }
                                        send_message(fd, &message, sizeof(message));
                                        senders.insert(client_id);
                                    }
                                }
                            }
                        }
                    }
                }
            } else if (fds[i].fd == tcp_fd
            && fds[i].revents & POLLIN) {
                struct sockaddr_in tcp_client;
                socklen_t length = sizeof(struct sockaddr_in);

                int new_socket = accept(fds[i].fd,
                                (struct sockaddr *)&tcp_client, &length);
                if (new_socket < 0) {
                    perror("Error at accept system call");
                    close_server(fds);
                    exit(1);
                }

                int option = 1;
                if (setsockopt(new_socket, IPPROTO_TCP, TCP_NODELAY, &option, sizeof(option)) < 0) {
                    perror("Error at setsockopt");
                    close(new_socket);
                    close_server(fds);
                    exit(1);
                }

                char buffer[MAX_SIZE_PACKET_SERVER];
                int bytes = recv_all(new_socket, buffer, sizeof(struct tcp_packet_to_server));
                if (bytes < 0) {
                    perror("Error at recv_all");
                    close(new_socket);
                    close_server(fds);
                    exit(1);
                }
                if (bytes == 0) {
                    close(new_socket);
                    continue;
                }
                struct tcp_packet_to_server *request = (struct tcp_packet_to_server *)buffer;
                if (findId(clients, std::string(request->id)) != -1) {
                    printf("Client %s already connected.\n", request->id);
                    struct tcp_packet_to_client tmp;
                    tmp.type = EXIT_CODE;
                    send_message(new_socket, &tmp, sizeof(tmp));
                    close(new_socket);
                    continue;
                } else {
                    uint16_t port = ntohs(tcp_client.sin_port);
                    printf("New client %s connected from %s:%d.\n", request->id, inet_ntoa(tcp_client.sin_addr), port);
                    fds.push_back({new_socket, POLLIN, 0});
                    connected_clients[request->id] = new_socket;
                    clients.push_back(std::string(request->id));
                }

            } else if (fds[i].revents & POLLIN) {
                char buffer[MAX_SIZE_PACKET_SERVER];
                int bytes = recv_all(fds[i].fd, buffer, sizeof(struct tcp_packet_to_server));
                if (bytes < 0) {
                    perror("Error at recv_all");
                    close_server(fds);
                    exit(1);
                }
                if (bytes == 0) {
                    for (auto &pair : connected_clients) {
                        if (pair.second == fds[i].fd) {
                            printf("Client %s disconnected.\n", pair.first.c_str());
                            connected_clients.erase(pair.first);
                            break;
                        }
                    }
                    fds.erase(fds.begin() + i);
                    close(fds[i].fd);
                }
                struct tcp_packet_to_server *request = (struct tcp_packet_to_server *) buffer;
                int j;
                switch (request->type) {
                    case EXIT:
                        j = findId(clients, std::string(request->id));
                        if (j != -1) {
                            close(fds[i].fd);
                            connected_clients.erase(request->id);
                            fds.erase(fds.begin() + i);
                            printf("Client %s disconnected.\n", request->id);
                            clients.erase(clients.begin() + j);
                        } else {
                            WRITE_ERROR("The client isn't connected\n");
                        }
                        break;
                    case SUBSCRIBE:
                        client_topics[request->topic].insert(request->id);
                        break;
                    case UNSUBSCRIBE:
                        client_topics[request->topic].erase(request->id);
                        break;
                    default:
                        break;
                }
            }
        }
    }
}

int main(int argc, char* argv[]) {
    if (argc != 2) {
        WRITE_ERROR("Usage: ./server <PORT_NUMBER>: Invalid number of arguments\n");
        return 1;
    }
    setvbuf(stdout, NULL, _IONBF, BUFSIZ);

    int tcp_socket_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (tcp_socket_fd < 0) {
        WRITE_ERROR("Problem at tcp socket\n");
        return 1;
    }

    int udp_socket_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (udp_socket_fd < 0) {
        WRITE_ERROR("Problem at the udp socket\n");
        close(tcp_socket_fd);
        return 1;
    }

    const short port = atoi(argv[1]);
    struct sockaddr_in socket_address;
    socket_address.sin_family = AF_INET;
    socket_address.sin_port = htons(port);
    socket_address.sin_addr.s_addr = htonl(INADDR_ANY);

    int option = 1;
    if (setsockopt(udp_socket_fd, SOL_SOCKET, SO_REUSEADDR, &option, sizeof(option)) < 0) {
        perror("Error at setsockopt");
        close(udp_socket_fd);
        close(tcp_socket_fd);
        return 1;
    }

    if (setsockopt(tcp_socket_fd, SOL_SOCKET, SO_REUSEADDR, &option, sizeof(option)) < 0) {
        perror("Error at setsockopt");
        close(udp_socket_fd);
        close(tcp_socket_fd);
        return 1;
    }

    if (bind(tcp_socket_fd, (sockaddr *)&socket_address, sizeof(socket_address))) {
        perror("Error at tcp bind");
        close(udp_socket_fd);
        close(tcp_socket_fd);
        return 1;
    }

    if (bind(udp_socket_fd, (sockaddr *)&socket_address, sizeof(socket_address))) {
        perror("Error at udp bind");
        close(udp_socket_fd);
        close(tcp_socket_fd);
        return 1;
    }

    if (listen(tcp_socket_fd, SOMAXCONN) < 0) {
        perror("Error at listen");
        close(udp_socket_fd);
        close(tcp_socket_fd);
        return 1;
    }

    run_server(tcp_socket_fd, udp_socket_fd);
    close(tcp_socket_fd);
    close(udp_socket_fd);
    return 0;
}