blackbird/libhttp/src/HttpClient.cpp

#include <us/HttpClient.hpp>
#include <chrono>
#include <iostream>

namespace pmc::net {

namespace {

// write_callback 的用户态上下文：流式回调 + 响应大小上限
struct WriteCtx {
    std::function<void(const std::string&)>* callback = nullptr;  // 流式回调（非空时每块回调）
    std::string* full_response = nullptr;                          // 退化累积（callback 为空时）
    size_t received = 0;       // 已接收字节数
    size_t max = 0;            // 响应上限（0 表示不限制）
    bool overflow = false;     // 是否超限
};

} // namespace

HttpClient::HttpClient() {
    curl_ = curl_easy_init();
    if (!curl_) {
        throw std::runtime_error("Failed to initialize CURL");
    }
}

HttpClient::~HttpClient() {
    ws_close();
    if (curl_) {
        curl_easy_cleanup(curl_);
    }
    if (headers_) {
        curl_slist_free_all(headers_);
    }
}

void HttpClient::set_max_response_size(size_t bytes) {
    max_response_size_ = bytes;
}

size_t HttpClient::get_max_response_size() const {
    return max_response_size_;
}

void HttpClient::set_timeout(long seconds) {
    timeout_ = seconds;
}

long HttpClient::get_timeout() const {
    return timeout_;
}

void HttpClient::set_connect_timeout(long seconds) {
    connect_timeout_ = seconds;
}

long HttpClient::get_connect_timeout() const {
    return connect_timeout_;
}

void HttpClient::set_header(const std::string& header) {
    headers_ = curl_slist_append(headers_, header.c_str());
}

void HttpClient::clear_headers() {
    if (headers_) {
        curl_slist_free_all(headers_);
        headers_ = nullptr;
    }
}

const std::map<std::string, std::string>& HttpClient::get_response_headers() const {
    return response_headers_;
}

long HttpClient::get_response_code() const {
    return response_code_;
}

void HttpClient::request_stream_sync(
    HttpMethod method,
    const std::string& url,
    const std::string& data,
    std::function<void(const std::string&)> callback)
{
    if (!curl_) {
        throw std::runtime_error("CURL not initialized");
    }

    response_headers_.clear();
    response_code_ = 0;

    curl_easy_reset(curl_);
    curl_easy_setopt(curl_, CURLOPT_URL, url.c_str());
    set_http_method(method, data);
    curl_easy_setopt(curl_, CURLOPT_HTTPHEADER, headers_);
    curl_easy_setopt(curl_, CURLOPT_HEADERFUNCTION, header_callback);
    curl_easy_setopt(curl_, CURLOPT_HEADERDATA, this);

    std::string full_response;
    WriteCtx write_ctx;
    write_ctx.callback = &callback;
    write_ctx.full_response = &full_response;
    write_ctx.max = max_response_size_;
    curl_easy_setopt(curl_, CURLOPT_WRITEFUNCTION, write_callback);
    curl_easy_setopt(curl_, CURLOPT_WRITEDATA, &write_ctx);

    curl_easy_setopt(curl_, CURLOPT_TIMEOUT, timeout_);
    curl_easy_setopt(curl_, CURLOPT_CONNECTTIMEOUT, connect_timeout_);

    CURLcode res = curl_easy_perform(curl_);

    if (res != CURLE_OK) {
        if (res == CURLE_WRITE_ERROR && write_ctx.overflow) {
            throw std::runtime_error("HTTP response exceeded max size ("
                + std::to_string(max_response_size_) + " bytes)");
        }
        throw std::runtime_error(std::string("CURL error: ") + curl_easy_strerror(res));
    }

    curl_easy_getinfo(curl_, CURLINFO_RESPONSE_CODE, &response_code_);

    // 退化路径：调用方未传 callback 时，整块累积后一次性回调
    if (!full_response.empty()) {
        callback(full_response);
    }
}

void HttpClient::get_stream_sync(const std::string& url, std::function<void(const std::string&)> callback) {
    request_stream_sync(HttpMethod::GET, url, "", callback);
}

void HttpClient::post_stream_sync(const std::string& url, const std::string& data, std::function<void(const std::string&)> callback) {
    request_stream_sync(HttpMethod::POST, url, data, callback);
}

void HttpClient::put_stream_sync(const std::string& url, const std::string& data, std::function<void(const std::string&)> callback) {
    request_stream_sync(HttpMethod::PUT, url, data, callback);
}

void HttpClient::delete_stream_sync(const std::string& url, std::function<void(const std::string&)> callback) {
    request_stream_sync(HttpMethod::DELETE, url, "", callback);
}

bool HttpClient::ws_connect(const std::string& url, int timeout_ms) {
    std::string host, port, path;
    bool ssl = false;

    if (url.size() >= 6 && url.substr(0, 6) == "wss://") {
        ssl = true;
        parse_ws_url(url.substr(6), host, port, path, "443");
    } else if (url.size() >= 5 && url.substr(0, 5) == "ws://") {
        ssl = false;
        parse_ws_url(url.substr(5), host, port, path, "80");
    } else {
        return false;
    }

    return ws_connect(host, port, path, ssl, timeout_ms);
}

bool HttpClient::ws_connect(const std::string& host, const std::string& port,
                            const std::string& path, bool ssl, int timeout_ms) {
    ws_close();
    ws_state_ = std::make_unique<WsState>();
    ws_state_->ssl_enabled = ssl;
    ws_state_->host = host;
    ws_state_->port = port;
    ws_state_->path = path.empty() ? "/" : path;

    try {
        tcp::resolver resolver(ws_state_->ioc);
        auto const results = resolver.resolve(host, port);

        if (ssl) {
            ws_state_->ssl_ctx = std::make_unique<asio::ssl::context>(asio::ssl::context::tlsv12_client);
            ws_state_->ssl_ctx->set_verify_mode(asio::ssl::verify_none);

            ws_state_->tcp_stream = std::make_unique<beast::tcp_stream>(ws_state_->ioc);

            ws_state_->tcp_stream->expires_after(std::chrono::milliseconds(timeout_ms));
            asio::connect(ws_state_->tcp_stream->socket(), results);

            ws_state_->ssl_stream = std::make_unique<asio::ssl::stream<tcp::socket&>>(
                ws_state_->tcp_stream->socket(), *ws_state_->ssl_ctx);

            if (!SSL_set_tlsext_host_name(ws_state_->ssl_stream->native_handle(), host.c_str())) {
                throw std::runtime_error("Failed to set SNI hostname");
            }

            ws_state_->ssl_stream->handshake(asio::ssl::stream_base::client);

            ws_state_->wss_stream = std::make_unique<websocket::stream<asio::ssl::stream<tcp::socket&>&>>(
                *ws_state_->ssl_stream);

            ws_state_->tcp_stream->expires_never();
            ws_state_->wss_stream->set_option(websocket::stream_base::timeout::suggested(beast::role_type::client));
            ws_state_->wss_stream->set_option(websocket::stream_base::decorator(
                [](websocket::request_type& req) {
                    req.set(beast::http::field::user_agent, "UniServe-HttpClient");
                }));
            ws_state_->wss_stream->handshake(host, ws_state_->path);
        } else {
            ws_state_->tcp_stream = std::make_unique<beast::tcp_stream>(ws_state_->ioc);

            ws_state_->tcp_stream->expires_after(std::chrono::milliseconds(timeout_ms));
            asio::connect(ws_state_->tcp_stream->socket(), results);

            ws_state_->ws_stream = std::make_unique<websocket::stream<tcp::socket&>>(
                ws_state_->tcp_stream->socket());

            ws_state_->tcp_stream->expires_never();
            ws_state_->ws_stream->set_option(websocket::stream_base::timeout::suggested(beast::role_type::client));
            ws_state_->ws_stream->set_option(websocket::stream_base::decorator(
                [](websocket::request_type& req) {
                    req.set(beast::http::field::user_agent, "UniServe-HttpClient");
                }));
            ws_state_->ws_stream->handshake(host, ws_state_->path);
        }

        ws_state_->connected = true;
        return true;
    } catch (const std::exception& e) {
        std::cerr << "[HttpClient::ws_connect] " << e.what() << std::endl;
        ws_state_.reset();
        return false;
    }
}

bool HttpClient::ws_send(const std::string& message) {
    if (!ws_state_ || !ws_state_->connected) return false;
    try {
        if (ws_state_->ssl_enabled && ws_state_->wss_stream) {
            ws_state_->wss_stream->write(asio::buffer(message));
        } else if (ws_state_->ws_stream) {
            ws_state_->ws_stream->write(asio::buffer(message));
        }
        return true;
    } catch (const std::exception& e) {
        std::cerr << "[HttpClient::ws_send] " << e.what() << std::endl;
        ws_state_->connected = false;
        return false;
    }
}

bool HttpClient::ws_send_binary(const void* data, size_t size) {
    if (!ws_state_ || !ws_state_->connected) return false;
    try {
        if (ws_state_->ssl_enabled && ws_state_->wss_stream) {
            ws_state_->wss_stream->binary(true);
            ws_state_->wss_stream->write(asio::buffer(data, size));
            ws_state_->wss_stream->text(true);
        } else if (ws_state_->ws_stream) {
            ws_state_->ws_stream->binary(true);
            ws_state_->ws_stream->write(asio::buffer(data, size));
            ws_state_->ws_stream->text(true);
        }
        return true;
    } catch (const std::exception& e) {
        std::cerr << "[HttpClient::ws_send_binary] " << e.what() << std::endl;
        ws_state_->connected = false;
        return false;
    }
}

std::string HttpClient::ws_recv(int timeout_ms) {
    (void)timeout_ms;
    if (!ws_state_ || !ws_state_->connected) return "";
    try {
        beast::flat_buffer buffer;

        if (ws_state_->ssl_enabled && ws_state_->wss_stream) {
            ws_state_->wss_stream->read(buffer);
        } else if (ws_state_->ws_stream) {
            ws_state_->ws_stream->read(buffer);
        }
        return beast::buffers_to_string(buffer.data());
    } catch (const beast::system_error& e) {
        if (e.code() == websocket::error::closed) {
            ws_state_->connected = false;
        } else {
            std::cerr << "[HttpClient::ws_recv] " << e.what() << std::endl;
            ws_state_->connected = false;
        }
        return "";
    } catch (const std::exception& e) {
        std::cerr << "[HttpClient::ws_recv] " << e.what() << std::endl;
        ws_state_->connected = false;
        return "";
    }
}

bool HttpClient::ws_is_connected() const {
    return ws_state_ && ws_state_->connected;
}

void HttpClient::ws_close() {
    if (!ws_state_) return;
    try {
        if (ws_state_->connected) {
            beast::error_code ec;
            if (ws_state_->ssl_enabled && ws_state_->wss_stream) {
                ws_state_->wss_stream->close(websocket::close_code::normal, ec);
            } else if (ws_state_->ws_stream) {
                ws_state_->ws_stream->close(websocket::close_code::normal, ec);
            }
        }
    } catch (...) {
        // Ignore close errors
    }
    ws_state_.reset();
}

std::string HttpClient::ws_get_host() const { return ws_state_ ? ws_state_->host : ""; }
std::string HttpClient::ws_get_port() const { return ws_state_ ? ws_state_->port : ""; }
std::string HttpClient::ws_get_path() const { return ws_state_ ? ws_state_->path : ""; }

void HttpClient::set_http_method(HttpMethod method, const std::string& data) {
    switch (method) {
        case HttpMethod::GET:
            curl_easy_setopt(curl_, CURLOPT_HTTPGET, 1L);
            break;
        case HttpMethod::POST:
            curl_easy_setopt(curl_, CURLOPT_POST, 1L);
            if (!data.empty()) {
                curl_easy_setopt(curl_, CURLOPT_POSTFIELDS, data.c_str());
                curl_easy_setopt(curl_, CURLOPT_POSTFIELDSIZE, data.size());
            }
            break;
        case HttpMethod::PUT:
            curl_easy_setopt(curl_, CURLOPT_CUSTOMREQUEST, "PUT");
            if (!data.empty()) {
                curl_easy_setopt(curl_, CURLOPT_POSTFIELDS, data.c_str());
                curl_easy_setopt(curl_, CURLOPT_POSTFIELDSIZE, data.size());
            }
            break;
        case HttpMethod::DELETE:
            curl_easy_setopt(curl_, CURLOPT_CUSTOMREQUEST, "DELETE");
            break;
        case HttpMethod::PATCH:
            curl_easy_setopt(curl_, CURLOPT_CUSTOMREQUEST, "PATCH");
            if (!data.empty()) {
                curl_easy_setopt(curl_, CURLOPT_POSTFIELDS, data.c_str());
                curl_easy_setopt(curl_, CURLOPT_POSTFIELDSIZE, data.size());
            }
            break;
        case HttpMethod::HEAD:
            curl_easy_setopt(curl_, CURLOPT_NOBODY, 1L);
            break;
        case HttpMethod::OPTIONS:
            curl_easy_setopt(curl_, CURLOPT_CUSTOMREQUEST, "OPTIONS");
            break;
    }
}

size_t HttpClient::header_callback(char* buffer, size_t size, size_t nitems, void* userdata) {
    size_t total_size = size * nitems;
    HttpClient* client = static_cast<HttpClient*>(userdata);
    std::string header_line(buffer, total_size);
    size_t colon_pos = header_line.find(':');
    if (colon_pos != std::string::npos) {
        std::string key = header_line.substr(0, colon_pos);
        std::string value = header_line.substr(colon_pos + 1);
        key.erase(0, key.find_first_not_of(" \t"));
        key.erase(key.find_last_not_of(" \t") + 1);
        value.erase(0, value.find_first_not_of(" \t"));
        value.erase(value.find_last_not_of(" \t\r\n") + 1);
        if (!key.empty() && !value.empty()) {
            client->response_headers_[key] = value;
        }
    }
    return total_size;
}

size_t HttpClient::write_callback(char* ptr, size_t size, size_t nmemb, void* userdata) {
    size_t total_size = size * nmemb;
    auto* ctx = static_cast<WriteCtx*>(userdata);
    // 响应上限：超限返回 0，让 curl 以 CURLE_WRITE_ERROR 中止传输
    if (ctx->max > 0 && ctx->received + total_size > ctx->max) {
        ctx->overflow = true;
        return 0;
    }
    ctx->received += total_size;
    // 流式：callback 非空时每块即回调（不累积，避免二次内存）
    if (ctx->callback && *ctx->callback) {
        (*ctx->callback)(std::string(ptr, total_size));
    } else if (ctx->full_response) {
        // 退化：callback 为空时整块累积，供一次性回调
        ctx->full_response->append(ptr, total_size);
    }
    return total_size;
}

void HttpClient::parse_ws_url(const std::string& url_part, std::string& host, std::string& port,
                              std::string& path, const char* default_port) {
    size_t path_pos = url_part.find('/');
    std::string host_port;
    if (path_pos != std::string::npos) {
        host_port = url_part.substr(0, path_pos);
        path = url_part.substr(path_pos);
    } else {
        host_port = url_part;
        path = "/";
    }
    size_t port_pos = host_port.find(':');
    if (port_pos != std::string::npos) {
        host = host_port.substr(0, port_pos);
        port = host_port.substr(port_pos + 1);
    } else {
        host = host_port;
        port = default_port;
    }
}

} // namespace pmc::net