blackbird/libhttp/inc/us/HttpServer.hpp

#ifndef HTTPSERVER_HPP
#define HTTPSERVER_HPP

#include <boost/asio.hpp>
#include <boost/beast.hpp>
#include <boost/beast/http.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/beast/version.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/ssl.hpp>
#include <boost/asio/strand.hpp>
#include <boost/config.hpp>
#include <memory>
#include <string>
#include <functional>
#include <unordered_map>
#include <vector>
#include <iostream>
#include <thread>
#include <atomic>
#include <mutex>
#include <regex>
#include <deque>
#include <chrono>

namespace pmc {
namespace net {

namespace beast = boost::beast;
namespace http  = beast::http;
namespace websocket = beast::websocket;
namespace asio  = boost::asio;
using tcp = boost::asio::ip::tcp;

using HttpRequestHandler = std::function<http::response<http::string_body>(
    const http::request<http::string_body>&,
    const std::unordered_map<std::string, std::string>&)>;

using MiddlewareHandler = std::function<bool(
    const http::request<http::string_body>&,
    http::response<http::string_body>&,
    const std::unordered_map<std::string, std::string>&)>;

// WebSocket 回调类型
using WebSocketMessageHandler = std::function<void(
    const std::string& clientId,
    const std::string& message,
    std::function<void(const std::string&)> sendCallback)>;

using WebSocketConnectHandler = std::function<void(
    const std::string& clientId,
    const std::string& remoteAddress)>;

using WebSocketDisconnectHandler = std::function<void(
    const std::string& clientId)>;

// WebSocket 代理：握手回调接收 HTTP upgrade 请求，返回目标 ws:// 或 wss:// URL
// 返回空字符串表示拒绝（返回 403）
using WsProxyHandshakeHandler = std::function<std::string(
    const http::request<http::string_body>& upgradeReq)>;

struct RouteEntry {
    std::string path;
    std::regex pattern;
    HttpRequestHandler handler;
    bool is_wildcard;
};

struct WsRouteEntry {
    std::string path;
    WebSocketMessageHandler handler;
};

struct WsProxyRouteEntry {
    std::string path;
    WsProxyHandshakeHandler handshake;
};

class HttpServer {
public:
    explicit HttpServer(const std::string& address, unsigned short port,
                        unsigned int threads = 1);
    explicit HttpServer(const std::string& address, unsigned short port,
                        const std::string& cert_file, const std::string& key_file,
                        unsigned int threads = 1);
    ~HttpServer();

    void start();
    void stop();

    void get(const std::string& path, HttpRequestHandler handler);
    void post(const std::string& path, HttpRequestHandler handler);
    void put(const std::string& path, HttpRequestHandler handler);
    void del(const std::string& path, HttpRequestHandler handler);
    void options(const std::string& path, HttpRequestHandler handler);
    void use(MiddlewareHandler middleware);
    void setStaticDirectory(const std::string& path);

    // WebSocket 接口
    void ws(const std::string& path, WebSocketMessageHandler handler);
    void ws_proxy(const std::string& path, WsProxyHandshakeHandler handshake);
    void onWsConnect(WebSocketConnectHandler handler);
    void onWsDisconnect(WebSocketDisconnectHandler handler);
    bool sendToWsClient(const std::string& clientId, const std::string& message);
    size_t broadcastWs(const std::string& message);
    void disconnectWsClient(const std::string& clientId);
    size_t getWsClientCount() const;

    bool isRunning() const;
    unsigned short getPort() const;
    std::string getAddress() const;
    bool isSSLEnabled() const;

    void run();
    void showStatus() const;

    void setPrivateKeyPassword(const std::string& password);

    // 请求体大小上限（字节）。默认 10MB，与 WAF 对齐；超限返回 413。
    void setMaxBodySize(std::size_t bytes) { max_body_size_ = bytes; }
    std::size_t getMaxBodySize() const { return max_body_size_; }

    // WebSocket 单条消息大小上限（字节）。默认 16MB，与 Beast 默认一致（防 OOM）。
    // 超限时 Beast 拒绝该消息并关闭连接（P1-3）。
    void setMaxWsMessageSize(std::size_t bytes) { max_ws_message_size_ = bytes; }
    std::size_t getMaxWsMessageSize() const { return max_ws_message_size_; }

    // 入站读超时（防 Slowloris，P1-4）。默认与 WAF 对齐：头 10s、体 30s。
    // 头超时覆盖 TCP 连接建立 + SSL 握手 + 请求头读取；体超时覆盖请求体读取。
    void setHeaderTimeout(std::chrono::seconds t) { header_timeout_ = t; }
    std::chrono::seconds getHeaderTimeout() const { return header_timeout_; }
    void setReadTimeout(std::chrono::seconds t) { read_timeout_ = t; }
    std::chrono::seconds getReadTimeout() const { return read_timeout_; }

    // WS 发送/代理队列上限（条数）。默认 1000。防止慢消费者导致无界队列内存耗尽（P1-6）。
    // 超限时判定为慢消费者，关闭对应连接。
    void setMaxWsQueueSize(std::size_t n) { max_ws_queue_size_ = n; }
    std::size_t getMaxWsQueueSize() const { return max_ws_queue_size_; }

    class debug {
    public:
        static void print_all_header_fields(const http::request<http::string_body>& req);
    };

private:
    class Session;
    class Listener;

    std::string address_;
    unsigned short port_;
    unsigned int threads_;
    std::atomic<bool> running_{false};

    std::size_t max_body_size_ = 10 * 1024 * 1024; // 默认 10MB,与 WAF 对齐
    std::size_t max_ws_message_size_ = 16 * 1024 * 1024; // 默认 16MB,与 Beast 默认一致
    std::chrono::seconds header_timeout_{10}; // 头/握手超时，防 Slowloris（对齐 WAF ReadHeaderTimeout）
    std::chrono::seconds read_timeout_{30};   // 体超时（对齐 WAF ReadTimeout）
    std::size_t max_ws_queue_size_ = 1000;    // WS 发送/代理队列上限，防慢消费者 OOM（P1-6）

    bool ssl_enabled_{false};
    std::string cert_file_;
    std::string key_file_;
    std::string private_key_password_;

    std::unique_ptr<asio::ssl::context> ssl_ctx_;
    std::unique_ptr<asio::io_context> ioc_;

    std::shared_ptr<Listener> listener_;
    std::vector<std::thread> worker_threads_;

    std::vector<RouteEntry> get_handlers_;
    std::vector<RouteEntry> post_handlers_;
    std::vector<RouteEntry> put_handlers_;
    std::vector<RouteEntry> delete_handlers_;
    std::vector<RouteEntry> options_handlers_;
    std::vector<MiddlewareHandler> middlewares_;
    std::string static_directory_;

    // WebSocket 路由和处理
    std::vector<WsRouteEntry> ws_handlers_;
    std::vector<WsProxyRouteEntry> ws_proxy_handlers_;
    WebSocketConnectHandler ws_connect_handler_;
    WebSocketDisconnectHandler ws_disconnect_handler_;
    std::unordered_map<std::string, std::shared_ptr<Session>> ws_clients_;
    mutable std::mutex ws_clients_mutex_;
    std::atomic<uint64_t> ws_client_counter_{0};

    std::vector<std::weak_ptr<Session>> sessions_;
    std::mutex sessions_mutex_;

    std::unordered_map<std::string, std::string> parseQueryParams(const std::string& query);
    bool isSafeStaticPath(const std::string& path) const;
    std::regex wildcardToRegex(const std::string& pattern);
    bool matchPath(const std::string& path, const std::regex& pattern);
    http::response<http::string_body> handleRequest(const http::request<http::string_body>& req);
    HttpRequestHandler findRoute(const std::string& method, const std::string& path);
    bool routeExists(const std::string& method, const std::string& path);
    WebSocketMessageHandler findWsRoute(const std::string& path);
    WsProxyHandshakeHandler findWsProxyRoute(const std::string& path);

    // WebSocket 客户端管理
    std::string generateWsClientId();
    void addWsClient(const std::string& clientId, std::shared_ptr<Session> session);
    void removeWsClient(const std::string& clientId);
    std::shared_ptr<Session> getWsClient(const std::string& clientId);

    void addSession(const std::shared_ptr<Session>& session);
    void removeSession(const Session* session);
    void closeAllSessions();
    void initSSLContext();

    friend class Session;
};

} // namespace net
} // namespace pmc

#endif // HTTPSERVER_HPP