ConnectionHandler.cpp

#include "ConnectionHandler.hpp"
#include "Utils.hpp"
#include "Client.hpp"
#include "Server.hpp"
#include <iostream>
#include <cstring>
#include <unistd.h>
#include <fcntl.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <algorithm>
#include <netinet/tcp.h>

ConnectionHandler::ConnectionHandler() : total_connections(0) {
    debug("[ConnectionHandler]: Initialized");
}

ConnectionHandler::~ConnectionHandler() {
    debug("[ConnectionHandler]: Cleaning up connections");
    
    // Close any remaining client connections and clean up Client objects
    for (std::map<int, Client*>::iterator it = client_objects.begin(); it != client_objects.end(); ++it) {
        debug("[ConnectionHandler]: Closing fd: " + Utils::to_string(it->first));
        // Delete the Client object
        delete it->second;
        // Close the socket
        close(it->first);
    }
    
    // Clear all maps
    client_objects.clear();
    client_addrs.clear();
    client_activity.clear();
    client_ips.clear();
}

int ConnectionHandler::acceptClient(int server_fd, std::vector<pollfd>& poll_fds, Server& server) {
    debug("[ConnectionHandler]: Accepting new client from server fd: " + Utils::to_string(server_fd));
    
    // Accept the new connection
    struct sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);
    int client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &client_len);
    
    if (client_fd < 0) {
        debug("[ConnectionHandler]: Accept failed");
        return -1;
    }
    
    // Configure socket for better performance and stability
    if (!setSocketOptions(client_fd)) {
        debug("[ConnectionHandler]: Failed to set socket options, closing connection");
        close(client_fd);
        return -1;
    }
    
    // Store client address information
    client_addrs[client_fd] = client_addr;
    
    // Update activity timestamp
    updateActivity(client_fd);
    
    // Add to poll_fds for monitoring (initially watch for read events)
    pollfd client_pollfd = {client_fd, POLLIN, 0};
    poll_fds.push_back(client_pollfd);
    
    // Get client IP for stats
    char client_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(client_addr.sin_addr), client_ip, INET_ADDRSTRLEN);
    debug("[ConnectionHandler]: New client connected from IP: " + std::string(client_ip));
    
    // Update connection statistics
    updateStats(std::string(client_ip), true);
    
    // Create a new Client object for this connection
    Client* newClient = new Client(server, client_fd);
    
    // Store the Client object
    client_objects[client_fd] = newClient;
    
    debug("[ConnectionHandler]: Created new Client object for fd: " + Utils::to_string(client_fd));
    
    return client_fd;
}

// In ConnectionHandler::closeClient
bool ConnectionHandler::closeClient(int client_fd) {
    debug("[ConnectionHandler]: Closing client fd: " + Utils::to_string(client_fd));
    
    if (client_fd < 0 || client_addrs.find(client_fd) == client_addrs.end()) {
        debug("[ConnectionHandler]: Invalid client fd: " + Utils::to_string(client_fd));
        return false;
    }
    
    // Get client IP for stats before removing from client_addrs
    std::string client_ip = getClientIP(client_fd);
    
    // Delete the Client object if it exists
    std::map<int, Client*>::iterator client_it = client_objects.find(client_fd);
    if (client_it != client_objects.end()) {
        debug("[ConnectionHandler]: Deleting Client object for fd: " + Utils::to_string(client_fd));
        
        // Delete client safely
        Client* client = client_it->second;
        client_objects.erase(client_it);  // Remove from map BEFORE deleting
        
        if (client) {
            try {
                delete client;  // Only delete after removing from map
            } catch (const std::exception& e) {
                debug("[ConnectionHandler]: Exception deleting client: " + std::string(e.what()));
            }
        }
    }
    
    // Close the socket
    if (fcntl(client_fd, F_GETFD) != -1) {
        close(client_fd);
    }
    
    // Remove from client tracking
    client_addrs.erase(client_fd);
    client_activity.erase(client_fd);
    
    // Update connection statistics
    updateStats(client_ip, false);
    
    return true;
}

bool ConnectionHandler::isValidClient(int fd) const {
    return client_addrs.find(fd) != client_addrs.end();
}

std::string ConnectionHandler::getClientIP(int client_fd) const {
    std::map<int, sockaddr_in>::const_iterator it = client_addrs.find(client_fd);
    if (it == client_addrs.end()) {
        return "unknown";
    }
    
    char client_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(it->second.sin_addr), client_ip, INET_ADDRSTRLEN);
    return std::string(client_ip);
}

void ConnectionHandler::updateActivity(int client_fd) {
    client_activity[client_fd] = time(NULL);
}

time_t ConnectionHandler::getLastActivity(int client_fd) const {
    std::map<int, time_t>::const_iterator it = client_activity.find(client_fd);
    if (it == client_activity.end()) {
        return 0;
    }
    return it->second;
}

// Get Client object for a given file descriptor
Client* ConnectionHandler::getClient(int client_fd) {
    std::map<int, Client*>::iterator it = client_objects.find(client_fd);
    if (it == client_objects.end()) {
        return NULL;
    }
    return it->second;
}

// Add a Client object
void ConnectionHandler::addClient(int client_fd, Client* client) {
    client_objects[client_fd] = client;
}

// Remove a Client object (without deleting it)
void ConnectionHandler::removeClient(int client_fd) {
    client_objects.erase(client_fd);
}

// Update poll events based on client states
void ConnectionHandler::updatePollEvents(std::vector<pollfd>& poll_fds) {
    for (size_t i = 0; i < poll_fds.size(); ++i) {
        int fd = poll_fds[i].fd;
        
        // Skip non-client fds (they may be server sockets)
        if (client_objects.find(fd) == client_objects.end()) {
            continue;
        }
        
        Client* client = client_objects[fd];
        
        // Set events based on client state
        switch (client->getState()) {
            case STATE_SENDING_RESPONSE:
                // Need to check for writability
                poll_fds[i].events = POLLOUT;
                break;
                
            case STATE_READING_REQUEST:
            case STATE_IDLE:
            case STATE_PROCESSING:
                // Need to check for readability
                poll_fds[i].events = POLLIN;
                break;
                
            case STATE_ERROR:
                // No events, will be removed on next iteration
                poll_fds[i].events = 0;
                break;
        }
    }
}

int ConnectionHandler::closeInactiveClients(std::vector<pollfd>& poll_fds, int timeout_seconds) {
    debug("[ConnectionHandler]: Checking for inactive clients, timeout: " + Utils::to_string(timeout_seconds) + "s");
    
    time_t current_time = time(NULL);
    int closed_count = 0;
    
    std::vector<int> to_close;
    
    // Find clients that have been inactive for too long
    for (std::map<int, time_t>::const_iterator it = client_activity.begin(); it != client_activity.end(); ++it) {
        if (current_time - it->second > timeout_seconds) {
            to_close.push_back(it->first);
        }
    }
    
    // Close identified inactive clients
    for (std::vector<int>::const_iterator it = to_close.begin(); it != to_close.end(); ++it) {
        debug("[ConnectionHandler]: Closing inactive client: " + Utils::to_string(*it));
        
        // Remove from poll_fds
        for (std::vector<pollfd>::iterator pfd_it = poll_fds.begin(); pfd_it != poll_fds.end(); ) {
            if (pfd_it->fd == *it) {
                pfd_it = poll_fds.erase(pfd_it);
            } else {
                ++pfd_it;
            }
        }
        
        if (closeClient(*it)) {
            closed_count++;
        }
    }
    
    if (closed_count > 0) {
        debug("[ConnectionHandler]: Closed " + Utils::to_string(closed_count) + " inactive clients");
    }
    
    return closed_count;
}

size_t ConnectionHandler::getActiveClientsCount() const {
    return client_addrs.size();
}

size_t ConnectionHandler::getTotalConnectionsCount() const {
    return total_connections;
}

size_t ConnectionHandler::getUniqueIPsCount() const {
    return client_ips.size();
}

const std::map<std::string, int>& ConnectionHandler::getClientIPs() const {
    return client_ips;
}

void ConnectionHandler::updateStats(const std::string& client_ip, bool connecting) {
    if (connecting) {
        total_connections++;
        client_ips[client_ip]++;
        debug("[ConnectionHandler]: Stats - New connection from " + client_ip + 
              " (total: " + Utils::to_string(total_connections) + 
              ", active: " + Utils::to_string(client_addrs.size()) + ")");
    } else {
        std::map<std::string, int>::iterator it = client_ips.find(client_ip);
        if (it != client_ips.end()) {
            if (it->second > 0) {
                it->second--;
                if (it->second == 0) {
                    client_ips.erase(it);
                }
            }
        }
        debug("[ConnectionHandler]: Stats - Closed connection from " + client_ip + 
              " (total: " + Utils::to_string(total_connections) + 
              ", active: " + Utils::to_string(client_addrs.size()) + ")");
    }
}

bool ConnectionHandler::setSocketOptions(int fd) {
    // Set non-blocking mode
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags < 0 || fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) {
        debug("[ConnectionHandler]: Failed to set non-blocking mode");
        return false;
    }
    
    // Set TCP_NODELAY to disable Nagle's algorithm
    int yes = 1;
    if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &yes, sizeof(yes)) < 0) {
        debug("[ConnectionHandler]: Failed to set TCP_NODELAY" );
        // Not critical, continue
    }
    
    // Set SO_REUSEADDR to allow quick restart
    if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes)) < 0) {
        debug("[ConnectionHandler]: Failed to set SO_REUSEADDR" );
        // Not critical, continue
    }
    
    // Set receive and send timeouts
    struct timeval timeout;
    timeout.tv_sec = 30;
    timeout.tv_usec = 0;
    
    if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0) {
        debug("[ConnectionHandler]: Failed to set SO_RCVTIMEO");
        // Not critical, continue
    }
    
    if (setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)) < 0) {
        debug("[ConnectionHandler]: Failed to set SO_SNDTIMEO" );
        // Not critical, continue
    }
    
    return true;
}