server.cpp

// =============================
// server.cpp
// Build: g++ -std=c++17 -pthread  -lssl -lcrypto server.cpp protocol.cpp utils.cpp tpool.cpp tls_registry.cpp tls.cpp -o server
// Run:   ./server 8888
// =============================

#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <openssl/ssl.h>
#include <unistd.h>
#include <pthread.h>

#include <cstring>
#include <string>
#include <unordered_map>
#include <vector>
#include <algorithm>
#include <memory>

#include "protocol.h"
#include "utils.h"
#include "tpool.h"
#include "tls.h"


#define N_THREADS 100

using namespace std;


void* handle_client(void* arg);
static inline void logout_user_locked(User* user, const string& user_id);
static inline User* get_user_locked(const string& id);
static inline bool is_online(const string& id);
pthread_mutex_t* get_room_mutex(const string& gid);


unordered_map<string, User> users;
pthread_mutex_t users_mtx = PTHREAD_MUTEX_INITIALIZER;
unordered_map<string, vector<string>> chatrooms;
pthread_mutex_t chatrooms_mtx = PTHREAD_MUTEX_INITIALIZER;
unordered_map<string, unique_ptr<pthread_mutex_t>> chatroom_flows; // mutex for each chatroom in "chatrooms"
pthread_mutex_t chatroom_flows_mtx = PTHREAD_MUTEX_INITIALIZER;
static SSL_CTX* g_tls_server_ctx = nullptr;
static SSL_CTX* g_tls_client_ctx = nullptr;


int main(int argc, char **argv) {
    uint16_t listen_port = 8888;
    if (argc >= 2) {
        try {
            int p = stoi(argv[1]);
            if (p < 0 || p > 65535) // we allow port==0 here
                ERR_EXIT("invalid port");
            listen_port = static_cast<uint16_t>(p);
        } catch (...) {
            ERR_EXIT("invalid port");
        }
    }

    int listen_fd = init_server(listen_port); 
    sockaddr_in cli_addr{};
    socklen_t cli_len = sizeof(cli_addr);

    /* create thread pool */
    tpool *pool = tpool_init(N_THREADS);

    /* tls things */
    g_tls_server_ctx = tls_make_server_ctx("certs/server.cert.pem", "certs/server.key.pem");
    if (!g_tls_server_ctx) ERR_EXIT("tls server ctx");

    g_tls_client_ctx = tls_make_client_ctx_noverify();
    if (!g_tls_client_ctx) ERR_EXIT("tls client ctx");
    
    while (true) {
        int conn_fd = accept(listen_fd, reinterpret_cast<sockaddr*>(&cli_addr), &cli_len);
        if (conn_fd < 0) {
            if (errno == EINTR) continue;  // try again
            if (errno == ENFILE) {
                fprintf(stderr, "out of file descriptor table ...");
                continue;
            }
            ERR_THROW("accept");
        }

        int *arg = new int(conn_fd);
        tpool_add(pool, handle_client, static_cast<void*>(arg));
    }
    
    // tpool_wait(pool);
    tpool_destroy(pool);
    
    /* tls things */
    if (g_tls_server_ctx) SSL_CTX_free(g_tls_server_ctx);
    if (g_tls_client_ctx) SSL_CTX_free(g_tls_client_ctx);

    tls_cleanup_fd(listen_fd);
    close(listen_fd);
}










void* handle_client(void* arg) {
    int fd = *static_cast<int*>(arg);
    delete static_cast<int*>(arg);

    /* tls things */
    if (!tls_accept_fd(g_tls_server_ctx, fd)) {
        tls_cleanup_fd(fd);
        close(fd);
        return nullptr;
    }

    string client_ip = peer_ip(fd);
    fprintf(stderr, "Client %.80s connected\n", client_ip.c_str());

    Cmd cmd;
    void *req_body = nullptr;
    string current_id = ""; // TODO: remove the use of current_id

    while (recv_request(fd, cmd, req_body)) {
        switch (cmd){
            case CMD_REGISTER: {
                auto in = static_cast<RegisterReq*>(req_body);
                string id(in->id, strnlen(in->id, ID_LEN));
                string pw(in->password, strnlen(in->password, PASSWORD_LEN));

                Status stat = STATUS_ERROR;
                pthread_mutex_lock(&users_mtx);
                if(get_user_locked(id)==nullptr){
                    // new user
                    users[id] = User{.password=pw};
                    stat = STATUS_OK;
                }
                pthread_mutex_unlock(&users_mtx);
                if(!send_response(fd, cmd, stat, nullptr, 0)) goto disconnect;
                if(stat==STATUS_OK)
                    fprintf(stderr, "Registered user %s (password: %s)\n", id.c_str(), pw.c_str());
                break;
            }

            case CMD_LOGIN: {
                if (!current_id.empty()) { // already logged in
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }

                auto in = static_cast<LoginReq*>(req_body);
                string id(in->id, strnlen(in->id, ID_LEN));
                string pw(in->password, strnlen(in->password, PASSWORD_LEN));
                uint16_t port = ntohs(in->port); // 0 <= port <= 65535
                uint16_t notify_port = ntohs(in->notify_port); // 0 <= port <= 65535
                
                pthread_mutex_lock(&users_mtx);
                User* user = get_user_locked(id);
                
                if (user == nullptr           // unknown id
                    || user->password != pw   // wrong password
                    || user->online){         // user is alreaedy online
                    pthread_mutex_unlock(&users_mtx);
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }

                bool addr_unique = true;
                for (const auto& [uid,u] : users)
                    if (u.online && u.ip == client_ip && u.port == port) {
                        addr_unique = false;
                        break;
                    }
                if(!addr_unique){
                    pthread_mutex_unlock(&users_mtx);
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }

                user->ip = client_ip;
                user->port = port;
                user->online = true;
                user->notify_fd = -1;
                current_id = id;
                pthread_mutex_unlock(&users_mtx);
                
                if(!send_response(fd, cmd, STATUS_OK, nullptr, 0)) goto disconnect;
                
                int notify_fd = init_client(client_ip, notify_port); 

                /* tls things */
                if (!tls_connect_fd(g_tls_client_ctx, notify_fd)){
                    // tls failed
                    tls_cleanup_fd(notify_fd);
                    close(notify_fd);
                    fprintf(stderr, "Server fails to connect to User %s via TLS\n", id.c_str());
                    
                    pthread_mutex_lock(&users_mtx);
                    logout_user_locked(get_user_locked(id), id);
                    pthread_mutex_unlock(&users_mtx);
                    break;
                }

                // save notify_fd
                pthread_mutex_lock(&users_mtx);
                User* u2 = get_user_locked(id);
                if (u2 && u2->online) u2->notify_fd = notify_fd;
                pthread_mutex_unlock(&users_mtx);

                fprintf(stderr, "User %s logged in (port: %d)\n", id.c_str(), port);
                break;
            }

            case CMD_LOGOUT: {
                auto in = static_cast<LogoutReq*>(req_body);
                string id(in->id, strnlen(in->id, ID_LEN));

                if (current_id.empty() || // not logged in yet
                    current_id != id) {   // logout id != current user id
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }

                pthread_mutex_lock(&users_mtx);
                User* user = get_user_locked(id);
                if(user == nullptr || !user->online){
                    pthread_mutex_unlock(&users_mtx);
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }
                logout_user_locked(user, id);
                pthread_mutex_unlock(&users_mtx);
                
                current_id = "";
                if(!send_response(fd, cmd, STATUS_OK, nullptr, 0)) goto disconnect;
                break;
            }

            case CMD_LIST: {
                vector<UserInfo> online;
                
                pthread_mutex_lock(&users_mtx);
                for (const auto& [uid,u] : users){
                    if (!u.online) continue;

                    UserInfo ui{};
                    strncpy(ui.id, uid.c_str(), ID_LEN);
                    strncpy(ui.ip, u.ip.c_str(), INET_ADDRSTRLEN);
                    ui.port = htons(u.port);
                    online.push_back(ui);
                }
                pthread_mutex_unlock(&users_mtx);
               

                uint32_t cnt = static_cast<uint32_t>(online.size());
                size_t body_len = sizeof(cnt) + online.size()*sizeof(UserInfo);
                ListRes* res = static_cast<ListRes*>(malloc(body_len));
                if (res == nullptr) ERR_THROW("malloc");

                res->count = htonl(cnt); // host order -> network order
                if (cnt)
                    memcpy(res->user_list, online.data(), online.size()*sizeof(UserInfo));

                if(!send_response(fd, cmd, STATUS_OK, static_cast<const void *>(res), body_len)) goto disconnect;
                free(res);
                break;
            }

            case CMD_JOIN: {
                auto in = static_cast<JoinReq*>(req_body);
                string id(in->id, strnlen(in->id, ID_LEN));
                string gid(in->gid, strnlen(in->gid, ID_LEN));

                if (current_id.empty() || current_id != id || !is_online(id)) {
                   if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }

                pthread_mutex_lock(&chatrooms_mtx);
                if (chatrooms.count(gid) != 0 && find(chatrooms[gid].begin(), chatrooms[gid].end(), id) != chatrooms[gid].end()) {
                    // user has already joined the requested chatroom
                    pthread_mutex_unlock(&chatrooms_mtx);
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;  
                }
                chatrooms[gid].push_back(id); // new chatroom is created if key does not exist
                pthread_mutex_unlock(&chatrooms_mtx);
                if(!send_response(fd, cmd, STATUS_OK, nullptr, 0)) goto disconnect;
                break; 
            }

            case CMD_LEAVE: {
                auto in = static_cast<LeaveReq*>(req_body);
                string id(in->id, strnlen(in->id, ID_LEN));
                string gid(in->gid, strnlen(in->gid, ID_LEN));

                if (current_id.empty() || current_id != id || !is_online(id)) {
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }

                pthread_mutex_lock(&chatrooms_mtx);
                auto chatroom_it = chatrooms.find(gid);
                if (chatroom_it == chatrooms.end()){
                    pthread_mutex_unlock(&chatrooms_mtx);
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }
                auto& roommates = chatroom_it->second;
                auto roommate_it = find(roommates.begin(), roommates.end(), id);
                if (roommate_it == roommates.end()) {
                    // user is not in the requested chatroom
                    pthread_mutex_unlock(&chatrooms_mtx);
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break; 
                }
                roommates.erase(roommate_it);

                // remove empty chatroom
                if (roommates.empty()) {
                    chatrooms.erase(chatroom_it);
                }

                pthread_mutex_unlock(&chatrooms_mtx);
                if(!send_response(fd, cmd, STATUS_OK, nullptr, 0)) goto disconnect;
                break; 
            }

            case CMD_GROUP: {
                auto in = static_cast<GroupReq*>(req_body);
                string id(in->id, strnlen(in->id, ID_LEN));
                string gid(in->gid, strnlen(in->gid, ID_LEN));

                if (current_id.empty() || current_id != id || !is_online(id)) {
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }
                
                pthread_mutex_lock(&chatrooms_mtx);
                auto chatroom_it = chatrooms.find(gid);
                if (chatroom_it == chatrooms.end()) {
                    // requested chatroom does not exist
                    pthread_mutex_unlock(&chatrooms_mtx);
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break;
                }
                auto& members = chatroom_it->second;
                if (find(members.begin(), members.end(), id) == members.end()) {
                    // user is not in the requested chatroom
                    pthread_mutex_unlock(&chatrooms_mtx);
                    if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                    break; 
                }
                vector<string> roommates = members;

                pthread_mutex_unlock(&chatrooms_mtx);
                pthread_mutex_t* room_mutex = get_room_mutex(gid);
                pthread_mutex_lock(room_mutex);
                vector<int> relay_fds;
                relay_fds.reserve(roommates.size());

                pthread_mutex_lock(&users_mtx);
                for (const auto& s : roommates) {
                    if (s == id) continue;
                    User* u = get_user_locked(s);
                    if (u==nullptr || !u->online) continue;
                    relay_fds.push_back(u->notify_fd);
                }
                pthread_mutex_unlock(&users_mtx);

                for (int relay_fd : relay_fds)
                    if (!send_request(relay_fd, CMD_GROUP, static_cast<const void*>(in), sizeof(*in))){
                        // ignore failed notifications
                    }

                pthread_mutex_unlock(room_mutex);
                if(!send_response(fd, cmd, STATUS_OK, nullptr, 0)) goto disconnect;
                break;
            }
            default: {
                if(!send_response(fd, cmd, STATUS_ERROR, nullptr, 0)) goto disconnect;
                break;
            }   
        }
        free(req_body);
        req_body = nullptr;
    }

disconnect:
    if (req_body != nullptr){ 
        free(req_body);
        req_body = nullptr;
    }
    tls_cleanup_fd(fd);
    close(fd);
    
    /* force the user to logout */
    pthread_mutex_lock(&users_mtx);
    logout_user_locked(get_user_locked(current_id), current_id);
    pthread_mutex_unlock(&users_mtx);
    
    fprintf(stderr, "Client %.80s (id: %s) disconnected\n", client_ip.c_str(), current_id.c_str());
    return nullptr;
}

static inline void logout_user_locked(User* user, const string& user_id){
    if(user==nullptr || !user->online) return;
    user->ip = "";
    user->port = 0;
    user->online = false;
    if (user->notify_fd >= 0){
        tls_cleanup_fd(user->notify_fd);
        close(user->notify_fd);
        user->notify_fd = -1;
    }
    fprintf(stderr, "User %s logged out\n", user_id.c_str());
}

static inline User* get_user_locked(const string& id) {
    // WARNING: user_mtx must be locked first
    auto it = users.find(id);
    if (it == users.end()) return nullptr;
    return &it->second;
}

static inline bool is_online(const string& id){
    bool online;
    pthread_mutex_lock(&users_mtx);
    User* user = get_user_locked(id);
    online = user!=nullptr && user->online;
    pthread_mutex_unlock(&users_mtx);
    return online;
}

pthread_mutex_t* get_room_mutex(const string& gid) {
    pthread_mutex_lock(&chatroom_flows_mtx);
    auto& ptr = chatroom_flows[gid];
    if (!ptr) {
        ptr = make_unique<pthread_mutex_t>();
        pthread_mutex_init(ptr.get(), nullptr);
    }
    auto* m = ptr.get();
    pthread_mutex_unlock(&chatroom_flows_mtx);
    return m;
}