streamable_http_server.cpp

#include "fastmcpp/server/streamable_http_server.hpp"

#include "fastmcpp/exceptions.hpp"
#include "fastmcpp/util/json.hpp"

#include <algorithm>
#include <cctype>
#include <chrono>
#include <httplib.h>
#include <iomanip>
#include <random>
#include <sstream>

namespace fastmcpp::server
{

StreamableHttpServerWrapper::StreamableHttpServerWrapper(
    McpHandler handler, std::string host, int port, std::string mcp_path, std::string auth_token,
    std::string cors_origin, std::unordered_map<std::string, std::string> response_headers)
    : handler_(std::move(handler)), host_(std::move(host)), requested_port_(port),
      mcp_path_(std::move(mcp_path)), auth_token_(std::move(auth_token)),
      response_headers_(std::move(response_headers))
{
    if (!cors_origin.empty() &&
        response_headers_.find("Access-Control-Allow-Origin") == response_headers_.end())
        response_headers_["Access-Control-Allow-Origin"] = std::move(cors_origin);

    for (const auto& [name, value] : response_headers_)
    {
        std::string lower_name = name;
        std::transform(lower_name.begin(), lower_name.end(), lower_name.begin(),
                       [](unsigned char c) { return static_cast<char>(std::tolower(c)); });

        if (lower_name == "content-type")
            throw std::invalid_argument("response_headers must not override '" + name + "'");
    }
}

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

bool StreamableHttpServerWrapper::check_auth(const std::string& auth_header) const
{
    // If no auth token configured, allow all requests
    if (auth_token_.empty())
        return true;

    // Check for "Bearer <token>" format
    if (auth_header.find("Bearer ") != 0)
        return false;

    std::string provided_token = auth_header.substr(7); // Skip "Bearer "
    return provided_token == auth_token_;
}

void StreamableHttpServerWrapper::apply_additional_response_headers(httplib::Response& res) const
{
    for (const auto& [name, value] : response_headers_)
        res.set_header(name, value);
}

std::string StreamableHttpServerWrapper::generate_session_id()
{
    // Generate cryptographically secure random session ID (128 bits = 32 hex chars)
    std::random_device rd;
    std::mt19937_64 gen(rd());
    std::uniform_int_distribution<uint64_t> dis;

    uint64_t high = dis(gen);
    uint64_t low = dis(gen);

    std::ostringstream oss;
    oss << std::hex << std::setfill('0') << std::setw(16) << high << std::setw(16) << low;
    return oss.str();
}

void StreamableHttpServerWrapper::run_server()
{
    if (requested_port_ == 0) // Request any available port from the operating system.
    {
        const int bound_port = svr_->bind_to_any_port(host_.c_str());
        if (bound_port != -1) // Returns -1 if some error occured.
        {
            bound_port_.store(bound_port);
            svr_->listen_after_bind();
        }
    }
    else
    {
        const bool success = svr_->bind_to_port(host_.c_str(), requested_port_);
        if (success)
        {
            bound_port_.store(requested_port_);
            svr_->listen_after_bind();
        }
    }
    running_ = false;
}

bool StreamableHttpServerWrapper::start()
{
    if (running_)
        return false;

    bound_port_.store(0); // Reset the bound port's value.
    svr_ = std::make_unique<httplib::Server>();

    // Security: Set payload and timeout limits to prevent DoS
    svr_->set_payload_max_length(10 * 1024 * 1024); // 10MB max payload
    svr_->set_read_timeout(30, 0);                  // 30 second read timeout
    svr_->set_write_timeout(30, 0);                 // 30 second write timeout

    // Handle OPTIONS for CORS preflight
    svr_->Options(mcp_path_,
                  [this](const httplib::Request&, httplib::Response& res)
                  {
                      res.set_header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS");
                      res.set_header("Access-Control-Allow-Headers",
                                     "Content-Type, Authorization, Mcp-Session-Id");
                      apply_additional_response_headers(res);
                      res.status = 204;
                  });

    // Set up MCP endpoint (POST)
    svr_->Post(
        mcp_path_,
        [this](const httplib::Request& req, httplib::Response& res)
        {
            // Apply CORS / additional headers up-front so they are present on every
            // response, including early returns (401, 503, 400, 404) and any exception
            // propagated to the catch handlers below.
            apply_additional_response_headers(res);

            // Expose response headers that cross-origin JS clients legitimately need to
            // read. Without this, browsers hide Mcp-Session-Id from response.headers.get()
            // even though it is sent on the wire, because browsers only expose a small
            // whitelist of "safe" response headers by default.
            res.set_header("Access-Control-Expose-Headers", "Mcp-Session-Id");

            try
            {
                // Security: Check authentication if configured
                if (!auth_token_.empty())
                {
                    auto auth_it = req.headers.find("Authorization");
                    if (auth_it == req.headers.end() || !check_auth(auth_it->second))
                    {
                        res.status = 401;
                        res.set_content("{\"error\":\"Unauthorized\"}", "application/json");
                        return;
                    }
                }

                // Parse JSON-RPC message
                auto message = fastmcpp::util::json::parse(req.body);

                // Check for Mcp-Session-Id header
                std::string session_id;
                auto session_it = req.headers.find("Mcp-Session-Id");
                if (session_it != req.headers.end())
                    session_id = session_it->second;

                // Handle initialize request - creates new session
                bool is_initialize = message.contains("method") &&
                                     message["method"].get<std::string>() == "initialize";

                if (is_initialize)
                {
                    // Security: Check session limit before creating new session
                    {
                        std::lock_guard<std::mutex> lock(sessions_mutex_);
                        if (sessions_.size() >= MAX_SESSIONS)
                        {
                            res.status = 503; // Service Unavailable
                            res.set_content("{\"error\":\"Maximum sessions reached\"}",
                                            "application/json");
                            return;
                        }
                    }

                    // Generate new session ID
                    session_id = generate_session_id();

                    // Create ServerSession for this session
                    // Note: For streamable HTTP, responses go back in HTTP response,
                    // so the send callback is not used for normal responses.
                    // It could be used for server-initiated requests in the future.
                    auto session = std::make_shared<ServerSession>(session_id, nullptr);

                    {
                        std::lock_guard<std::mutex> lock(sessions_mutex_);
                        sessions_[session_id] = session;
                    }
                }
                else if (session_id.empty())
                {
                    // Non-initialize request without session ID
                    res.status = 400;
                    res.set_content("{\"error\":\"Mcp-Session-Id header required\"}",
                                    "application/json");
                    return;
                }
                else
                {
                    // Verify session exists
                    std::lock_guard<std::mutex> lock(sessions_mutex_);
                    if (sessions_.find(session_id) == sessions_.end())
                    {
                        res.status = 404;
                        res.set_content("{\"error\":\"Invalid or expired session\"}",
                                        "application/json");
                        return;
                    }
                }

                // Inject session_id into request meta for handler access
                if (!message.contains("params"))
                    message["params"] = Json::object();
                if (!message["params"].contains("_meta"))
                    message["params"]["_meta"] = Json::object();
                message["params"]["_meta"]["session_id"] = session_id;

                // Check if this is a response to a server-initiated request
                if (ServerSession::is_response(message))
                {
                    // Get the session and route the response
                    std::shared_ptr<ServerSession> session;
                    {
                        std::lock_guard<std::mutex> lock(sessions_mutex_);
                        auto it = sessions_.find(session_id);
                        if (it != sessions_.end())
                            session = it->second;
                    }

                    if (session)
                    {
                        bool handled = session->handle_response(message);
                        if (handled)
                        {
                            res.set_header("Mcp-Session-Id", session_id);
                            res.set_content("{\"status\":\"ok\"}", "application/json");
                            res.status = 200;
                            return;
                        }
                    }

                    // Response not handled (unknown request ID)
                    res.status = 400;
                    res.set_content("{\"error\":\"Unknown response ID\"}", "application/json");
                    return;
                }

                // Check if this is a notification (no "id" field means notification)
                // JSON-RPC 2.0 spec: server MUST NOT reply to notifications
                bool is_notification = !message.contains("id") || message["id"].is_null();

                if (is_notification)
                {
                    // For notifications, call handler but don't send response body
                    // This is required by JSON-RPC 2.0 spec and MCP protocol
                    try
                    {
                        handler_(message); // Process but ignore result
                    }
                    catch (...)
                    {
                        // Silently ignore errors for notifications
                    }
                    res.set_header("Mcp-Session-Id", session_id);
                    res.status = 202; // Accepted, no content
                    return;
                }

                // Normal request - process with handler
                auto response = handler_(message);

                // Set session ID header in response
                res.set_header("Mcp-Session-Id", session_id);

                // Return JSON response
                res.set_content(response.dump(), "application/json");
                res.status = 200;
            }
            catch (const fastmcpp::NotFoundError& e)
            {
                // Method/tool not found → -32601
                fastmcpp::Json error_response;
                error_response["jsonrpc"] = "2.0";
                try
                {
                    auto request = fastmcpp::util::json::parse(req.body);
                    if (request.contains("id"))
                        error_response["id"] = request["id"];
                }
                catch (...)
                {
                }
                error_response["error"] = {{"code", -32601}, {"message", std::string(e.what())}};

                res.set_content(error_response.dump(), "application/json");
                res.status = 200; // JSON-RPC errors are still 200 OK at HTTP level
            }
            catch (const fastmcpp::ValidationError& e)
            {
                // Invalid params → -32602
                fastmcpp::Json error_response;
                error_response["jsonrpc"] = "2.0";
                try
                {
                    auto request = fastmcpp::util::json::parse(req.body);
                    if (request.contains("id"))
                        error_response["id"] = request["id"];
                }
                catch (...)
                {
                }
                error_response["error"] = {{"code", -32602}, {"message", std::string(e.what())}};

                res.set_content(error_response.dump(), "application/json");
                res.status = 200;
            }
            catch (const std::exception& e)
            {
                // Internal error → -32603
                fastmcpp::Json error_response;
                error_response["jsonrpc"] = "2.0";
                try
                {
                    auto request = fastmcpp::util::json::parse(req.body);
                    if (request.contains("id"))
                        error_response["id"] = request["id"];
                }
                catch (...)
                {
                }
                error_response["error"] = {{"code", -32603}, {"message", std::string(e.what())}};

                res.set_content(error_response.dump(), "application/json");
                res.status = 500;
            }
        });

    // Handle GET request to return 405 Method Not Allowed
    svr_->Get(
        mcp_path_,
        [this](const httplib::Request&, httplib::Response& res)
        {
            // CORS / additional headers must be applied on every response, including
            // this 405. Without this, browsers reject the response with a misleading
            // "No 'Access-Control-Allow-Origin' header is present" error.
            apply_additional_response_headers(res);

            res.status = 405;
            res.set_header("Allow", "POST, DELETE, OPTIONS");
            res.set_header("Content-Type", "application/json");

            fastmcpp::Json error_response = {
                {"error", "Method Not Allowed"},
                {"message", "The MCP endpoint only supports POST, DELETE, and OPTIONS requests."}};

            res.set_content(error_response.dump(), "application/json");
        });

    // Handle DELETE request for session termination (MCP Streamable HTTP spec).
    // Without this handler, httplib would fall back to its default 404 response,
    // which does not carry the configured CORS headers - causing browsers to report
    // a "No 'Access-Control-Allow-Origin' header is present" error.
    svr_->Delete(mcp_path_,
                 [this](const httplib::Request& req, httplib::Response& res)
                 {
                     apply_additional_response_headers(res);

                     // Security: Check authentication if configured
                     if (!auth_token_.empty())
                     {
                         auto auth_it = req.headers.find("Authorization");
                         if (auth_it == req.headers.end() || !check_auth(auth_it->second))
                         {
                             res.status = 401;
                             res.set_content("{\"error\":\"Unauthorized\"}", "application/json");
                             return;
                         }
                     }

                     auto session_it = req.headers.find("Mcp-Session-Id");
                     if (session_it == req.headers.end() || session_it->second.empty())
                     {
                         res.status = 400;
                         res.set_content("{\"error\":\"Mcp-Session-Id header required\"}",
                                         "application/json");
                         return;
                     }

                     const std::string& session_id = session_it->second;
                     bool did_remove = false;
                     {
                         std::lock_guard<std::mutex> lock(sessions_mutex_);
                         did_remove = sessions_.erase(session_id) > 0;
                     }

                     if (did_remove)
                     {
                         res.status = 204; // No Content
                     }
                     else
                     {
                         res.status = 404;
                         res.set_content("{\"error\":\"Invalid or expired session\"}",
                                         "application/json");
                     }
                 });

    running_ = true;

    thread_ = std::thread([this]() { run_server(); });

    // Wait for server to be ready using GET (returns 405, but shows server is up)
    for (int attempt = 0; attempt < 20; ++attempt)
    {
        if (running_)
        {
            if (const int bp = port(); bp > 0)
            {
                httplib::Client probe(host_.c_str(), bp);
                probe.set_connection_timeout(std::chrono::seconds(2));
                probe.set_read_timeout(std::chrono::seconds(2));
                auto res = probe.Get(mcp_path_.c_str());
                if (res)
                    return true;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
        }
        else
        {
            stop();
            return false; // thread_ signalled failure.
        }
    }

    return true;
}

void StreamableHttpServerWrapper::stop()
{
    running_ = false;

    // Clear sessions
    {
        std::lock_guard<std::mutex> lock(sessions_mutex_);
        sessions_.clear();
    }

    if (svr_)
        svr_->stop();
    if (thread_.joinable())
        thread_.join();

    bound_port_.store(0); // Reset the bound port's value.
}

} // namespace fastmcpp::server