WSCommModifyService/WebServer.cpp

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <shlwapi.h>

#include "WebServer.h"
#include "WebSocketSession.h"
#include "Logger.h"
#include <fstream>
#include <sstream>
#include <iostream>

WebServer::WebServer(int port, bool use_dll) : port_(port), acceptor_(ioc_), running_(false)
{
    websocket_handler_ = std::make_shared<WebSocketHandler>(use_dll);
}

WebServer::~WebServer()
{
    Stop();
}

bool WebServer::Start()
{
    try
    {
        // 初始化WebSocketHandler（包括ActiveX控件）
        if (!websocket_handler_->Initialize())
        {
            LOG_ERROR("Failed to initialize WebSocketHandler");
            return false;
        }

        // 设置端点
        tcp::endpoint endpoint(net::ip::make_address("0.0.0.0"), port_);

        // 打开acceptor
        acceptor_.open(endpoint.protocol());
        acceptor_.set_option(net::socket_base::reuse_address(true));
        acceptor_.bind(endpoint);
        acceptor_.listen(net::socket_base::max_listen_connections);

        running_ = true;

        // 开始接受连接
        DoAccept();

        // 在单独线程中运行io_context
        server_thread_ = std::thread([this]()
                                     { ioc_.run(); });

        LOG_INFO("WebServer started on port " + std::to_string(port_));
        return true;
    }
    catch (const std::exception &e)
    {
        LOG_ERROR("Failed to start WebServer: " + std::string(e.what()));
        return false;
    }
}

void WebServer::Stop()
{
    if (running_)
    {
        running_ = false;

        // 停止acceptor
        beast::error_code ec;
        acceptor_.close(ec);

        // 停止io_context
        ioc_.stop();

        // 等待线程结束
        if (server_thread_.joinable())
        {
            server_thread_.join();
        }

        // 清理WebSocketHandler
        if (websocket_handler_)
        {
            websocket_handler_->Cleanup();
        }

        LOG_INFO("WebServer stopped");
    }
}

void WebServer::DoAccept()
{
    acceptor_.async_accept(
        net::make_strand(ioc_),
        [self = shared_from_this()](beast::error_code ec, tcp::socket socket)
        {
            self->OnAccept(ec, std::move(socket));
        });
}

void WebServer::OnAccept(beast::error_code ec, tcp::socket socket)
{
    if (ec)
    {
        if (ec != net::error::operation_aborted)
        {
            LOG_ERROR("Accept error: " + ec.message());
        }
        return;
    }

    // 创建HTTP会话来处理连接
    std::make_shared<HttpSession>(std::move(socket), shared_from_this())->run();

    // 继续接受新连接
    if (running_)
    {
        DoAccept();
    }
}

std::string WebServer::GetStaticFileContent(const std::string &path)
{
    // 获取可执行文件所在目录
    char buffer[MAX_PATH];
    GetModuleFileNameA(NULL, buffer, MAX_PATH);
    std::string::size_type pos = std::string(buffer).find_last_of("\\");
    std::string exe_path = std::string(buffer).substr(0, pos);

    std::string file_path;

    if (path == "/" || path == "/index.html")
    {
        file_path = exe_path + "\\web\\index.html";
    }
    else if (path == "/style.css")
    {
        file_path = exe_path + "\\web\\style.css";
    }
    else if (path == "/script.js")
    {
        file_path = exe_path + "\\web\\script.js";
    }
    else
    {
        return "";
    }

    // 以二进制模式打开文件，避免编码问题
    std::ifstream file(file_path, std::ios::binary);
    if (!file.is_open())
    {
        LOG_ERROR("Failed to open file: " + file_path);
        return "";
    }

    // 读取文件内容
    std::stringstream content_buffer;
    content_buffer << file.rdbuf();
    return content_buffer.str();
}

std::string WebServer::GetMimeType(const std::string &path)
{
    if (path == "/" || path == "/index.html" || path.ends_with(".html"))
    {
        return "text/html; charset=utf-8";
    }
    if (path.ends_with(".css"))
    {
        return "text/css; charset=utf-8";
    }
    if (path.ends_with(".js"))
    {
        return "application/javascript; charset=utf-8";
    }
    if (path.ends_with(".json"))
    {
        return "application/json; charset=utf-8";
    }
    return "text/plain; charset=utf-8";
}

void WebServer::HandleHttpRequest(http::request<http::dynamic_body> &&req,
                                  std::function<void(http::response<http::dynamic_body>)> send)
{
    // 处理普通HTTP请求
    auto const bad_request = [&req](beast::string_view why)
    {
        http::response<http::dynamic_body> res{http::status::bad_request, req.version()};
        res.set(http::field::server, BOOST_BEAST_VERSION_STRING);
        res.set(http::field::content_type, "text/html; charset=utf-8");
        res.keep_alive(req.keep_alive());
        beast::ostream(res.body()) << "Bad request: " << why;
        res.prepare_payload();
        return res;
    };

    auto const not_found = [&req](beast::string_view target)
    {
        http::response<http::dynamic_body> res{http::status::not_found, req.version()};
        res.set(http::field::server, BOOST_BEAST_VERSION_STRING);
        res.set(http::field::content_type, "text/html; charset=utf-8");
        res.keep_alive(req.keep_alive());
        beast::ostream(res.body()) << "The resource '" << target << "' was not found.";
        res.prepare_payload();
        return res;
    };

    // 确保我们可以处理这个方法
    if (req.method() != http::verb::get && req.method() != http::verb::head)
    {
        return send(bad_request("Unknown HTTP-method"));
    }

    // 请求路径不能包含".."
    if (req.target().empty() || req.target()[0] != '/' || req.target().find("..") != beast::string_view::npos)
    {
        return send(bad_request("Illegal request-target"));
    }

    // 获取文件内容
    std::string path = std::string(req.target());
    std::string content = GetStaticFileContent(path);

    if (content.empty())
    {
        return send(not_found(req.target()));
    }

    // 创建响应
    http::response<http::dynamic_body> res{http::status::ok, req.version()};
    res.set(http::field::server, BOOST_BEAST_VERSION_STRING);
    res.set(http::field::content_type, GetMimeType(path));
    res.keep_alive(req.keep_alive());

    // 设置响应体
    beast::ostream(res.body()) << content;
    res.prepare_payload();

    return send(std::move(res));
}

// HttpSession实现
HttpSession::HttpSession(tcp::socket &&socket, std::shared_ptr<WebServer> server)
    : stream_(std::move(socket)), server_(server)
{
}

void HttpSession::run()
{
    net::dispatch(stream_.get_executor(),
                  beast::bind_front_handler(&HttpSession::do_read, shared_from_this()));
}

void HttpSession::do_read()
{
    req_ = {};

    stream_.expires_after(std::chrono::seconds(30));

    http::async_read(stream_, buffer_, req_,
                     beast::bind_front_handler(&HttpSession::on_read, shared_from_this()));
}

void HttpSession::on_read(beast::error_code ec, std::size_t bytes_transferred)
{
    boost::ignore_unused(bytes_transferred);

    if (ec == http::error::end_of_stream)
    {
        return do_close();
    }

    if (ec)
    {
        LOG_ERROR("HTTP read error: " + ec.message());
        return;
    }

    // 检查是否是WebSocket升级请求
    if (websocket::is_upgrade(req_))
    {
        // 释放socket并传递给WebSocket会话
        tcp::socket socket = stream_.release_socket();

        // 创建WebSocket会话
        auto ws_session = std::make_shared<WebSocketSession>(std::move(socket), std::move(req_));

        // 设置WebSocketHandler引用
        ws_session->set_websocket_handler(server_->GetWebSocketHandler());

        // 设置消息处理回调
        ws_session->set_message_handler([server = server_](std::shared_ptr<WebSocketSession> session, const std::string &message)
                                        { server->GetWebSocketHandler()->HandleMessage(session, message); });

        // 添加到WebSocket处理器
        server_->GetWebSocketHandler()->AddClient(ws_session);

        ws_session->run();
        return;
    }

    // 处理普通HTTP请求
    server_->HandleHttpRequest(std::move(req_),
                               [self = shared_from_this()](http::response<http::dynamic_body> res)
                               {
                                   // 发送响应
                                   auto sp = std::make_shared<http::response<http::dynamic_body>>(std::move(res));

                                   http::async_write(self->stream_, *sp,
                                                     [self, sp](beast::error_code ec, std::size_t bytes_transferred)
                                                     {
                                                         self->on_write(sp->need_eof(), ec, bytes_transferred);
                                                     });
                               });
}

void HttpSession::on_write(bool close, beast::error_code ec, std::size_t bytes_transferred)
{
    boost::ignore_unused(bytes_transferred);

    if (ec)
    {
        LOG_ERROR("HTTP write error: " + ec.message());
        return;
    }

    if (close)
    {
        return do_close();
    }

    // 读取另一个请求
    do_read();
}

void HttpSession::do_close()
{
    beast::error_code ec;
    stream_.socket().shutdown(tcp::socket::shutdown_send, ec);
}