simple-object-server/src/server.cpp

#include <filesystem>
#include <iostream>
#include <sstream>
#include <fstream>
#include <random>     // For random number generation
#include <chrono>     // For seeding random number generator
#include <vector> // For getAllKeys
#include <string_view> // For litecask values
#include <stdexcept> // For std::runtime_error
#include <sstream>

#include "server.hpp"
#include "hash.hpp"
#include "compress.hpp"
#include "string_utils.hpp" // Include the new utility header
#include "put_handler.hpp"
#include "utils.hpp"
#include "welcome_page.hpp"

namespace simple_object_storage {


bool Server::init_db() {
    try {
        std::filesystem::path db_path = config_.object_store_path / "index.db";
        db_ = std::make_unique<Database>(db_path);
        return true;
    } catch (const std::runtime_error& e) {
        std::cerr << "Database initialization error: " << e.what() << std::endl;
        return false;
    }
}

bool Server::validate_write_request(const httplib::Request &req, httplib::Response &res, const std::vector<std::string> &required_params, std::map<std::string, std::string> &params)
{
    std::string client_ip = req.remote_addr;

    // Check if the client is already over the limit (do NOT increment)
    if (auth_rate_limiter_->is_over_limit(client_ip)) {
        res.status = 429;
        nlohmann::json response = {{"result", "error"}, {"error", "Too many authentication attempts. Please try again later."}};
        res.set_content(response.dump(), "application/json");
        return false;
    }

    // Get token from Authorization header
    std::string token;
    if (req.has_header("Authorization")) {
        const auto& auth_header = req.get_header_value("Authorization");
        if (auth_header.substr(0, 7) == "Bearer ") {
            token = auth_header.substr(7);
        }
    }

    if (token.empty()) {
        // Only count failed attempt (increment the limiter)
        auth_rate_limiter_->is_allowed(client_ip); // This will increment the count
        res.status = 401;
        nlohmann::json response = {{"result", "error"}, {"error", "Missing or invalid Authorization header"}};
        res.set_content(response.dump(), "application/json");
        return false;
    }

    bool write_token_valid = std::find(config_.write_tokens.begin(), config_.write_tokens.end(), token) != config_.write_tokens.end();
    if (!write_token_valid) {
        // Only count failed attempt (increment the limiter)
        auth_rate_limiter_->is_allowed(client_ip); // This will increment the count
        res.status = 403;
        nlohmann::json response = {{"result", "error"}, {"error", "Invalid write token"}};
        res.set_content(response.dump(), "application/json");
        return false;
    }

    // If authentication is successful, do not increment rate limiter

    for (const auto& param : req.params) {
        params[param.first] = param.second;
    }

    for (const auto& param : required_params) {
        if (!req.has_param(param)) {
            res.status = 400;
            nlohmann::json response = {{"result", "error"}, {"error", "Missing required query parameter: " + param}};
            res.set_content(response.dump(), "application/json");
            return false;
        }
    }

    return true;
}

Server::Server(const ServerConfig& config)
    : config_(config), running_(false) {

    if (!std::filesystem::exists(config_.object_store_path)) {
        std::cerr << "Object store directory does not exist: " << config_.object_store_path << std::endl;
        return;
    }
    
    // Initialize the database
    if (!init_db()) {
        // Error already printed in init_db
        // Consider throwing or setting an error state
    }

    // Initialize the put handler
    put_handler_ = std::make_unique<PutHandler>(*this);

    // Initialize rate limiter
    auth_rate_limiter_ = std::make_unique<RateLimiter>(
        config_.auth_rate_limit,
        std::chrono::seconds(config_.auth_window_seconds)
    );
}

Server::~Server() {
    stop();
}

bool Server::start() {
    if (!db_) { // Check if DB initialization failed
        std::cerr << "Database is not initialized. Cannot start server." << std::endl;
        return false;
    }
    setup_routes();

    running_ = true;
    if (!server_.listen(config_.host.c_str(), config_.port)) {
        running_ = false;
        std::cerr << "Failed to listen on " << config_.host << ":" << config_.port << std::endl;
        return false;
    }
    return true; // Should not be reached if listen is blocking and successful
}

void Server::stop() {
    if (running_) {
        server_.stop();
        running_ = false;
        std::cout << "Server stopped." << std::endl;
    }
}

void Server::setup_routes() {
    // Add CORS preflight handler for all routes
    server_.Options(".*", [this](const httplib::Request& req, httplib::Response& res) {
        handle_cors_preflight(req, res);
    });

    // Add CORS headers to all responses
    server_.set_post_routing_handler([this](const httplib::Request& req, httplib::Response& res) {
        add_cors_headers(req, res);
    });

    
    server_.Get("/index.html", [](const httplib::Request&, httplib::Response& res) {
        res.set_content(welcome_page(), "text/html");
    });

    // Get hash for label:tag
    server_.Get("/hash/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        handle_get_hash(req, res);
    });

    // Get version for label:tag
    server_.Get("/version/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        handle_get_version(req, res);
    });


    // Check if object exists by hash or label:tag
    server_.Get("/exists/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        handle_exists(req, res);
    });

    // Get directory listing
    server_.Get("/dir", [this](const httplib::Request& req, httplib::Response& res) {
        handle_get_directory(req, res);
    });

    // Upload object with streaming support
    server_.Put("/upload", [this](const httplib::Request& req, httplib::Response& res) {
        put_handler_->handle_put_object(req, res);
    });

    // Handle PUT requests to other paths
    server_.Put("/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        res.status = 404;
        nlohmann::json response = {{"result", "error"}, {"error", "Not found - put requests must be to /upload"}};
        res.set_content(response.dump(), "application/json");
    });

    // Get metadata for label:tag
    server_.Get("/meta/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        handle_get_metadata(req, res);
    });

    // Delete an object (and all tags on that object)
    server_.Get("/deleteobject", [this](const httplib::Request& req, httplib::Response& res) {
        handle_delete_object(req, res);
    });

    server_.Get("/status", [this](const httplib::Request& req, httplib::Response& res) {
        res.set_content(nlohmann::json({{"result", "success"}, {"status", "ok"}}).dump(), "application/json");
    });

    // Get object by hash or label:tag
    server_.Get("/object/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        handle_get_object(req, res);
    });

    // Welcome page and download object.
    server_.Get("/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        if (req.path == "/") {
            res.set_content(welcome_page(), "text/html");
            return;
        }
        // if the path is not /, then it's a hash or label:tag
        handle_get_object(req, res);
    });

}

void Server::handle_cors_preflight(const httplib::Request& req, httplib::Response& res) {
    add_cors_headers(req, res);
    res.status = 204;  // No content
}

void Server::add_cors_headers(const httplib::Request& req, httplib::Response& res) {
    // Get the origin from the request
    std::string origin = req.get_header_value("Origin");
    
    // If no origin header, no CORS headers needed
    if (origin.empty()) {
        return;
    }

    // Check if origin is allowed
    bool origin_allowed = false;
    if (config_.allowed_origins.empty() || 
        std::find(config_.allowed_origins.begin(), config_.allowed_origins.end(), "*") != config_.allowed_origins.end()) {
        origin_allowed = true;
    } else {
        origin_allowed = std::find(config_.allowed_origins.begin(), config_.allowed_origins.end(), origin) != config_.allowed_origins.end();
    }

    if (origin_allowed) {
        res.set_header("Access-Control-Allow-Origin", origin);
        
        // Add other CORS headers
        std::string methods = join(config_.allowed_methods, ", ");
        res.set_header("Access-Control-Allow-Methods", methods);
        
        std::string headers = join(config_.allowed_headers, ", ");
        res.set_header("Access-Control-Allow-Headers", headers);
        
        if (config_.allow_credentials) {
            res.set_header("Access-Control-Allow-Credentials", "true");
        }
        
        // Add max age for preflight requests
        res.set_header("Access-Control-Max-Age", "86400");  // 24 hours
    }
}

std::string Server::join(const std::vector<std::string>& strings, const std::string& delimiter) {
    if (strings.empty()) return "";
    
    std::string result = strings[0];
    for (size_t i = 1; i < strings.size(); ++i) {
        result += delimiter + strings[i];
    }
    return result;
}

void Server::handle_get_object(const httplib::Request& req, httplib::Response& res) {
    const auto& key = req.matches[1].str();
    std::string hash_str = key;

    // first check if the key matches.
    dbEntry entry;
    if (!db_->get(key, entry)) {
        res.status = 404;
        nlohmann::json response = {{"result", "error"}, {"error", "Couldn't find: " + key}};
        res.set_content(response.dump(), "application/json");
        return;
    }

    // Construct the file path using the hash string
    std::filesystem::path file_path = config_.object_store_path / entry.hash;
    if (!std::filesystem::exists(file_path) || !std::filesystem::is_regular_file(file_path)) {
        res.status = 404;
        nlohmann::json response = {{"result", "error"}, {"error", "Hash recognised, but object missing for: " + entry.hash}};
        res.set_content(response.dump(), "application/json");
        return;
    }

    // Send file using Response::set_file_content
    std::string content_type = "application/octet-stream"; // Basic default
    res.set_file_content(file_path.string(), content_type);
    // No JSON response for file content
}

void Server::handle_get_hash(const httplib::Request& req, httplib::Response& res) {
    const auto& labeltag = req.matches[1].str();

    dbEntry entry;
    if (!db_->get(labeltag, entry)) {
        res.status = 404;
        nlohmann::json response = {{"result", "error"}, {"error", "Label:tag not found"}};
        res.set_content(response.dump(), "application/json");
        return;
    }

    nlohmann::json response = {{"result", "success"}, {"hash", entry.hash}};
    res.set_content(response.dump(), "application/json");
}

void Server::handle_get_directory(const httplib::Request& /*req*/, httplib::Response& res) {
    std::vector<dbEntry> entries;
    
    if (!db_->list(entries)) {
        res.status = 500;
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to retrieve directory listing"}};
        res.set_content(response.dump(), "application/json");
        return;
    }
    
    // output the labeltags as an array of label:tag pairs
    nlohmann::json entries_array = nlohmann::json::array();
    for (const auto& entry : entries) {
        entries_array.push_back({{"labeltags", entry.labeltags}, {"hash", entry.hash}});
    }

    nlohmann::json response = {{"result", "success"}, {"entries", entries_array}};
    res.set_content(response.dump(), "application/json");
}

void Server::handle_get_metadata(const httplib::Request& req, httplib::Response& res) {
    const auto& key = req.matches[1].str();
    std::string hash_str = key;

    // Check if the key is a label:tag format
    dbEntry entry;
    nlohmann::json response;

    if (db_->get(key, entry)) {
        // Got it from label:tag, use the hash
        try {
            response = {{"result", "success"}, {"metadata", entry.metadata}};
        } catch (const nlohmann::json::exception& e) {
            std::cerr << "Error serializing metadata for hash " << hash_str << ": " << e.what() << std::endl;
            res.status = 500;
            response = {{"result", "error"}, {"error", "Internal server error: Failed to serialize metadata"}};
        }
    }
    else
    {
        res.status = 404;
        response = {{"result", "error"}, {"error", "Invalid hash: " + hash_str}};
    }
    res.set_content(response.dump(), "application/json");
}

std::pair<std::string, std::string> Server::parse_labeltag(const std::string& labeltag) const {
    size_t colon_pos = labeltag.find(':');
    if (colon_pos == std::string::npos || colon_pos == 0 || colon_pos == labeltag.length() - 1) {
        return {"", ""};
    }
    return {labeltag.substr(0, colon_pos), labeltag.substr(colon_pos + 1)};
}

void Server::handle_delete_object(const httplib::Request& req, httplib::Response& res) {
    dbEntry entry;

    {
        std::map<std::string, std::string> params;
        if (!validate_write_request(req, res, {"hash"}, params)) {
            return;
        }

        if (!db_->get(params["hash"], entry)) {
            res.status = 404;
            nlohmann::json response = {{"result", "error"}, {"error", "Object not found for: " + params["hash"]}};
            res.set_content(response.dump(), "application/json");
            return;
        }
    } // we only use sanitised data from here on out.

    // Remove all tags that reference this hash
    if (!db_->remove_by_hash(entry.hash)) {
        res.status = 500;
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to remove some or all associated tags"}};
        res.set_content(response.dump(), "application/json");
        return;
    }

    // Remove the file
    std::filesystem::path file_path = config_.object_store_path / entry.hash;
    if (std::filesystem::exists(file_path) && std::filesystem::is_regular_file(file_path)) {
        try {
            std::filesystem::remove(file_path);
        } catch (const std::filesystem::filesystem_error& e) {
            std::cerr << "Error deleting object file: " << e.what() << std::endl;
            res.status = 500;
            nlohmann::json response = {{"result", "error"}, {"error", "Failed to delete object file: " + std::string(e.what())}};
            res.set_content(response.dump(), "application/json");
            return;
        }
    }

    nlohmann::json response = {{"result", "success"}};
    res.set_content(response.dump(), "application/json");
}

void Server::handle_exists(const httplib::Request& req, httplib::Response& res) {
    const auto& key = req.matches[1].str();
    nlohmann::json response;
    dbEntry entry;
    if (db_->get(key, entry)) {
        response = {{"result", "success"}, {"exists", true}};
    } else {
        response = {{"result", "success"}, {"exists", false}};
    }
    res.set_content(response.dump(), "application/json");
}

void Server::handle_get_version(const httplib::Request& req, httplib::Response& res) {
    const auto& key = req.matches[1].str();
    std::string hash_str = key;

    // Check if the key is a label:tag format
    dbEntry entry;
    if (!db_->get(key, entry)) {
        nlohmann::json response = {{"result", "failed"}, {"error", "Failed to get version for: " + key}};
        res.set_content(response.dump(), "application/json");
        return;
    }

    nlohmann::json response;
    if (entry.metadata.contains("version")) {
        response = {{"result", "success"}, {"version", entry.metadata["version"]}};
    } else {
        response = {{"result", "failed"}, {"error", "No version found for: " + key}};
    }
    res.set_content(response.dump(), "application/json");
}

} // namespace simple_object_storage