fix(pymsg): 增加连接数与线程池队列上限,防资源耗尽

- P1-1: MsgServer 新增 max_connections(默认1000),_accept_client 超限
  直接拒绝并关闭新连接,防连接耗尽攻击
- P1-2: 新增 _BoundedThreadPoolExecutor + max_queue_size(默认1000),
  慢 handler 时任务排队超限即拒绝并断开连接,防内存无限增长
同步更新审计报告 P1-1/P1-2 状态,修复率 18.75% -> 25%
This commit is contained in:
qing committed 2026-08-03 22:25:32 +08:00
1 parent 47ab9f2fe6
commit 87f209af73
2 files changed
+68 -21

No files matched your search

+14 -14
View File
@@ -16,10 +16,10 @@
| 等级 | 总数 | 已修复 | 部分修复 | 未修复 | 修复率 | | 等级 | 总数 | 已修复 | 部分修复 | 未修复 | 修复率 |
|------|------|--------|----------|--------|--------| |------|------|--------|----------|--------|--------|
| P0 | 3 | 2 | 0 | 1 | 67% | | P0 | 3 | 2 | 0 | 1 | 67% |
| P1 | 5 | 0 | 0 | 5 | 0% | | P1 | 5 | 2 | 0 | 3 | 40% |
| P2 | 4 | 0 | 0 | 4 | 0% | | P2 | 4 | 0 | 0 | 4 | 0% |
| P3 | 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-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-2 | ✅ 已修复 | `__init__.py` L113-114, L462-463 | 双重防护:`recv()` 将 `UnicodeDecodeError` 包装为 `ConnectionError`;`_handle_client` 防御性捕获 `UnicodeDecodeError` |
| P0-3 | ❌ 未修复 | `rpc.py` L80-171 | 无认证机制,仅依赖 127.0.0.1 网络层隔离 | | P0-3 | ❌ 未修复 | `rpc.py` L80-171 | 无认证机制,仅依赖 127.0.0.1 网络层隔离 |
| P1-1 | ❌ 未修复 | `__init__.py` L359 | `listen(10)` 仅限制 SYN 队列,不限制已建立连接数 | | 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` L362-365 | `ThreadPoolExecutor` 使用默认无界队列 | | 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-3 | ❌ 未修复 | `rpc.py` L160, L178 | `_dispatch` 将 `str(e)` 和 `f"JSON 解析失败: {e}"` 直接返回客户端 |
| P1-4 | ❌ 未修复 | 全模块 | 明文 TCP,无 TLS 支持 | | P1-4 | ❌ 未修复 | 全模块 | 明文 TCP,无 TLS 支持 |
| P1-5 | ❌ 未修复 | `rpc.py` L158 | `json.loads(text)` 无深度限制 | | P1-5 | ❌ 未修复 | `rpc.py` L158 | `json.loads(text)` 无深度限制 |
@@ -60,7 +60,7 @@
| 优先级 | 总数 | 已修复 | 待修复 | | 优先级 | 总数 | 已修复 | 待修复 |
|--------|------|--------|--------| |--------|------|--------|--------|
| P0 | 3 | 2 | 1 | | P0 | 3 | 2 | 1 |
| P1 | 5 | 0 | 5 | | P1 | 5 | 2 | 3 |
| P2 | 4 | 0 | 4 | | P2 | 4 | 0 | 4 |
| P3 | 4 | 0 | 4 | | P3 | 4 | 0 | 4 |
@@ -105,13 +105,13 @@
## 🟠 P1 — 高危问题 ## 🟠 P1 — 高危问题
| # | 问题 | 位置 | | # | 问题 | 位置 | 状态 |
|---|------|------| |---|------|------|------|
| P1-1 | 无连接数限制,连接耗尽攻击 | `__init__.py:334` `listen(10)` 不限制已建立连接数 | | P1-1 | 无连接数限制,连接耗尽攻击 | `__init__.py:342-350,467-484` | ✅ 已修复(`max_connections` 默认 1000) |
| P1-2 | 线程池无队列限制,慢 handler 耗尽内存 | `__init__.py:337-340` ThreadPoolExecutor 默认无界队列 | | 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-3 | 异常信息直接返回客户端,泄露内部细节 | `rpc.py:144,162` 直接将 `str(e)` 返回调用方 | ❌ 未修复 |
| P1-4 | 无 TLS/加密传输,内网流量可被窃听 | 全模块明文 TCP | | P1-4 | 无 TLS/加密传输,内网流量可被窃听 | 全模块明文 TCP | ❌ 未修复 |
| P1-5 | JSON 解析无深度限制,深度嵌套导致栈溢出/CPU 100% | `recv()` 后 `json.loads()` 无限制 | | P1-5 | JSON 解析无深度限制,深度嵌套导致栈溢出/CPU 100% | `recv()` 后 `json.loads()` 无限制 | ❌ 未修复 |
--- ---
@@ -168,8 +168,8 @@
- [x] 修复 P0-2: 捕获 UnicodeDecodeError/JSONDecodeError ✅ 已完成 - [x] 修复 P0-2: 捕获 UnicodeDecodeError/JSONDecodeError ✅ 已完成
### 短期(72小时) ### 短期(72小时)
- [ ] 添加最大连接数限制 - [x] 添加最大连接数限制 ✅ 已完成
- [ ] 线程池改用有界队列 - [x] 线程池改用有界队列 ✅ 已完成
- [ ] 异常信息脱敏 - [ ] 异常信息脱敏
- [ ] JSON 解析深度限制 - [ ] JSON 解析深度限制
+54 -7
View File
@@ -18,6 +18,7 @@ Usage:
""" """
import contextlib import contextlib
import queue as _queue
import select import select
import socket as _socket import socket as _socket
import struct import struct
@@ -312,6 +313,26 @@ class MsgClient:
self._sock = None 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: class MsgServer:
"""基于 pymsg 长度前缀协议的多客户端 TCP 服务器。 """基于 pymsg 长度前缀协议的多客户端 TCP 服务器。
@@ -335,16 +356,26 @@ class MsgServer:
``workers > 0`` 时回调提交到线程池并发执行(I/O 密集回调可获得真实 ``workers > 0`` 时回调提交到线程池并发执行(I/O 密集回调可获得真实
并行度)。同一连接上的消息仍严格串行处理;``reply()`` / ``broadcast()`` 并行度)。同一连接上的消息仍严格串行处理;``reply()`` / ``broadcast()``
经每连接发送锁串行化,不会发生帧交错。 经每连接发送锁串行化,不会发生帧交错。
``max_connections`` 限制同时建立的活跃连接数,超过上限的新连接被直接
拒绝并关闭,防止连接耗尽攻击。
``max_queue_size`` 限制线程池待办队列深度(``workers > 0`` 时),慢
handler 导致任务排队超限时拒绝并断开该连接,防止内存无限增长。
""" """
_RECV_TIMEOUT = 10.0 # 服务端接收单条消息的超时(秒),防止 Slowloris 式半开连接卡死主线程 _RECV_TIMEOUT = 10.0 # 服务端接收单条消息的超时(秒),防止 Slowloris 式半开连接卡死主线程
def __init__(self, addr: str, port: int, *, workers: int = 0, 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.addr = addr
self.port = port self.port = port
self._workers = workers self._workers = workers
self._recv_timeout = recv_timeout self._recv_timeout = recv_timeout
self._max_connections = max_connections
self._max_queue_size = max_queue_size
self._running = False self._running = False
self._server_sock: "_socket.socket | None" = None self._server_sock: "_socket.socket | None" = None
self._clients: "dict" = {} # socket -> (ip, port) self._clients: "dict" = {} # socket -> (ip, port)
@@ -375,8 +406,9 @@ class MsgServer:
servsock.listen(10) servsock.listen(10)
self._server_sock = servsock self._server_sock = servsock
if self._workers > 0: if self._workers > 0:
self._pool = ThreadPoolExecutor( self._pool = _BoundedThreadPoolExecutor(
max_workers=self._workers, max_workers=self._workers,
max_queue_size=self._max_queue_size,
thread_name_prefix=f"pymsg-{self.port}", thread_name_prefix=f"pymsg-{self.port}",
) )
@@ -465,10 +497,19 @@ class MsgServer:
def _accept_client(self, servsock, socketlist): def _accept_client(self, servsock, socketlist):
try: try:
client_sock, client_addr = servsock.accept() client_sock, client_addr = servsock.accept()
if client_sock.fileno() != -1: if client_sock.fileno() == -1:
socketlist.append(client_sock) return
self._clients[client_sock] = client_addr # 活跃连接数上限控制:超过 max_connections 直接拒绝,
self._send_locks[client_sock] = threading.Lock() # 防止连接耗尽攻击(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: except OSError:
pass pass
@@ -493,7 +534,13 @@ class MsgServer:
# 再把回调提交到线程池执行 # 再把回调提交到线程池执行
with self._lock: with self._lock:
self._busy.add(s) 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): def _run_task(self, s, addr, text, on_message):
"""工作线程:执行用户回调。 """工作线程:执行用户回调。