server.cpp

#include <iostream>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <cstring>
#include <fstream>
#include <sstream>
#include <thread>
#include <algorithm>

#include "server.h"
#include "HttpRequest.h"
#include "HttpResponse.h"
#include "FileSystem.h"
#include "MimeTypes.h"
#include "router.h"
#include "ThreadPool.h"
#include "Logger.h"
#include "RequestInspector.h"

using namespace std;

#pragma comment(lib, "Ws2_32.lib")

const DWORD KEEP_ALIVE_TIMEOUT = 10000;   // 10 seconds
const int MAX_KEEP_ALIVE_REQUESTS = 100;

void Server::use(std::shared_ptr<Middleware> middleware)
{
    pipeline.use(middleware);
}

void Server::use(const string& prefix, shared_ptr<Middleware> middleware)
{
    pipeline.use(prefix, middleware);
}

Server::Server(int port) : port(port), threadPool(4)
{
    this->port = port;
    serverSocket = INVALID_SOCKET;
    memset(&serverAddress, 0, sizeof(serverAddress));
}

Server::~Server()
{
    Logger::info("Server stopped");

    if (serverSocket != INVALID_SOCKET)
    {
        closesocket(serverSocket);
    }
    if (clientSocket != INVALID_SOCKET)
    {
        closesocket(clientSocket);
    }

    WSACleanup();
}

bool Server::initialize()
{
    // Initialize Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        Logger::error("WSAStartup failed!");
        return false;
    }

    Logger::info("Winsock initialized!");

    // Create TCP Socket
    serverSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

    if (serverSocket == INVALID_SOCKET)
    {
        Logger::error("Socket failed");
        WSACleanup();
        return false;
    }

    Logger::debug("Socket created successfully");

    // Configure Address
    serverAddress.sin_family = AF_INET;
    serverAddress.sin_port = htons(port);
    serverAddress.sin_addr.s_addr = INADDR_ANY;

    // Bind
    if (bind(serverSocket,
             (sockaddr*)&serverAddress,
             sizeof(serverAddress)) == SOCKET_ERROR)
    {
        Logger::error("Bind failed");

        closesocket(serverSocket);
        WSACleanup();

        return false;
    }

    Logger::debug("Bind successful");

    // Listen
    if (listen(serverSocket, SOMAXCONN) == SOCKET_ERROR)
    {
        Logger::error("Listen failed");

        closesocket(serverSocket);
        WSACleanup();

        return false;
    }

    Logger::info("Listening on port " + to_string(port));

    return true;
}

string loadFile(const string& path)
{
    std::ifstream file(path);

    if (!file.is_open())
    {
        return "";
    }

    std::stringstream buffer;
    buffer << file.rdbuf();

    return buffer.str();
}

void Server::get(const string& path, Router::Handler handler)
{
    router.get(path, handler);
}

void Server::post(const string& path, Router::Handler handler)
{
    router.post(path, handler);
}

void Server::put(const string& path, Router::Handler handler)
{
    router.put(path, handler);
}

void Server::patch(const string& path, Router::Handler handler)
{
    router.patch(path, handler);
}

void Server::del(const string& path, Router::Handler handler)
{
    router.del(path, handler);
}

void Server::head(const string& path, Router::Handler handler)
{
    router.head(path, handler);
}

void Server::options(const string& path, Router::Handler handler)
{
    router.options(path, handler);
}

void Server::query(const string& path, Router::Handler handler)
{
    router.query(path, handler);
}

bool sendAll(SOCKET socket, const char* data, size_t length)
{
    size_t total = 0;

    while (total < length)
    {
        int sent = send(
            socket,
            data + total,
            (int)(length - total),
            0
        );

        if (sent == SOCKET_ERROR)
        {
            return false;
        }

        total += sent;
    }

    return true;
}

void Server::handelClient(SOCKET clientSocket, string clientIP)
{
    stringstream threadIdStream;
    threadIdStream << std::this_thread::get_id();
    Logger::debug("Worker Thread started: " + threadIdStream.str());

    char buffer[4096];
    int requestCount = 0;

    while (true)
    {
        requestCount++;

        if (requestCount > MAX_KEEP_ALIVE_REQUESTS)
        {
            Logger::warn("Keep-Alive request limit reached for client " + clientIP);
            break;
        }

        int bytesReceived = recv(
            clientSocket,
            buffer,
            sizeof(buffer) - 1,
            0
        );

        if (bytesReceived == SOCKET_ERROR)
        {
            int error = WSAGetLastError();

            if (error == WSAETIMEDOUT)
            {
                Logger::warn("Keep-Alive timeout for client " + clientIP);
            }
            else
            {
                Logger::error("Receive failed for client " + clientIP + ": " + to_string(error));
            }

            break;
        }

        if (bytesReceived == 0)
        {
            Logger::info("Client closed connection: " + clientIP);
            break;
        }

        buffer[bytesReceived] = '\0';

        HttpRequest request;

        if (!request.parse(buffer))
        {
            Logger::error("Invalid HTTP Request from " + clientIP);
            closesocket(clientSocket);
            return;
        }

        request.setClientIP(clientIP);

        static atomic<uint64_t> globalReqCounter{0};
        uint64_t reqId = ++globalReqCounter;

        stringstream ssMethod;
        ssMethod << request.getMethod();

        RequestInspector::current().start(reqId, clientIP, ssMethod.str(), request.getPath(), request.getVersion(), threadIdStream.str());

        string acceptEncodingHeader = request.getHeader("Accept-Encoding");
        if(!acceptEncodingHeader.empty()) RequestInspector::current().addHeader("Accept-Encoding", acceptEncodingHeader);

        string ifNoneMatchHeader = request.getHeader("If-None-Match");
        if(!ifNoneMatchHeader.empty()) RequestInspector::current().addHeader("If-None-Match", ifNoneMatchHeader);

        string ifModifiedHeader = request.getHeader("If-Modified-Since");
        if(!ifModifiedHeader.empty()) RequestInspector::current().addHeader("If-Modified-Since", ifModifiedHeader);

        string rangeHeaderVal = request.getHeader("Range");
        if(!rangeHeaderVal.empty()) RequestInspector::current().addHeader("Range", rangeHeaderVal);

        string connection = request.getHeader("Connection");
        string range = request.getHeader("Range");

        if (!range.empty())
        {
            Logger::debug("Range requested: " + range);
        }

        bool keepAlive = true;

        // HTTP/1.0 defaults to close
        if (request.getVersion() == "HTTP/1.0")
        {
            keepAlive = (connection == "keep-alive");
        }
        // HTTP/1.1 defaults to keep-alive
        else
        {
            keepAlive = (connection != "close");
        }

        HttpResponse response;

        pipeline.execute(
            request,
            response,
            [&]()
            {
                router.handle(request, response);
            }
        );

        if (keepAlive)
        {
            response.setHeader("Connection", "keep-alive");
        }
        else
        {
            response.setHeader("Connection", "close");
        }

        if (response.isChunked())
        {
            response.setHeader("Transfer-Encoding", "chunked");

            string headers = response.buildHeaders();

            if (!sendAll(clientSocket, headers.c_str(), headers.size()))
            {
                Logger::error("Failed to send chunked headers to " + clientIP);
                break;
            }

            const string& body = response.getBody();
            const size_t CHUNK_SIZE = 1024;
            size_t pos = 0;
            bool sendFailed = false;

            while (pos < body.size())
            {
                size_t len = min(CHUNK_SIZE, body.size() - pos);
                string chunk;
                stringstream ss;

                ss << hex << len;
                chunk += ss.str();
                chunk += "\r\n";
                chunk.append(body, pos, len);
                chunk += "\r\n";

                if (!sendAll(clientSocket, chunk.c_str(), chunk.size()))
                {
                    Logger::error("Failed to send chunk body to " + clientIP);
                    sendFailed = true;
                    break;
                }

                pos += len;
            }

            if (sendFailed) break;

            if (!sendAll(clientSocket, "0\r\n\r\n", 5))
            {
                Logger::error("Failed to send chunk terminal sequence to " + clientIP);
                break;
            }

            Logger::debug("Chunked response sent to " + clientIP);
        }
        else
        {
            string rawResponse = response.build();

            bool ok = sendAll(
                clientSocket,
                rawResponse.c_str(),
                rawResponse.size()
            );

            if (!ok)
            {
                Logger::error("Send failed to " + clientIP + ": " + to_string(WSAGetLastError()));
            }
            else
            {
                Logger::debug("Response sent to " + clientIP + " (" + to_string(rawResponse.size()) + " bytes)");
            }
        }

        if (!keepAlive)
        {
            RequestInspector::current().setResponse(
                                            response.getStatus(), 
                                            response.getHeader("Content-Type"), 
                                            response.getBody().size(), 
                                            "close"
                                        );
            RequestInspector::current().finishAndLog();
            break;
        }
        else
        {
            RequestInspector::current().setResponse(
    response.getStatus(), 
    response.getHeader("Content-Type"), 
    response.getBody().size(), 
    "keep-alive"
);
            RequestInspector::current().finishAndLog();
        }
    }

    closesocket(clientSocket);
    Logger::info("Client disconnected: " + clientIP);
}

void Server::start()
{
    Logger::info("Server started");

    bool running = true;

    while (running)
    {
        sockaddr_in clientAddress;
        int clientSize = sizeof(clientAddress);

        SOCKET clientSocket = accept(
            serverSocket,
            (sockaddr*)&clientAddress,
            &clientSize
        );

        if (clientSocket == INVALID_SOCKET)
        {
            Logger::error("Socket failed");
            continue;
        }

        char* ip = inet_ntoa(clientAddress.sin_addr);
        string clientIP = ip ? ip : "Unknown";

        setsockopt(
            clientSocket,
            SOL_SOCKET,
            SO_RCVTIMEO,
            (const char*)&KEEP_ALIVE_TIMEOUT,
            sizeof(KEEP_ALIVE_TIMEOUT)
        );

        Logger::info("Client connected: " + clientIP);

        threadPool.submit(
            [this, clientSocket, clientIP]()
            {
                handelClient(clientSocket, clientIP);
            }
        );
    }
}