sweeper_video/src/tunnel_client.cpp

#include "tunnel_client.hpp"

#include <atomic>
#include <chrono>
#include <cstdio>
#include <mutex>
#include <nlohmann/json.hpp>
#include <string>
#include <thread>
#include <vector>
#include <websocketpp/client.hpp>
#include <websocketpp/config/asio_no_tls_client.hpp>

#include "httplib.h"
#include "logger.hpp"

using json = nlohmann::json;
using ws_client = websocketpp::client<websocketpp::config::asio_client>;

namespace
{
constexpr size_t kChunkSize = 64 * 1024;  // 64KB
constexpr int kHttpTimeoutSec = 30;

static bool looks_binary_content_type(const std::string& ct)
{
    if (ct.find("video/") == 0) return true;
    if (ct.find("application/octet-stream") != std::string::npos) return true;
    if (ct.find("image/") == 0) return true;
    return false;
}

// 小工具：拼 log 字符串
static inline std::string kv(const std::string& k, const std::string& v) { return k + "=" + v; }
static inline std::string kv(const std::string& k, int v) { return k + "=" + std::to_string(v); }

}  // namespace

TunnelClient::TunnelClient(const std::string& vid, const std::string& server_ws_url, int local_http_port)
    : vid_(vid), ws_url_(server_ws_url + "?vid=" + vid), local_port_(local_http_port)
{
}

void TunnelClient::start()
{
    running_.store(true);
    th_ = std::thread(&TunnelClient::run_loop, this);
}

void TunnelClient::stop()
{
    running_.store(false);

    ws_client* cli = client_;  // copy to avoid race
    if (cli)
    {
        cli->get_io_service().post(
            [this, cli]()
            {
                try
                {
                    websocketpp::lib::error_code ec;
                    cli->close(hdl_, websocketpp::close::status::going_away, "stop", ec);
                }
                catch (...)
                {
                }
            });
    }

    if (th_.joinable()) th_.join();
}

// =========================== send helpers (thread-safe) ===========================

void TunnelClient::send_text_safe(const std::string& s)
{
    if (!client_) return;

    client_->get_io_service().post(
        [this, s]()
        {
            websocketpp::lib::error_code ec;
            client_->send(hdl_, s, websocketpp::frame::opcode::text, ec);
            if (ec)
            {
                LOG_ERROR(std::string("[Tunnel] send_text failed: ") + ec.message());
            }
        });
}

void TunnelClient::send_binary_safe(const void* data, size_t len)
{
    if (!client_) return;

    std::string payload(reinterpret_cast<const char*>(data), len);  // copy to avoid buffer lifetime issue
    client_->get_io_service().post(
        [this, payload]()
        {
            websocketpp::lib::error_code ec;
            client_->send(hdl_, payload, websocketpp::frame::opcode::binary, ec);
            if (ec)
            {
                LOG_ERROR(std::string("[Tunnel] send_binary failed: ") + ec.message());
            }
        });
}

// =========================== core: proxy local HTTP and reply ===========================

void TunnelClient::handle_request_and_reply(const json& req)
{
    const std::string req_id = req.value("req_id", "");
    const std::string method = req.value("method", "GET");
    const std::string path = req.value("path", "/");
    const std::string body = req.value("body", "");

    LOG_INFO("[Tunnel] " + kv("req_id", req_id) + " " + method + " " + path);

    httplib::Client cli("127.0.0.1", local_port_);
    cli.set_connection_timeout(kHttpTimeoutSec, 0);
    cli.set_read_timeout(kHttpTimeoutSec, 0);
    cli.set_write_timeout(kHttpTimeoutSec, 0);
    cli.set_keep_alive(true);

    if (method == "GET")
    {
        bool started_header = false;
        int status_code = 0;
        std::string content_type = "application/octet-stream";

        auto ok = cli.Get(
            path.c_str(),
            // on_response: got headers
            [&](const httplib::Response& res)
            {
                status_code = res.status;
                auto it = res.headers.find("Content-Type");
                if (it != res.headers.end()) content_type = it->second;

                json hdr = {
                    {"type", "header"},
                    {"req_id", req_id},
                    {"status", status_code},
                    {"content_type", content_type},
                };
                send_text_safe(hdr.dump());
                started_header = true;

                LOG_INFO("[Tunnel] local " + kv("status", status_code) + " " + kv("ct", content_type));
                return true;  // continue
            },
            // on_data: stream body in binary frames
            [&](const char* data, size_t data_length)
            {
                if (!started_header)
                {
                    json hdr = {
                        {"type", "header"},
                        {"req_id", req_id},
                        {"status", 200},
                        {"content_type", content_type},
                    };
                    send_text_safe(hdr.dump());
                    started_header = true;
                }

                send_binary_safe(data, data_length);
                return true;  // continue
            });

        if (!ok)
        {
            json err = {{"type", "json"}, {"req_id", req_id}, {"status", 500}, {"body", "local http error"}};
            send_text_safe(err.dump());
            LOG_ERROR("[Tunnel] local http error (GET) " + kv("req_id", req_id));
            return;
        }

        json end = {{"type", "end"}, {"req_id", req_id}};
        send_text_safe(end.dump());
        LOG_INFO("[Tunnel] done " + kv("req_id", req_id));
        return;
    }

    if (method == "POST")
    {
        auto res = cli.Post(path.c_str(), body, "application/json");
        if (!res)
        {
            json err = {{"type", "json"}, {"req_id", req_id}, {"status", 500}, {"body", "local http error"}};
            send_text_safe(err.dump());
            LOG_ERROR("[Tunnel] local http error (POST) " + kv("req_id", req_id));
            return;
        }

        std::string ct = "application/json";
        auto it = res->headers.find("Content-Type");
        if (it != res->headers.end()) ct = it->second;

        json out = {
            {"type", "json"}, {"req_id", req_id}, {"status", res->status}, {"content_type", ct}, {"body", res->body},
        };
        send_text_safe(out.dump());
        LOG_INFO("[Tunnel] POST done " + kv("req_id", req_id) + " " + kv("status", res->status));
        return;
    }

    json out = {{"type", "json"}, {"req_id", req_id}, {"status", 405}, {"body", "unsupported method"}};
    send_text_safe(out.dump());
    LOG_WARN("[Tunnel] unsupported method " + kv("req_id", req_id) + " " + kv("method", method));
}

// =========================== run loop ===========================

void TunnelClient::run_loop()
{
    while (running_.load())
    {
        ws_client c;
        c.clear_access_channels(websocketpp::log::alevel::all);
        c.clear_error_channels(websocketpp::log::elevel::all);

        c.init_asio();
        c.start_perpetual();

        // expose current client to stop()
        client_ = &c;

        c.set_open_handler(
            [this](websocketpp::connection_hdl h)
            {
                hdl_ = h;
                LOG_INFO("[Tunnel] Connected to server");
            });

        // IMPORTANT: do not print reconnect related logs
        c.set_close_handler(
            [this, &c](websocketpp::connection_hdl)
            {
                c.stop_perpetual();
                c.stop();
            });

        c.set_fail_handler(
            [this, &c](websocketpp::connection_hdl)
            {
                c.stop_perpetual();
                c.stop();
            });

        c.set_message_handler(
            [this](websocketpp::connection_hdl, ws_client::message_ptr msg)
            {
                try
                {
                    if (msg->get_opcode() != websocketpp::frame::opcode::text)
                    {
                        LOG_WARN("[Tunnel] ignore non-text msg from server");
                        return;
                    }

                    json req = json::parse(msg->get_payload());

                    std::thread(
                        [this, req]()
                        {
                            try
                            {
                                handle_request_and_reply(req);
                            }
                            catch (const std::exception& e)
                            {
                                LOG_ERROR(std::string("[Tunnel] handler exception: ") + e.what());
                            }
                            catch (...)
                            {
                                LOG_ERROR("[Tunnel] handler unknown exception");
                            }
                        })
                        .detach();
                }
                catch (const std::exception& e)
                {
                    LOG_ERROR(std::string("[Tunnel] JSON parse error: ") + e.what());
                }
                catch (...)
                {
                    LOG_ERROR("[Tunnel] JSON parse unknown error");
                }
            });

        websocketpp::lib::error_code ec;
        auto conn = c.get_connection(ws_url_, ec);
        if (ec)
        {
            client_ = nullptr;
            // do not print reconnect logs; just sleep and retry
            std::this_thread::sleep_for(std::chrono::seconds(2));
            continue;
        }

        c.connect(conn);

        try
        {
            c.run();
        }
        catch (const std::exception& e)
        {
            LOG_ERROR(std::string("[Tunnel] run exception: ") + e.what());
        }
        catch (...)
        {
            LOG_ERROR("[Tunnel] run unknown exception");
        }

        client_ = nullptr;

        if (!running_.load()) break;

        // do not print reconnect logs
        std::this_thread::sleep_for(std::chrono::seconds(2));
    }

    client_ = nullptr;
    LOG_INFO("[Tunnel] Loop exit");
}