chatServer.cpp

#include <iostream>
#include <unistd.h>
#include <bits/stdc++.h>

#include <arpa/inet.h>
#include <netdb.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>

#include <pthread.h>
#include <time.h>

#define BUFLEN 2048
#define SERVER_IP_ADDRESS "127.0.0.1"
#define PORT 5300

// Structure to manage client info
struct client{
    char username[20];
    int uid;
    int client_sock_fd;
};

// Keep track of FDs of all connected clients
std::vector<int> sock_fds;

// Keep track of all users and their sock_fds
std::unordered_map<int, int> users;

std::unordered_map<std::string, int> usernames;

// Available commands
std::vector<std::string> commands = {"whisper", "view"};

// Mutex for sync
pthread_mutex_t sock_fds_mutex PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t users_mutex PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t usernames_mutex PTHREAD_MUTEX_INITIALIZER;

// Trims strings of whitespaces
std::string trim(const std::string& str) {
    size_t start = str.find_first_not_of(" \n\r\t");
    size_t end = str.find_last_not_of(" \n\r\t");
    return (start == std::string::npos || end == std::string::npos) ? "" : str.substr(start, end - start + 1);
}

// Broadcasts received messages to all clients
void broadcastMsg(const char *data, const char *username, int uid){
    char message[BUFLEN];
    snprintf(message, sizeof(message), "%s [%d]: %s", username, uid, data);

    pthread_mutex_lock(&sock_fds_mutex);
    for(auto sock_fd : sock_fds){
        write(sock_fd, message, strlen(message));
    }
    pthread_mutex_unlock(&sock_fds_mutex);
}

// Whisper private message to a client
int whisperMsg(int uid, const char *username, const char *private_msg){
    int sock_fd;
    char message[BUFLEN];
    snprintf(message, sizeof(message), "[WHISPER] %s [%d]: %s", username, uid, private_msg);

    pthread_mutex_lock(&users_mutex);
    sock_fd = users[uid];
    pthread_mutex_unlock(&users_mutex);

    if(write(sock_fd, message, sizeof(message))==-1){
        std::cout << "Failed to send whisper response message to client!" << std::endl;
        return 0;
    }
    return 1;
}

int viewAllUsers(int sock_fd){
    char message[BUFLEN];
    bzero(message, sizeof(message));
    strcpy(message, "\nAll connected users:\n");

    pthread_mutex_lock(&usernames_mutex);
    for(auto it: usernames){
        char user[100];
        bzero(user, sizeof(user));
        snprintf(user, sizeof(user), "%s : %d\n", it.first.c_str(), it.second);
        strcat(message, user);
    }
    pthread_mutex_unlock(&usernames_mutex);

    if(write(sock_fd, message, strlen(message))==-1){
        std::cout << "Failed to send view clients response to client!" << std::endl;
        return 0;
    }
    return 1;
}

int sendInvalidCommandError(int sock_fd){
    char message[BUFLEN];
    bzero(message, sizeof(message));

    snprintf(message, sizeof(message), "Invalid command!");
    if(write(sock_fd, message, sizeof(message))==-1){
        std::cout << "Failed to send invalid command response to client!" << std::endl;
        return 0;
    }
    return 1;
}

// Read data from clients
void *handleClient(void *client){
    char username[20];
    bzero(username, sizeof(username));
    struct client *c = (struct client*)client;
    int sock_fd = c->client_sock_fd;
    int uid = c->uid;
    strcpy(username, c->username);
    username[strcspn(username, "\n")] = '\0';
    std::cout << username << " joined!\n" << std::endl;

    pthread_mutex_lock(&sock_fds_mutex);
    sock_fds.push_back(sock_fd);
    pthread_mutex_unlock(&sock_fds_mutex);

    char r_buff[BUFLEN];
    char private_msg[BUFLEN];

    while(1){
        // Recieve data from a single client
        bzero(r_buff, sizeof(r_buff));
        int n = read(sock_fd, r_buff, sizeof(r_buff));
        if(n<=0){
            std::cout << username << " disconnected!" << std::endl;
            break;
        }

        // Command is received
        if(r_buff[0]=='/'){
            char user_id[10];
            bzero(private_msg, sizeof(private_msg));
            
            char *ptr = r_buff + 1;
            char *temp;
            temp = strtok(ptr, " ");
            int args = 0;
            std::string cmd;
            while(temp!=NULL){
                if(args==0){
                    cmd = std::string(temp);
                    cmd = trim(cmd);
                    args++;
                }else if(args==1){
                    strcpy(user_id, temp);
                    args++;
                }else{
                    strcat(private_msg, temp);
                    strcat(private_msg, " ");
                }
                temp = strtok(NULL, " ");
            }

            if(cmd=="whisper"){
                whisperMsg(atoi(user_id), username, private_msg);
            }else if(cmd=="view"){
                viewAllUsers(sock_fd);
            }else{
                sendInvalidCommandError(sock_fd);
            }
        }else{
            // Message is received
            std::cout << username << "[" << uid << "]" <<  " says: " << r_buff << std::endl;

            // Broadcast data to all clients
            broadcastMsg(r_buff, username, uid);
        }
    }

    pthread_mutex_lock(&sock_fds_mutex);
    sock_fds.erase(std::remove(sock_fds.begin(), sock_fds.end(), sock_fd), sock_fds.end());
    pthread_mutex_unlock(&sock_fds_mutex);
    
    pthread_mutex_lock(&users_mutex);
    users.erase(uid);
    pthread_mutex_unlock(&users_mutex);

    close(sock_fd);
    free(c);
    return NULL;
}


int main(int argc, char **argv){

    int sock_fd, client_fd;
    int yes = 1;
    char username[20];

    srand(time(NULL));

    // Internet socket address
    struct sockaddr_in server, client;
    socklen_t client_len;

    // Create a socket
    if((sock_fd = socket(AF_INET, SOCK_STREAM, 0))==-1){
        std::cout << "Failed to create a socket!" << std::endl;
        exit(1);
    }

    // Define server port, address and family
    server.sin_addr.s_addr = inet_addr(SERVER_IP_ADDRESS);
    server.sin_port = PORT;
    server.sin_family = AF_INET;

    // Set socket options
    if(setsockopt(sock_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes))==-1){
        std::cout << "Failed to set socket options" << std::endl;
        exit(1);
    }

    // Bind socket to server
    if(bind(sock_fd, (struct sockaddr *)&server, sizeof(server))==-1){
        std::cout << "Failed to bind socket to server" << std::endl;
        exit(1);
    }

    // Listen for connections
    listen(sock_fd, 5);
    std::cout << "Server started on port: " << PORT << "\n" << std::endl;

    while(1){
        // Accept any client connections
        client_len = sizeof(client);
        if((client_fd = accept(sock_fd, (struct sockaddr*)&client, &client_len))==-1){
            std::cout << "Failed to accept client!" << std::endl;
            exit(1);
        }

        // Read user details from client
        bzero(username, sizeof(username));
        if(read(client_fd, username, sizeof(username))==-1){
            std::cout << "Failed to read client details!" << std::endl;
        }
        int uid = rand()%100;
        struct client *c = (struct client *)malloc(sizeof(struct client));
        c->client_sock_fd = client_fd;
        c->uid = uid;
        strcpy(c->username, username);
        
        users[uid] = client_fd;
        std::string name = username;
        name = trim(name);
        usernames[name] = uid;

        // Create a new thread for every client
        pthread_t tid;
        if(pthread_create(&tid, NULL, handleClient, (void *)c)!=0){
            std::cout << "Failed to create client thread" << std::endl;
        }

        // After the thread terminates, storage for tid can be reclaimed
        pthread_detach(tid);
    }

    close(sock_fd);

    return 0;
}