qing
b2ecb8e12b
worker 收 EOF 失败后过早清除 busy 标记，事件循环对同一 fd 重复投递任务， 迟到的失败标记经 fd 复用落在新连接上被 _reap_failed 误杀（客户端收到空 EOF）。 失败路径改由 _reap_failed 统一清除 busy 并关闭；_handle_client 增加防重复投递守卫。 memdb 短连接压力测试错误率从 5-9% 降至 0。
657 lines
27 KiB
C++
657 lines
27 KiB
C++
#pragma once
|
||
|
||
/** MsgServer — 基于 pymsg 协议的多客户端 TCP 服务器（C++ 版，对应 Python 版 pymsg.MsgServer）
|
||
*
|
||
* 提供 select-based 事件循环 + 可选线程池并发处理，支持 reply()/broadcast() 线程安全发送。
|
||
*
|
||
* 设计要点（与 Python 版对齐）：
|
||
* - select 事件循环，支持多客户端并发连接
|
||
* - workers=0 同步模式（回调在主线程执行），workers>0 线程池并发模式
|
||
* - 同一连接消息严格串行：回调执行期间暂停该连接 select 监听
|
||
* - reply()/broadcast() 经每连接发送锁串行化，防止帧交错
|
||
* - max_connections 限制活跃连接数，防止连接耗尽攻击
|
||
* - max_queue_size 限制线程池任务队列深度，慢 handler 时断开连接防 OOM
|
||
* - deadline 模式 recv 超时，防 Slowloris 攻击
|
||
* - stop() 可从任意线程（包括回调内部）安全调用，优雅关闭
|
||
* - worker 线程出错仅标记 failed，主循环统一清理连接，避免 fd 并发操作
|
||
*
|
||
* Usage:
|
||
* #include <us/MsgServer.hpp>
|
||
*
|
||
* pmc::net::MsgServer server("0.0.0.0", 8003, {.workers=4});
|
||
* server.start([&](int fd, const std::string& addr, const std::string& msg) {
|
||
* server.reply(fd, "echo: " + msg); // 回复发送者
|
||
* server.broadcast(addr + ": " + msg, fd); // 广播给其他人
|
||
* });
|
||
*/
|
||
|
||
#include <us/Msg_Basic.hpp>
|
||
#include <us/Msg.hpp>
|
||
|
||
#include <algorithm>
|
||
#include <atomic>
|
||
#include <chrono>
|
||
#include <condition_variable>
|
||
#include <cstdio>
|
||
#include <cstring>
|
||
#include <deque>
|
||
#include <functional>
|
||
#include <memory>
|
||
#include <mutex>
|
||
#include <queue>
|
||
#include <string>
|
||
#include <thread>
|
||
#include <unordered_map>
|
||
#include <unordered_set>
|
||
#include <utility>
|
||
#include <vector>
|
||
|
||
#include <arpa/inet.h>
|
||
#include <fcntl.h>
|
||
#include <netinet/in.h>
|
||
#include <netinet/tcp.h>
|
||
#include <sys/socket.h>
|
||
#include <sys/select.h>
|
||
#include <sys/types.h>
|
||
#include <unistd.h>
|
||
|
||
#ifndef MSG_NOSIGNAL
|
||
#define MSG_NOSIGNAL 0
|
||
#endif
|
||
|
||
namespace pmc {
|
||
namespace net {
|
||
|
||
// ============================================================================
|
||
// MsgServer 配置
|
||
// ============================================================================
|
||
|
||
struct MsgServerConfig {
|
||
int workers = 0; ///< 线程池大小；0 = 同步模式（回调在主线程执行）
|
||
double recv_timeout = 10.0; ///< 接收单条消息的总超时（秒），防 Slowloris
|
||
int max_connections = 1000; ///< 最大并发连接数，超过则拒绝新连接
|
||
int max_queue_size = 1000; ///< 线程池任务队列上限（workers>0 时），满时断开连接
|
||
int listen_backlog = 10; ///< listen() backlog
|
||
double poll_timeout = 0.1; ///< select 轮询超时（秒），影响关闭响应速度
|
||
};
|
||
|
||
// ============================================================================
|
||
// ConnectionLock — 每连接发送锁 + 活跃标志
|
||
// ============================================================================
|
||
|
||
/// 每连接的发送锁和活跃标志。
|
||
///
|
||
/// [libmsg-audit P0-1 / P0-2 / P1-3] 组合 send_mtx（互斥锁）+ active（原子标志）彻底消除
|
||
/// CWE-362 竞态：
|
||
/// - 第一轮 P0-1：单客户端移除路径先持 send_mtx 再 close，reply/broadcast 先获取
|
||
/// send_mtx 则 _remove_client 等待，防止 fd 关闭后被 OS 复用导致数据发往错误客户端。
|
||
/// - 第二轮 P0-2：批量关闭路径 _close_all_clients 同样持 send_mtx 后 shutdown+close，
|
||
/// 修复第一轮只覆盖单路径的竞态回归。
|
||
/// - 第四轮 P1-3：active 标志堵住反向竞态——_remove_client 先获取锁并 close(fd) 后，
|
||
/// reply/broadcast 后获取锁时通过 active==false 得知 fd 已关闭，立即放弃发送。
|
||
struct ConnectionLock {
|
||
std::mutex send_mtx;
|
||
std::atomic<bool> active{true};
|
||
};
|
||
|
||
// ============================================================================
|
||
// MsgServer
|
||
// ============================================================================
|
||
|
||
/// 基于 pymsg 长度前缀协议的多客户端 TCP 服务器。
|
||
class MsgServer {
|
||
public:
|
||
/// 消息回调签名：(客户端 fd, 客户端地址 "ip:port", 消息内容)
|
||
using MessageCallback = std::function<void(int sockfd, const std::string& addr, const std::string& message)>;
|
||
|
||
explicit MsgServer(const std::string& addr, int port, MsgServerConfig config = {})
|
||
: _addr(addr), _port(port), _config(config), _server_fd(-1), _running(false) {
|
||
// [libmsg-audit P2-3 / P2-5] 全量边界校验：拒绝非法 port / workers / recv_timeout /
|
||
// max_connections / max_queue_size / listen_backlog / poll_timeout。防止 htons 截断产生
|
||
// 错误端口、负值导致所有连接被拒或队列无法提交、极大值（>300s）绕过 Slowloris 防护、
|
||
// poll_timeout=0 导致 select 忙轮询 CPU 100%。
|
||
detail::validatePort(port);
|
||
if (config.workers < 0 || config.workers > 256) {
|
||
throw MsgException("invalid workers: " + std::to_string(config.workers) + " (must be in [0, 256])");
|
||
}
|
||
detail::validateTimeout(config.recv_timeout);
|
||
if (config.max_connections < 1 || config.max_connections > 100000) {
|
||
throw MsgException("invalid max_connections: " + std::to_string(config.max_connections));
|
||
}
|
||
if (config.max_queue_size < 1 || config.max_queue_size > 1000000) {
|
||
throw MsgException("invalid max_queue_size: " + std::to_string(config.max_queue_size));
|
||
}
|
||
if (config.listen_backlog < 1 || config.listen_backlog > 4096) {
|
||
throw MsgException("invalid listen_backlog: " + std::to_string(config.listen_backlog));
|
||
}
|
||
if (config.poll_timeout <= 0 || config.poll_timeout > 5.0) {
|
||
throw MsgException("invalid poll_timeout: " + std::to_string(config.poll_timeout));
|
||
}
|
||
}
|
||
|
||
~MsgServer() { stop(); }
|
||
|
||
MsgServer(const MsgServer&) = delete;
|
||
MsgServer& operator=(const MsgServer&) = delete;
|
||
|
||
// ---- 启动（阻塞当前线程）------------------------------------------------
|
||
|
||
/** 启动服务器，阻塞当前线程直到 stop() 被调用或发生致命错误。
|
||
*
|
||
* @param on_message 消息回调，签名 void(int fd, const string& addr, const string& msg)。
|
||
* - 同步模式（workers=0）：回调在主线程执行，同一连接消息天然串行。
|
||
* - 线程池模式（workers>0）：回调提交到线程池并发执行；同一连接
|
||
* 上的消息仍严格串行（回调完成前该连接暂停 select 监听）。
|
||
* 回调内可安全调用 reply()/broadcast()/stop()。
|
||
*/
|
||
void start(MessageCallback on_message) {
|
||
// [libmsg-audit P2-8] 空回调会在 worker 线程 / RpcServer 分发中抛 std::bad_function_call，
|
||
// 提前校验并给出明确错误，避免行为不符合预期且错误信息含糊。
|
||
if (!on_message) {
|
||
throw MsgException("MsgServer::start(): on_message callback must not be null");
|
||
}
|
||
_on_message = std::move(on_message);
|
||
_running = true;
|
||
|
||
// 创建监听 socket
|
||
_server_fd.store(::socket(AF_INET, SOCK_STREAM, 0));
|
||
if (_server_fd.load() < 0) {
|
||
throw MsgConnectionException(std::string("socket() failed: ") + ::strerror(errno));
|
||
}
|
||
|
||
int optval = 1;
|
||
::setsockopt(_server_fd.load(), SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
|
||
|
||
struct sockaddr_in serv_addr {};
|
||
serv_addr.sin_family = AF_INET;
|
||
serv_addr.sin_port = htons(static_cast<uint16_t>(_port));
|
||
if (::inet_pton(AF_INET, _addr.c_str(), &serv_addr.sin_addr) <= 0) {
|
||
int fd = _server_fd.exchange(-1);
|
||
if (fd >= 0) ::close(fd);
|
||
throw MsgConnectionException("invalid bind address: " + _addr);
|
||
}
|
||
|
||
if (::bind(_server_fd.load(), reinterpret_cast<struct sockaddr*>(&serv_addr), sizeof(serv_addr)) < 0) {
|
||
int saved_errno = errno;
|
||
int fd = _server_fd.exchange(-1);
|
||
if (fd >= 0) ::close(fd);
|
||
throw MsgConnectionException(std::string("bind() failed: ") + ::strerror(saved_errno));
|
||
}
|
||
|
||
if (::listen(_server_fd.load(), _config.listen_backlog) < 0) {
|
||
int saved_errno = errno;
|
||
int fd = _server_fd.exchange(-1);
|
||
if (fd >= 0) ::close(fd);
|
||
throw MsgConnectionException(std::string("listen() failed: ") + ::strerror(saved_errno));
|
||
}
|
||
|
||
// 启动线程池（如果配置了 workers）
|
||
std::vector<std::thread> pool_threads;
|
||
if (_config.workers > 0) {
|
||
_pool_running = true;
|
||
pool_threads.reserve(_config.workers);
|
||
for (int i = 0; i < _config.workers; ++i) {
|
||
pool_threads.emplace_back(&MsgServer::_worker_loop, this);
|
||
}
|
||
}
|
||
|
||
// 主事件循环
|
||
try {
|
||
_event_loop();
|
||
} catch (...) {
|
||
// 确保线程池和连接被清理
|
||
_running = false;
|
||
_shutdown_pool(pool_threads);
|
||
_close_all_clients();
|
||
int fd = _server_fd.exchange(-1);
|
||
if (fd >= 0) { ::shutdown(fd, SHUT_RDWR); ::close(fd); }
|
||
throw;
|
||
}
|
||
|
||
_shutdown_pool(pool_threads);
|
||
_close_all_clients();
|
||
int fd = _server_fd.exchange(-1);
|
||
if (fd >= 0) { ::shutdown(fd, SHUT_RDWR); ::close(fd); }
|
||
}
|
||
|
||
// ---- 停止 --------------------------------------------------------------
|
||
|
||
/// 停止服务器（可从任意线程调用，包括回调内部）。
|
||
/// 关闭监听 socket 以解除 select 阻塞，事件循环退出前清理所有连接。
|
||
void stop() noexcept {
|
||
bool expected = true;
|
||
if (!_running.compare_exchange_strong(expected, false)) return;
|
||
int fd = _server_fd.exchange(-1);
|
||
if (fd >= 0) {
|
||
::shutdown(fd, SHUT_RDWR);
|
||
::close(fd);
|
||
}
|
||
}
|
||
|
||
// ---- 发送（线程安全）----------------------------------------------------
|
||
|
||
/** 向指定客户端发送一条应答消息（线程安全）。
|
||
*
|
||
* 经每连接发送锁串行化，workers>0 时可在回调线程安全调用。
|
||
* 如果连接已断开则静默忽略（fail-safe，不抛异常）。
|
||
*/
|
||
void reply(int sockfd, const std::string& msg) {
|
||
std::shared_ptr<ConnectionLock> conn_lock;
|
||
{
|
||
std::lock_guard<std::mutex> lock(_clients_mtx);
|
||
auto it = _send_locks.find(sockfd);
|
||
if (it == _send_locks.end()) return; // 连接已断开
|
||
conn_lock = it->second;
|
||
}
|
||
std::lock_guard<std::mutex> lock(conn_lock->send_mtx);
|
||
// [libmsg-audit P1-3] 检查 active：防止 _remove_client 先获取 send_mtx 并 close(fd) 后，
|
||
// reply 向已关闭/被 OS 复用的 fd 发送数据（CWE-362 反向竞态）。
|
||
if (!conn_lock->active.load()) return;
|
||
try {
|
||
pmsg_send(sockfd, msg);
|
||
} catch (const MsgException&) {
|
||
// [libmsg-audit P2-6] 发送失败（连接断开或发送超时），标记为 failed 由主循环清理。
|
||
// 与 broadcast() 保持一致，避免连接残留不清理（原 reply 静默忽略，连接滞留）。
|
||
std::lock_guard<std::mutex> flk(_failed_mtx);
|
||
_failed.insert(sockfd);
|
||
}
|
||
}
|
||
|
||
/** 向所有已连接客户端广播消息（线程安全）。
|
||
*
|
||
* @param msg 要广播的消息。
|
||
* @param except_fd 排除的 fd（通常是发送者自己），默认 -1 表示不排除。
|
||
*/
|
||
void broadcast(const std::string& msg, int except_fd = -1) {
|
||
// 先拷贝一份 fd + conn_lock 列表，避免持锁期间发送
|
||
std::vector<std::pair<int, std::shared_ptr<ConnectionLock>>> targets;
|
||
{
|
||
std::lock_guard<std::mutex> lock(_clients_mtx);
|
||
targets.reserve(_clients.size());
|
||
for (auto& [fd, addr] : _clients) {
|
||
if (fd == except_fd) continue;
|
||
auto it = _send_locks.find(fd);
|
||
if (it != _send_locks.end()) {
|
||
targets.emplace_back(fd, it->second);
|
||
}
|
||
}
|
||
}
|
||
for (auto& [fd, conn_lock] : targets) {
|
||
std::lock_guard<std::mutex> slock(conn_lock->send_mtx);
|
||
// [libmsg-audit P1-3] 检查 active：防止 _remove_client 先 close(fd) 后向复用 fd 发送（CWE-362 反向竞态）
|
||
if (!conn_lock->active.load()) continue;
|
||
try {
|
||
pmsg_send(fd, msg);
|
||
} catch (const MsgException&) {
|
||
// [libmsg-audit P2-6] 发送失败，标记为 failed 由主循环统一清理
|
||
std::lock_guard<std::mutex> flk(_failed_mtx);
|
||
_failed.insert(fd);
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---- 状态查询 ----------------------------------------------------------
|
||
|
||
/// 当前已连接客户端数量。
|
||
size_t client_count() const {
|
||
std::lock_guard<std::mutex> lock(_clients_mtx);
|
||
return _clients.size();
|
||
}
|
||
|
||
private:
|
||
// ---- 内部：线程池 -------------------------------------------------------
|
||
|
||
struct Task {
|
||
int fd;
|
||
std::string addr;
|
||
std::string message;
|
||
bool needs_recv = false; // true: worker 负责 pmsg_recv；false: 消息已就绪
|
||
};
|
||
|
||
void _worker_loop() {
|
||
while (true) {
|
||
Task task;
|
||
bool has_task = false;
|
||
{
|
||
std::unique_lock<std::mutex> lock(_pool_mtx);
|
||
_pool_cv.wait(lock, [this] { return !_pool_running || !_pool_queue.empty(); });
|
||
if (!_pool_queue.empty()) {
|
||
task = std::move(_pool_queue.front());
|
||
_pool_queue.pop();
|
||
has_task = true;
|
||
} else if (!_pool_running) {
|
||
return; // 队列空且已停止，退出
|
||
}
|
||
}
|
||
if (!has_task) continue;
|
||
|
||
// [libmsg-audit P1-1] workers > 0 模式下，消息接收在 worker 线程中完成（防事件循环阻塞 CWE-400），
|
||
// 避免攻击者以 max_connections 个半开连接阻塞事件循环长达数小时（Slowloris 变体）。
|
||
if (task.needs_recv) {
|
||
try {
|
||
task.message = pmsg_recv(task.fd, _config.recv_timeout);
|
||
} catch (const MsgException&) {
|
||
// 不在此清除 busy：fd 保持 busy 直到 _reap_failed 关闭连接。
|
||
// 旧实现提前清除 busy，EOF 可读性会让事件循环为同一 fd 重复
|
||
// 投递任务，迟到的失败标记经 fd 复用可能落在新连接上误杀之。
|
||
std::lock_guard<std::mutex> flk(_failed_mtx);
|
||
_failed.insert(task.fd);
|
||
continue;
|
||
}
|
||
}
|
||
|
||
// 执行回调
|
||
bool failed = false;
|
||
try {
|
||
_on_message(task.fd, task.addr, std::move(task.message));
|
||
} catch (const MsgConnectionException&) {
|
||
failed = true;
|
||
} catch (const std::exception& e) {
|
||
std::fprintf(stderr, "[MsgServer] callback error on fd=%d: %s\n",
|
||
task.fd, e.what());
|
||
std::fflush(stderr);
|
||
failed = true;
|
||
} catch (...) {
|
||
failed = true;
|
||
}
|
||
if (failed) {
|
||
// 失败连接交由 _reap_failed 统一清除 busy 并关闭，
|
||
// 不提前恢复 select 监听（同 recv 失败路径，防 fd 复用竞态）。
|
||
std::lock_guard<std::mutex> flk(_failed_mtx);
|
||
_failed.insert(task.fd);
|
||
} else {
|
||
// 清除 busy 标记，恢复 select 监听
|
||
std::lock_guard<std::mutex> lock(_busy_mtx);
|
||
_busy.erase(task.fd);
|
||
}
|
||
}
|
||
}
|
||
|
||
void _shutdown_pool(std::vector<std::thread>& pool_threads) {
|
||
// 停止接收新任务，唤醒所有 worker 排空队列后退出
|
||
_pool_running = false;
|
||
_pool_cv.notify_all();
|
||
for (auto& t : pool_threads) {
|
||
if (t.joinable()) t.join();
|
||
}
|
||
pool_threads.clear();
|
||
// worker 已排空队列，此时 busy 集合中的任务都已完成
|
||
// 任何仍标记为 failed 的连接由 _close_all_clients 统一清理
|
||
}
|
||
|
||
// ---- 内部：主事件循环 ---------------------------------------------------
|
||
|
||
void _event_loop() {
|
||
while (_running) {
|
||
// 清理 failed 连接
|
||
_reap_failed();
|
||
|
||
// 构建 fd_set：server_fd + 非 busy 的客户端 fd
|
||
fd_set readfds, exceptfds;
|
||
FD_ZERO(&readfds);
|
||
FD_ZERO(&exceptfds);
|
||
int maxfd = -1;
|
||
|
||
int srv_fd = _server_fd.load();
|
||
if (srv_fd >= 0) {
|
||
FD_SET(srv_fd, &readfds);
|
||
maxfd = srv_fd;
|
||
}
|
||
|
||
std::vector<int> watched_clients;
|
||
{
|
||
std::lock_guard<std::mutex> clk(_clients_mtx);
|
||
std::lock_guard<std::mutex> blk(_busy_mtx);
|
||
watched_clients.reserve(_clients.size());
|
||
for (auto& [fd, addr] : _clients) {
|
||
if (_busy.count(fd)) continue; // 回调处理中，暂停监听
|
||
FD_SET(fd, &readfds);
|
||
FD_SET(fd, &exceptfds);
|
||
if (fd > maxfd) maxfd = fd;
|
||
watched_clients.push_back(fd);
|
||
}
|
||
}
|
||
|
||
if (maxfd < 0) {
|
||
// 没有任何 fd 可监听（server_fd 已关闭且无客户端），退出
|
||
break;
|
||
}
|
||
|
||
struct timeval tv;
|
||
tv.tv_sec = 0;
|
||
tv.tv_usec = static_cast<long>(_config.poll_timeout * 1'000'000.0);
|
||
|
||
int ready = ::select(maxfd + 1, &readfds, nullptr, &exceptfds, &tv);
|
||
if (ready < 0) {
|
||
if (errno == EINTR) continue;
|
||
break; // EBADF 等致命错误
|
||
}
|
||
if (ready == 0) continue;
|
||
|
||
// 处理异常 fd
|
||
for (int fd : watched_clients) {
|
||
if (FD_ISSET(fd, &exceptfds)) {
|
||
_remove_client(fd);
|
||
}
|
||
}
|
||
|
||
// 处理新连接（重新加载 srv_fd，防止 stop() 期间关闭）
|
||
srv_fd = _server_fd.load();
|
||
if (srv_fd >= 0 && FD_ISSET(srv_fd, &readfds)) {
|
||
_accept_client();
|
||
}
|
||
|
||
// 处理客户端消息
|
||
for (int fd : watched_clients) {
|
||
if (FD_ISSET(fd, &readfds)) {
|
||
// 检查 fd 是否仍在客户端列表中（可能已被异常处理移除）
|
||
{
|
||
std::lock_guard<std::mutex> lock(_clients_mtx);
|
||
if (!_clients.count(fd)) continue;
|
||
}
|
||
_handle_client(fd);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
void _accept_client() {
|
||
struct sockaddr_in client_addr {};
|
||
socklen_t addrlen = sizeof(client_addr);
|
||
int srv_fd = _server_fd.load();
|
||
int client_fd = ::accept(srv_fd, reinterpret_cast<struct sockaddr*>(&client_addr), &addrlen);
|
||
if (client_fd < 0) {
|
||
// EINVAL/EBADF: server_fd 已被 stop() 关闭；其他错误忽略
|
||
return;
|
||
}
|
||
if (client_fd >= FD_SETSIZE) { ::close(client_fd); return; }
|
||
|
||
// [libmsg-audit P3-3] 禁用 Nagle 算法和设置发送超时（在加入 _clients 之前，确保连接配置完整），
|
||
// 先完成所有 socket 配置再暴露给其他线程，避免半配置连接被回复。
|
||
int flag = 1;
|
||
::setsockopt(client_fd, IPPROTO_TCP, TCP_NODELAY, &flag, sizeof(flag));
|
||
|
||
// [libmsg-audit P1-2] 设置发送超时：防止慢客户端不读数据导致发送缓冲区填满，
|
||
// 使 reply()/broadcast() 的 send() 永久阻塞（拒绝服务，CWE-400）。
|
||
// 超时后 sendAll() 抛 MsgTimeoutException，被 reply/broadcast 捕获并标记 failed 清理。
|
||
struct timeval sndtv {};
|
||
constexpr double kServerSendTimeout = 5.0; // 秒
|
||
sndtv.tv_sec = static_cast<time_t>(kServerSendTimeout);
|
||
sndtv.tv_usec = static_cast<long>(
|
||
(kServerSendTimeout - static_cast<double>(sndtv.tv_sec)) * 1'000'000.0);
|
||
::setsockopt(client_fd, SOL_SOCKET, SO_SNDTIMEO, &sndtv, sizeof(sndtv));
|
||
|
||
// 检查连接数上限并注册
|
||
{
|
||
std::lock_guard<std::mutex> lock(_clients_mtx);
|
||
if (static_cast<int>(_clients.size()) >= _config.max_connections) {
|
||
::close(client_fd);
|
||
return;
|
||
}
|
||
char ip_str[INET_ADDRSTRLEN];
|
||
::inet_ntop(AF_INET, &client_addr.sin_addr, ip_str, sizeof(ip_str));
|
||
std::string addr = std::string(ip_str) + ":" + std::to_string(ntohs(client_addr.sin_port));
|
||
_clients[client_fd] = addr;
|
||
_send_locks[client_fd] = std::make_shared<ConnectionLock>();
|
||
}
|
||
}
|
||
|
||
void _handle_client(int fd) {
|
||
std::string addr;
|
||
{
|
||
std::lock_guard<std::mutex> lock(_clients_mtx);
|
||
auto it = _clients.find(fd);
|
||
if (it == _clients.end()) return;
|
||
addr = it->second;
|
||
}
|
||
|
||
if (_config.workers == 0) {
|
||
// 同步模式：在主线程接收并执行回调
|
||
std::string message;
|
||
try {
|
||
message = pmsg_recv(fd, _config.recv_timeout);
|
||
} catch (const MsgException&) {
|
||
_remove_client(fd);
|
||
return;
|
||
}
|
||
try {
|
||
_on_message(fd, addr, std::move(message));
|
||
} catch (const MsgConnectionException&) {
|
||
_remove_client(fd);
|
||
} catch (const std::exception& e) {
|
||
std::fprintf(stderr, "[MsgServer] callback error on fd=%d: %s\n", fd, e.what());
|
||
std::fflush(stderr);
|
||
_remove_client(fd);
|
||
}
|
||
return;
|
||
}
|
||
|
||
// [libmsg-audit P1-1] 线程池模式：标记 busy，提交任务（由 worker 负责 recv + callback，
|
||
// 事件循环线程不阻塞接收，防 Slowloris DoS CWE-400）
|
||
{
|
||
std::lock_guard<std::mutex> blk(_busy_mtx);
|
||
if (_busy.count(fd)) return; // 已有在途任务，防同一 fd 重复投递（fd 复用竞态）
|
||
_busy.insert(fd);
|
||
}
|
||
{
|
||
std::lock_guard<std::mutex> plk(_pool_mtx);
|
||
if (static_cast<int>(_pool_queue.size()) >= _config.max_queue_size) {
|
||
// 队列已满：拒绝并断开连接（防 OOM）
|
||
std::lock_guard<std::mutex> blk(_busy_mtx);
|
||
_busy.erase(fd);
|
||
_remove_client(fd);
|
||
return;
|
||
}
|
||
_pool_queue.push(Task{fd, std::move(addr), {}, true});
|
||
}
|
||
_pool_cv.notify_one();
|
||
}
|
||
|
||
void _reap_failed() {
|
||
std::unordered_set<int> failed;
|
||
{
|
||
std::lock_guard<std::mutex> flk(_failed_mtx);
|
||
failed.swap(_failed);
|
||
}
|
||
for (int fd : failed) {
|
||
{
|
||
std::lock_guard<std::mutex> blk(_busy_mtx);
|
||
_busy.erase(fd);
|
||
}
|
||
_remove_client(fd);
|
||
}
|
||
}
|
||
|
||
void _remove_client(int fd) {
|
||
std::shared_ptr<ConnectionLock> conn_lock;
|
||
{
|
||
std::lock_guard<std::mutex> clk(_clients_mtx);
|
||
if (_clients.erase(fd) == 0) return; // 已被移除，防止 double-close
|
||
auto it = _send_locks.find(fd);
|
||
if (it != _send_locks.end()) {
|
||
conn_lock = it->second;
|
||
_send_locks.erase(it);
|
||
}
|
||
}
|
||
// [libmsg-audit P0-1 / P1-3] 持有 send_mtx 后设置 active=false 再 close(fd)：
|
||
// 1. 防止 broadcast/reply 在 fd 关闭后向 OS 复用的新连接发送数据（CWE-362，第一轮 P0-1）
|
||
// 2. active=false 让 reply/broadcast 在获取锁后得知 fd 已关闭，放弃发送（第四轮 P1-3 反向竞态）
|
||
if (conn_lock) {
|
||
std::lock_guard<std::mutex> slk(conn_lock->send_mtx);
|
||
conn_lock->active.store(false);
|
||
::shutdown(fd, SHUT_RDWR);
|
||
::close(fd);
|
||
} else {
|
||
::shutdown(fd, SHUT_RDWR);
|
||
::close(fd);
|
||
}
|
||
}
|
||
|
||
void _close_all_clients() {
|
||
std::vector<std::pair<int, std::shared_ptr<ConnectionLock>>> to_close;
|
||
{
|
||
std::lock_guard<std::mutex> clk(_clients_mtx);
|
||
to_close.reserve(_clients.size());
|
||
for (auto& [fd, addr] : _clients) {
|
||
auto it = _send_locks.find(fd);
|
||
if (it != _send_locks.end()) {
|
||
to_close.emplace_back(fd, it->second);
|
||
} else {
|
||
to_close.emplace_back(fd, nullptr);
|
||
}
|
||
}
|
||
_clients.clear();
|
||
_send_locks.clear();
|
||
}
|
||
// [libmsg-audit P0-2 / P1-3] 持有每个连接的 send_mtx 后设置 active=false 再 close(fd)：
|
||
// 防止 stop()/析构时 reply/broadcast 正在发送，导致数据发到 OS 复用的新连接（CWE-362）。
|
||
// 第二轮 P0-2 修复批量关闭路径的竞态回归，第四轮 P1-3 的 active 标志堵住反向竞态。
|
||
for (auto& [fd, conn_lock] : to_close) {
|
||
if (conn_lock) {
|
||
std::lock_guard<std::mutex> slk(conn_lock->send_mtx);
|
||
conn_lock->active.store(false);
|
||
::shutdown(fd, SHUT_RDWR);
|
||
::close(fd);
|
||
} else {
|
||
::shutdown(fd, SHUT_RDWR);
|
||
::close(fd);
|
||
}
|
||
}
|
||
std::lock_guard<std::mutex> blk(_busy_mtx);
|
||
_busy.clear();
|
||
std::lock_guard<std::mutex> flk(_failed_mtx);
|
||
_failed.clear();
|
||
}
|
||
|
||
// ---- 成员 --------------------------------------------------------------
|
||
|
||
std::string _addr;
|
||
int _port;
|
||
MsgServerConfig _config;
|
||
std::atomic<int> _server_fd; ///< [libmsg-audit P2-7] 监听 socket fd（atomic 防止 stop() 与 _event_loop 数据竞争 UB，CWE-362/662）
|
||
std::atomic<bool> _running;
|
||
MessageCallback _on_message;
|
||
|
||
// 客户端管理
|
||
mutable std::mutex _clients_mtx;
|
||
std::unordered_map<int, std::string> _clients; ///< fd -> "ip:port"
|
||
std::unordered_map<int, std::shared_ptr<ConnectionLock>> _send_locks; ///< fd -> 发送锁+活跃标志
|
||
|
||
// 同连接串行
|
||
std::mutex _busy_mtx;
|
||
std::unordered_set<int> _busy; ///< 回调执行中的 fd
|
||
|
||
// 失败连接（worker 线程标记，主循环清理）
|
||
std::mutex _failed_mtx;
|
||
std::unordered_set<int> _failed;
|
||
|
||
// 线程池
|
||
std::atomic<bool> _pool_running{false};
|
||
std::mutex _pool_mtx;
|
||
std::condition_variable _pool_cv;
|
||
std::queue<Task> _pool_queue;
|
||
};
|
||
|
||
} // namespace net
|
||
} // namespace pmc
|