WSCommModifyService/WebSocketSession.cpp

#include "WebSocketSession.h"
#include "WebSocketHandler.h"
#include "Logger.h"
#include <boost/beast/version.hpp>

WebSocketSession::WebSocketSession(tcp::socket&& socket, http::request<http::dynamic_body> req)
    : ws_(std::move(socket)), req_(std::move(req)) {
}

WebSocketSession::~WebSocketSession() {
    close();
    
    // 注意：不应该在这里调用RemoveClient，因为如果WebSocketHandler仍然持有shared_ptr，
    // 析构函数永远不会被调用。正确的做法是在RemoveClient时触发析构。
}

void WebSocketSession::run() {
    // 设置建议的超时设置
    beast::get_lowest_layer(ws_).expires_never();
    
    // 设置WebSocket超时参数
    ws_.set_option(websocket::stream_base::timeout{
        std::chrono::seconds(30),  // handshake timeout
        std::chrono::seconds(30),  // idle timeout
        true                      // 取消关闭连接前的静默超时
    });
    
    // 设置装饰器来改变服务器握手
    ws_.set_option(websocket::stream_base::decorator(
        [](websocket::response_type& res) {
            res.set(http::field::server, std::string(BOOST_BEAST_VERSION_STRING) + " websocket-server-async");
        }));
    
    // 接受WebSocket握手
    ws_.async_accept(
        req_,
        beast::bind_front_handler(&WebSocketSession::on_accept, shared_from_this()));
}

void WebSocketSession::on_accept(beast::error_code ec) {
    if (ec) {
        LOG_ERROR("WebSocket accept failed: " + ec.message());
        return;
    }
    
    LOG_INFO("WebSocket connection established");
    
    // 开始读取消息
    do_read();
}

void WebSocketSession::do_read() {
    if (closed_.load()) {
        return;
    }
    
    // 读取消息
    ws_.async_read(
        buffer_,
        beast::bind_front_handler(&WebSocketSession::on_read, shared_from_this()));
}

void WebSocketSession::on_read(beast::error_code ec, std::size_t bytes_transferred) {
    boost::ignore_unused(bytes_transferred);
    
    if (ec) {
        if (ec != websocket::error::closed) {
            LOG_ERROR("WebSocket read failed: " + ec.message());
        }
        
        // 通知处理器移除客户端
        if (websocket_handler_) {
            websocket_handler_->RemoveClient(shared_from_this());
        }
        return;
    }
    
    // 处理接收到的消息
    std::string message = beast::buffers_to_string(buffer_.data());
    buffer_.clear();
    
    // 调用消息处理回调
    if (message_handler_) {
        message_handler_(shared_from_this(), message);
    }
    
    // 继续读取下一个消息
    do_read();
}

void WebSocketSession::safe_send(std::shared_ptr<std::string> message) {
    if (closed_.load()) {
        return;
    }
    
    std::lock_guard<std::mutex> lock(send_mutex_);
    send_queue_.push(message);
    
    if (!is_sending_) {
        do_write();
    }
}

void WebSocketSession::do_write() {
    // 注意：此方法应该在持有send_mutex_锁的情况下调用
    if (closed_.load() || send_queue_.empty()) {
        is_sending_ = false;
        return;
    }
    
    is_sending_ = true;
    current_message_ = send_queue_.front();
    send_queue_.pop();
    
    // 使用成员变量current_message_确保在异步操作期间消息不会被销毁
    ws_.async_write(
        net::buffer(*current_message_),
        beast::bind_front_handler(&WebSocketSession::on_write, shared_from_this()));
}

void WebSocketSession::on_write(beast::error_code ec, std::size_t bytes_transferred) {
    boost::ignore_unused(bytes_transferred);
    
    std::lock_guard<std::mutex> lock(send_mutex_);
    current_message_.reset(); // 清理当前消息
    
    if (ec) {
        LOG_ERROR("WebSocket write failed: " + ec.message());
        is_sending_ = false;
        
        // 如果写入失败，关闭连接并通知处理器移除客户端
        if (websocket_handler_) {
            close();
            websocket_handler_->RemoveClient(shared_from_this());
        }
        return;
    }
    
    if (!send_queue_.empty()) {
        do_write();
    } else {
        is_sending_ = false;
    }
}

void WebSocketSession::close() {
    if (closed_.exchange(true)) {
        return; // 已经关闭
    }
    
    beast::error_code ec;
    ws_.close(websocket::close_code::normal, ec);
    if (ec) {
        LOG_ERROR("WebSocket close failed: " + ec.message());
    }
}

void WebSocketSession::set_message_handler(MessageHandler handler) {
    message_handler_ = handler;
}