From 87f209af73d3ba3c3544298c9be2b8bd036d5ff1 Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Mon, 3 Aug 2026 22:25:32 +0800 Subject: [PATCH] =?UTF-8?q?fix(pymsg):=20=E5=A2=9E=E5=8A=A0=E8=BF=9E?= =?UTF-8?q?=E6=8E=A5=E6=95=B0=E4=B8=8E=E7=BA=BF=E7=A8=8B=E6=B1=A0=E9=98=9F?= =?UTF-8?q?=E5=88=97=E4=B8=8A=E9=99=90=EF=BC=8C=E9=98=B2=E8=B5=84=E6=BA=90?= =?UTF-8?q?=E8=80=97=E5=B0=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - P1-1: MsgServer 新增 max_connections（默认1000），_accept_client 超限 直接拒绝并关闭新连接，防连接耗尽攻击 - P1-2: 新增 _BoundedThreadPoolExecutor + max_queue_size（默认1000）， 慢 handler 时任务排队超限即拒绝并断开连接，防内存无限增长 同步更新审计报告 P1-1/P1-2 状态，修复率 18.75% -> 25% --- .../security/pymsg_SECURITY_AUDIT_20260801.md | 28 ++++----- pymsg/pymsg/__init__.py | 61 ++++++++++++++++--- 2 files changed, 68 insertions(+), 21 deletions(-) diff --git a/docs/security/pymsg_SECURITY_AUDIT_20260801.md b/docs/security/pymsg_SECURITY_AUDIT_20260801.md index 825c0ae..ee47987 100644 --- a/docs/security/pymsg_SECURITY_AUDIT_20260801.md +++ b/docs/security/pymsg_SECURITY_AUDIT_20260801.md @@ -16,10 +16,10 @@ | 等级 | 总数 | 已修复 | 部分修复 | 未修复 | 修复率 | |------|------|--------|----------|--------|--------| | P0 | 3 | 2 | 0 | 1 | 67% | -| P1 | 5 | 0 | 0 | 5 | 0% | +| P1 | 5 | 2 | 0 | 3 | 40% | | P2 | 4 | 0 | 0 | 4 | 0% | | P3 | 4 | 0 | 0 | 4 | 0% | -| **合计** | **16** | **2** | **0** | **14** | **12.5%** | +| **合计** | **16** | **4** | **0** | **12** | **25%** | ### 逐条验证 @@ -28,8 +28,8 @@ | P0-1 | ✅ 已修复 | `__init__.py` L103-120, L140-170 | `recv()` 改为 deadline 模式：`_recv_exact` 接收 `time.monotonic()+timeout` 作为总截止时刻，每次 `recv_into` 前按剩余时间收紧 socket 超时；header 与 payload 共享同一总超时，逐字节发送无法刷新 deadline，半开连接被及时关闭（已实测：0.6s 超时 0.60s 即抛 TimeoutError） | | P0-2 | ✅ 已修复 | `__init__.py` L113-114, L462-463 | 双重防护：`recv()` 将 `UnicodeDecodeError` 包装为 `ConnectionError`；`_handle_client` 防御性捕获 `UnicodeDecodeError` | | P0-3 | ❌ 未修复 | `rpc.py` L80-171 | 无认证机制，仅依赖 127.0.0.1 网络层隔离 | -| P1-1 | ❌ 未修复 | `__init__.py` L359 | `listen(10)` 仅限制 SYN 队列，不限制已建立连接数 | -| P1-2 | ❌ 未修复 | `__init__.py` L362-365 | `ThreadPoolExecutor` 使用默认无界队列 | +| P1-1 | ✅ 已修复 | `__init__.py` L342-350, L467-484 | `MsgServer` 新增 `max_connections` 参数（默认 1000），`_accept_client` 在 `len(_clients) >= max_connections` 时直接拒绝并关闭新连接；已实测 max_connections=3 时第 4 个连接被关闭（ConnectionResetError），handler 只处理前 3 条 | +| P1-2 | ✅ 已修复 | `__init__.py` L316-333, L369, L405-410, L530-540 | 新增 `_BoundedThreadPoolExecutor`（`submit` 时检查 `_work_queue.qsize()`，超限抛 `queue.Full`）；`MsgServer` 新增 `max_queue_size` 参数（默认 1000），`_handle_client` 捕获 `queue.Full` 后拒绝并断开该连接；已实测 workers=1 + max_queue_size=1 时，慢 handler 阻塞下第 3 个连接被拒绝关闭，handler 只处理前 2 个任务 | | P1-3 | ❌ 未修复 | `rpc.py` L160, L178 | `_dispatch` 将 `str(e)` 和 `f"JSON 解析失败: {e}"` 直接返回客户端 | | P1-4 | ❌ 未修复 | 全模块 | 明文 TCP，无 TLS 支持 | | P1-5 | ❌ 未修复 | `rpc.py` L158 | `json.loads(text)` 无深度限制 | @@ -60,7 +60,7 @@ | 优先级 | 总数 | 已修复 | 待修复 | |--------|------|--------|--------| | P0 | 3 | 2 | 1 | -| P1 | 5 | 0 | 5 | +| P1 | 5 | 2 | 3 | | P2 | 4 | 0 | 4 | | P3 | 4 | 0 | 4 | @@ -105,13 +105,13 @@ ## 🟠 P1 — 高危问题 -| # | 问题 | 位置 | -|---|------|------| -| P1-1 | 无连接数限制，连接耗尽攻击 | `__init__.py:334` `listen(10)` 不限制已建立连接数 | -| P1-2 | 线程池无队列限制，慢 handler 耗尽内存 | `__init__.py:337-340` ThreadPoolExecutor 默认无界队列 | -| P1-3 | 异常信息直接返回客户端，泄露内部细节 | `rpc.py:144,162` 直接将 `str(e)` 返回调用方 | -| P1-4 | 无 TLS/加密传输，内网流量可被窃听 | 全模块明文 TCP | -| P1-5 | JSON 解析无深度限制，深度嵌套导致栈溢出/CPU 100% | `recv()` 后 `json.loads()` 无限制 | +| # | 问题 | 位置 | 状态 | +|---|------|------|------| +| P1-1 | 无连接数限制，连接耗尽攻击 | `__init__.py:342-350,467-484` | ✅ 已修复（`max_connections` 默认 1000） | +| P1-2 | 线程池无队列限制，慢 handler 耗尽内存 | `__init__.py:316-333,369,405-410,530-540` | ✅ 已修复（`_BoundedThreadPoolExecutor` + `max_queue_size` 默认 1000） | +| P1-3 | 异常信息直接返回客户端，泄露内部细节 | `rpc.py:144,162` 直接将 `str(e)` 返回调用方 | ❌ 未修复 | +| P1-4 | 无 TLS/加密传输，内网流量可被窃听 | 全模块明文 TCP | ❌ 未修复 | +| P1-5 | JSON 解析无深度限制，深度嵌套导致栈溢出/CPU 100% | `recv()` 后 `json.loads()` 无限制 | ❌ 未修复 | --- @@ -168,8 +168,8 @@ - [x] 修复 P0-2: 捕获 UnicodeDecodeError/JSONDecodeError ✅ 已完成 ### 短期（72小时） -- [ ] 添加最大连接数限制 -- [ ] 线程池改用有界队列 +- [x] 添加最大连接数限制 ✅ 已完成 +- [x] 线程池改用有界队列 ✅ 已完成 - [ ] 异常信息脱敏 - [ ] JSON 解析深度限制 diff --git a/pymsg/pymsg/__init__.py b/pymsg/pymsg/__init__.py index f6984c6..849d239 100644 --- a/pymsg/pymsg/__init__.py +++ b/pymsg/pymsg/__init__.py @@ -18,6 +18,7 @@ Usage: """ import contextlib +import queue as _queue import select import socket as _socket import struct @@ -312,6 +313,26 @@ class MsgClient: self._sock = None +class _BoundedThreadPoolExecutor(ThreadPoolExecutor): + """有界线程池：提交队列有上限，满时抛 ``queue.Full``。 + + 标准 ``ThreadPoolExecutor`` 使用无界队列，慢 handler 时新任务会无限 + 堆积在内存中直至 OOM。此子类在 ``submit`` 时检查排队深度，超限即抛 + ``queue.Full``，由调用方拒绝并断开该连接，防止内存无限增长。 + """ + + def __init__(self, max_workers: int, max_queue_size: int, + thread_name_prefix: str) -> None: + super().__init__(max_workers=max_workers, + thread_name_prefix=thread_name_prefix) + self._max_queue_size = max_queue_size + + def submit(self, fn, *args, **kwargs): + if self._work_queue.qsize() >= self._max_queue_size: + raise _queue.Full("thread pool task queue is full") + return super().submit(fn, *args, **kwargs) + + class MsgServer: """基于 pymsg 长度前缀协议的多客户端 TCP 服务器。 @@ -335,16 +356,26 @@ class MsgServer: ``workers > 0`` 时回调提交到线程池并发执行（I/O 密集回调可获得真实 并行度）。同一连接上的消息仍严格串行处理；``reply()`` / ``broadcast()`` 经每连接发送锁串行化，不会发生帧交错。 + + ``max_connections`` 限制同时建立的活跃连接数，超过上限的新连接被直接 + 拒绝并关闭，防止连接耗尽攻击。 + + ``max_queue_size`` 限制线程池待办队列深度（``workers > 0`` 时），慢 + handler 导致任务排队超限时拒绝并断开该连接，防止内存无限增长。 """ _RECV_TIMEOUT = 10.0 # 服务端接收单条消息的超时（秒），防止 Slowloris 式半开连接卡死主线程 def __init__(self, addr: str, port: int, *, workers: int = 0, - recv_timeout: float = _RECV_TIMEOUT) -> None: + recv_timeout: float = _RECV_TIMEOUT, + max_connections: int = 1000, + max_queue_size: int = 1000) -> None: self.addr = addr self.port = port self._workers = workers self._recv_timeout = recv_timeout + self._max_connections = max_connections + self._max_queue_size = max_queue_size self._running = False self._server_sock: "_socket.socket | None" = None self._clients: "dict" = {} # socket -> (ip, port) @@ -375,8 +406,9 @@ class MsgServer: servsock.listen(10) self._server_sock = servsock if self._workers > 0: - self._pool = ThreadPoolExecutor( + self._pool = _BoundedThreadPoolExecutor( max_workers=self._workers, + max_queue_size=self._max_queue_size, thread_name_prefix=f"pymsg-{self.port}", ) @@ -465,10 +497,19 @@ class MsgServer: def _accept_client(self, servsock, socketlist): try: client_sock, client_addr = servsock.accept() - if client_sock.fileno() != -1: - socketlist.append(client_sock) - self._clients[client_sock] = client_addr - self._send_locks[client_sock] = threading.Lock() + if client_sock.fileno() == -1: + return + # 活跃连接数上限控制：超过 max_connections 直接拒绝， + # 防止连接耗尽攻击（listen backlog 只限制 SYN 队列，不限制已建立连接数） + if len(self._clients) >= self._max_connections: + try: + client_sock.close() + except OSError: + pass + return + socketlist.append(client_sock) + self._clients[client_sock] = client_addr + self._send_locks[client_sock] = threading.Lock() except OSError: pass @@ -493,7 +534,13 @@ class MsgServer: # 再把回调提交到线程池执行 with self._lock: self._busy.add(s) - self._pool.submit(self._run_task, s, addr, text, on_message) + try: + self._pool.submit(self._run_task, s, addr, text, on_message) + except _queue.Full: + # 有界队列已满：慢 handler 时任务不应无限堆积，拒绝并断开该连接 + with self._lock: + self._busy.discard(s) + self._remove_client(s, socketlist) def _run_task(self, s, addr, text, on_message): """工作线程：执行用户回调。