"""pymsg - Length-prefixed variable-length socket message framing. Protocol: [4-byte big-endian uint32 length][UTF-8 payload] Usage: import socket from pymsg import send, recv # sender send(sock, "hello world") # receiver msg = recv(sock) # with timeout (seconds) send(sock, "hello", timeout=5.0) msg = recv(sock, timeout=5.0) """ import contextlib import queue as _queue import select import socket as _socket import struct import threading import time import traceback from concurrent.futures import ThreadPoolExecutor __all__ = ["send", "recv", "TimeoutError", "MsgClient", "MsgServer"] _HEADER = struct.Struct("!I") # 4 bytes, big-endian unsigned int # 消息大小安全上限（64 MiB），与 C++ 端 Msg.hpp 保持一致。 # 防止恶意客户端发送超大长度头导致 OOM。 _MAX_MSG_SIZE = 64 * 1024 * 1024 class TimeoutError(Exception): """Raised when a send or recv operation exceeds its timeout.""" @contextlib.contextmanager def _timeout(sock: "_socket.socket", seconds: "float | None"): """Temporarily set the socket timeout, restoring it on exit.""" if seconds is None: yield return original = sock.gettimeout() sock.settimeout(seconds) try: yield finally: sock.settimeout(original) def send( sock: "_socket.socket", data: str, *, timeout: "float | None" = None, ) -> None: """Send a string over a socket using length-prefixed framing. Args: sock: A connected socket object. data: The string message to send. timeout: Optional timeout in seconds. Raises TimeoutError on expiry. """ payload = data.encode("utf-8") if len(payload) > _MAX_MSG_SIZE: raise ValueError( f"payload too large: {len(payload)} bytes (max {_MAX_MSG_SIZE})" ) header = _HEADER.pack(len(payload)) with _timeout(sock, timeout): try: _send_all(sock, header, payload) except _socket.timeout: raise TimeoutError( f"send timed out after {timeout}s" ) from None def recv( sock: "_socket.socket", *, timeout: "float | None" = None, ) -> str: """Receive a length-prefixed string message from a socket. Args: sock: A connected socket object. timeout: Optional timeout in seconds. Raises TimeoutError on expiry. Returns: The decoded string message. Raises: ConnectionError: If the connection is closed before a complete message is received. TimeoutError: If the operation exceeds the timeout. """ # [SECURITY AUDIT P0-1 ✅ 已修复] recv 总超时（deadline 模式） # 原问题：sock.settimeout() 是「单次系统调用」超时而非整条消息总超时， # 攻击者可每次在快超时时发 1 字节绕过，造成 Slowloris 型连接耗尽 DoS。 # 修复：header 与 payload 共享同一总截止时刻 deadline，每次 recv_into # 前按剩余时间收紧 socket 超时，半开连接被及时关闭（实测 0.6s 超时即终止）。 # 详见 docs/security/pymsg_SECURITY_AUDIT_20260801.md deadline = None if timeout is not None: deadline = time.monotonic() + timeout try: header = _recv_exact(sock, _HEADER.size, deadline) (length,) = _HEADER.unpack(header) if length > _MAX_MSG_SIZE: raise ConnectionError( f"message too large: {length} bytes (max {_MAX_MSG_SIZE})" ) payload = _recv_exact(sock, length, deadline) return payload.decode("utf-8") # [SECURITY AUDIT P0-2 ✅ 已修复] UnicodeDecodeError 导致同步服务崩溃 # 原问题：payload.decode("utf-8") 抛出的 UnicodeDecodeError（ValueError 子类） # 不在 _handle_client 捕获列表中，同步模式收到无效 UTF-8 会令整个进程崩溃。 # 修复（双重防护）：此处将 UnicodeDecodeError 包装为 ConnectionError（第一层）， # _handle_client 再防御性捕获兜底（第二层）。db/confs/smtp/captcha 等服务均受保护。 except UnicodeDecodeError: raise ConnectionError("invalid UTF-8 payload") from None except (TimeoutError, _socket.timeout): raise TimeoutError( f"recv timed out after {timeout}s" ) from None def _send_all(sock: "_socket.socket", header: bytes, payload: bytes) -> None: """通过 sendmsg 散射-聚集 I/O 发送 header + payload，避免拼接拷贝。 内核直接从两块独立内存读取数据，无需在用户态拼接 header + payload。 处理 partial send：sendmsg 可能只发出部分数据，需循环补发。 """ header_len = len(header) total = header_len + len(payload) sent = sock.sendmsg([memoryview(header), memoryview(payload)]) while sent < total: if sent < header_len: sent += sock.sendmsg( [memoryview(header)[sent:], memoryview(payload)] ) else: offset = sent - header_len sent += sock.sendmsg([memoryview(payload)[offset:]]) def _recv_exact( sock: "_socket.socket", n: int, deadline: "float | None" = None, ) -> bytearray: """Read exactly n bytes from sock, raising on premature close. 使用 recv_into 直接写入预分配 buffer，避免 bytearray 逐步增长和 最终 bytes() 转换的额外拷贝。 deadline 模式：传入 ``time.monotonic() + 超时`` 作为接收的**总截止时刻**， 每次 recv_into 前按剩余时间收紧 socket 超时。若剩余时间耗尽则抛 ``TimeoutError``。这防止 Slowloris 型攻击：逐字节发送会不断刷新单次 ``sock.settimeout()``，但无法刷新总 deadline，半开连接会被及时关闭。 """ # [SECURITY AUDIT P0-1 ✅ 已修复] 本函数为 P0-1 的核心修复实现， # 与上方 recv() 的 deadline 计算配合，构成整条消息的总超时防护。 # 详见 docs/security/pymsg_SECURITY_AUDIT_20260801.md buf = bytearray(n) view = memoryview(buf) got = 0 while got < n: if deadline is not None: remaining = deadline - time.monotonic() if remaining <= 0: raise TimeoutError("recv timed out") sock.settimeout(remaining) chunk = sock.recv_into(view[got:], n - got) if chunk == 0: raise ConnectionError( f"connection closed: expected {n} bytes, got {got}" ) got += chunk return buf # --------------------------------------------------------------------------- # 高层类：封装基于 pymsg 的客户端 / 服务器 # --------------------------------------------------------------------------- class MsgClient: """TCP 聊天客户端：封装套接字创建、连接、发送、监听与关闭。 ``listen()`` 在调用线程内阻塞运行，可另开线程调用 ``stop()`` 来终止。 Usage:: def on_msg(text: str) -> None: print(f"收到: {text}") c = MsgClient("localhost", 8003, "Alice") c.connect() c.send("hello") # 发送单条消息 c.listen(on_msg) # 阻塞当前线程，持续回调 """ def __init__(self, addr: str, port: int, name: str) -> None: self.addr = addr self.port = port self.name = name self._sock: "_socket.socket | None" = None self._running = False # ---- 连接 ------------------------------------------------------------ def connect(self) -> None: """创建 TCP 套接字并连接到服务器。""" self._sock = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM) self._sock.connect((self.addr, self.port)) # ---- 发送单条消息（短连接，不等待回复）-------------------------------- def send(self, msg: str) -> None: """发送一条消息到服务器（使用短连接，pymsg 协议）。""" try: with _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM) as s: s.settimeout(3) s.connect((self.addr, self.port)) send(s, f"{self.name}: {msg}") except ConnectionRefusedError: raise ConnectionRefusedError( f"Could not connect to {self.addr}:{self.port}" ) # ---- 请求-应答（短连接）------------------------------------------------ def call(self, msg: str, *, timeout: "float | None" = None, retry_on_refused: bool = False, retry_interval: float = 1.0, _log_every: float = 5.0) -> str: """发送一条消息并阻塞等待单条回复（短连接：connect → send → recv → close）。 适合请求-应答模式，如 RPC 调用。不会给消息添加 name 前缀。 Args: msg: 要发送的原始字符串。 timeout: 超时秒数，默认 10 秒。 retry_on_refused: 目标服务未启动（ConnectionRefusedError）时是否阻塞重试， 直到连接成功。适用于启动顺序不确定的服务发现场景。 默认 False（连接拒绝立即抛异常）。 retry_interval: 重试间隔秒数，默认 1.0。 _log_every: 重试期间打印等待提示的间隔（秒），默认每 5 秒打印一次，避免刷屏。 Returns: str: 服务器返回的字符串。 Raises: ConnectionRefusedError: 目标服务不可达且 ``retry_on_refused=False``。 TimeoutError: 等待响应超时。 ConnectionError: 连接在收全数据前断开。 """ _tout = timeout if timeout is not None else 10.0 _last_log = 0.0 _attempt = 0 while True: try: with _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM) as s: s.settimeout(_tout) s.connect((self.addr, self.port)) send(s, msg) return recv(s, timeout=_tout) except ConnectionRefusedError: if not retry_on_refused: raise ConnectionRefusedError( f"Could not connect to {self.addr}:{self.port}" ) _attempt += 1 now = time.monotonic() if now - _last_log >= _log_every or _attempt == 1: print( f"[pymsg] 等待服务 {self.addr}:{self.port} 就绪... " f"(已重试 {_attempt} 次)", file=__import__('sys').stderr, ) _last_log = now time.sleep(retry_interval) # ---- 持续监听（阻塞当前线程，回调模式）------------------------------- def listen(self, on_message: "callable") -> None: """阻塞当前线程，持续接收服务器消息并回调。 Args: on_message: 回调 ``on_message(text: str) -> None``。 """ if self._sock is None: raise RuntimeError("必须先调用 connect()") self._running = True socklist = [self._sock] while self._running: try: readable, _, exception = select.select( socklist, [], socklist, 0.1 ) if readable: try: text = recv(self._sock) on_message(text) except ConnectionError: break if exception: break except (ConnectionResetError, OSError): break # ---- 关闭 ------------------------------------------------------------ def stop(self) -> None: """设置停止标志并断开 TCP 连接。 可在另一个线程中调用，使 ``listen()`` 退出阻塞。 """ self._running = False if self._sock is not None: try: self._sock.close() except OSError: pass self._sock = None # [SECURITY AUDIT P1-2 ✅ 已修复] 线程池无队列限制 → 慢 handler 耗尽内存 # 原问题：标准 ThreadPoolExecutor 使用无界队列，慢 handler 时任务无限堆积直至 OOM。 # 修复：本类在 submit 时检查排队深度，超限即抛 queue.Full；MsgServer 新增 # max_queue_size 参数（默认 1000），_handle_client 捕获 queue.Full 后拒绝并断开连接。 # 详见 docs/security/pymsg_SECURITY_AUDIT_20260801.md 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 服务器。 收到消息时调用回调函数 —— 转发、过滤等逻辑全部在回调中实现。 Usage:: def on_msg(sender_sock, addr, text): # 转发给所有人（排除发送者） server.broadcast(f"[{addr}] {text}", except_sock=sender_sock) server = MsgServer("0.0.0.0", 8003) server.start(on_msg) # 阻塞当前线程 # 在另一线程中调用 server.stop() 可优雅退出 并发模式:: server = MsgServer("0.0.0.0", 8003, workers=4) ``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, # [SECURITY AUDIT P1-1 ✅ 已修复] 无连接数限制 → 连接耗尽攻击 # 原问题：服务端无并发连接上限，攻击者可耗尽文件描述符/线程资源。 # 修复：新增 max_connections（默认 1000），_accept_client 在 # len(_clients) >= max_connections 时直接拒绝并关闭新连接。 # 详见 docs/security/pymsg_SECURITY_AUDIT_20260801.md max_connections: int = 1000, # [SECURITY AUDIT P1-2 ✅ 已修复] 见上方 _BoundedThreadPoolExecutor 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) self._send_locks: "dict" = {} # socket -> Lock（防并发发送帧交错） self._pool: "ThreadPoolExecutor | None" = None self._lock = threading.Lock() # 保护 _busy / _failed self._busy: "set" = set() # 回调执行中的 socket（暂停 select 监听） self._failed: "set" = set() # 回调失败的 socket（主循环统一移除） # ---- 启动（阻塞）---------------------------------------------------- def start(self, on_message: "callable") -> None: """启动服务器，阻塞当前线程。 Args: on_message: 回调 ``on_message(sock, addr, text)`` - sock: 发送消息的客户端套接字 - addr: (ip, port) 元组 - text: 解码后的字符串消息 ``workers > 0`` 时回调提交到线程池并发执行；同一连接的消息保持串行 （回调完成前该连接暂停 select 监听，恢复粒度为一个 select 周期， 最坏约 100ms——对短连接 RPC 模式无影响）。 """ self._running = True servsock = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM) servsock.setsockopt(_socket.SOL_SOCKET, _socket.SO_REUSEADDR, 1) servsock.bind((self.addr, self.port)) servsock.listen(10) self._server_sock = servsock if self._workers > 0: self._pool = _BoundedThreadPoolExecutor( max_workers=self._workers, max_queue_size=self._max_queue_size, thread_name_prefix=f"pymsg-{self.port}", ) socketlist = [servsock] try: while self._running: self._reap_failed(socketlist) with self._lock: watched = [s for s in socketlist if s not in self._busy] try: readable, _, exception = select.select( watched, [], watched, 0.1 ) except (ValueError, OSError): break for s in exception: self._remove_client(s, socketlist) for s in readable: if s is servsock: self._accept_client(servsock, socketlist) else: self._handle_client(s, socketlist, on_message) finally: if self._pool is not None: self._pool.shutdown(wait=True) self._pool = None self._shutdown(servsock, socketlist) # ---- 广播 ----------------------------------------------------------- def broadcast( self, msg: str, except_sock: "_socket.socket | None" = None ) -> None: """向所有已连接客户端广播消息（pymsg 协议）。 Args: msg: 要发送的字符串。 except_sock: 排除的套接字（通常为发送者），不向其发送。 """ disconnected = [] for client in list(self._clients): if client is except_sock: continue try: self._locked_send(client, msg) except (BrokenPipeError, ConnectionResetError, OSError): disconnected.append(client) for client in disconnected: self._clients.pop(client, None) self._send_locks.pop(client, None) try: client.close() except OSError: pass # ---- 应答 ----------------------------------------------------------- def reply(self, sock: "_socket.socket", msg: str) -> None: """向指定客户端套接字发送一条应答消息（pymsg 协议）。 用于请求-应答模式：服务器收到消息后，通过同一 socket 回复处理结果。 经每连接发送锁串行化，``workers > 0`` 时可在回调线程安全调用。 """ self._locked_send(sock, msg) def _locked_send(self, sock: "_socket.socket", msg: str) -> None: """带每连接锁的发送：防止多线程向同一 socket 并发写入导致帧交错。""" lock = self._send_locks.get(sock) if lock is None: send(sock, msg) return with lock: send(sock, msg) # ---- 停止 ----------------------------------------------------------- def stop(self) -> None: """设置停止标志并关闭监听套接字（可从另一线程调用）。""" self._running = False if self._server_sock is not None: try: self._server_sock.close() except OSError: pass # ---- 内部方法 ------------------------------------------------------- def _accept_client(self, servsock, socketlist): try: client_sock, client_addr = servsock.accept() if client_sock.fileno() == -1: return # [SECURITY AUDIT P1-1 ✅ 已修复] 活跃连接数上限控制：超过 # max_connections 直接拒绝，防止连接耗尽攻击（listen backlog 只限制 # SYN 队列，不限制已建立连接数）。本段即 P1-1 的落地实现。 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 def _handle_client(self, s, socketlist, on_message): try: text = recv(s, timeout=self._recv_timeout) except (ConnectionError, ConnectionResetError, OSError, TimeoutError, UnicodeDecodeError): # [SECURITY AUDIT P0-2 ✅ 已修复（第二层兜底）] UnicodeDecodeError： # recv() 已将其包装为 ConnectionError，此处防御性捕获防止任何遗漏。 # TimeoutError：客户端在超时时间内未发送完整消息（Slowloris 防护，P0-1），关闭连接 self._remove_client(s, socketlist) return addr = self._clients.get(s, ("unknown", 0)) if self._pool is None: # 同步模式（默认）：回调在主循环线程内执行，行为与旧版一致 try: on_message(s, addr, text) except (ConnectionError, ConnectionResetError, OSError): self._remove_client(s, socketlist) return # 线程池模式：先标记 busy（暂停 select 监听该连接，保证同连接串行）， # 再把回调提交到线程池执行 with self._lock: self._busy.add(s) try: self._pool.submit(self._run_task, s, addr, text, on_message) except _queue.Full: # [SECURITY AUDIT P1-2 ✅ 已修复] 有界队列已满：慢 handler 时任务不应 # 无限堆积，拒绝并断开该连接，防止内存无限增长。 with self._lock: self._busy.discard(s) self._remove_client(s, socketlist) def _run_task(self, s, addr, text, on_message): """工作线程：执行用户回调。 出错时只打标记（``_failed``），实际的连接关闭与移除由主循环统一 执行，避免 worker 线程与 accept/select 并发操作 fd。 """ try: on_message(s, addr, text) except (ConnectionError, ConnectionResetError, OSError): with self._lock: self._failed.add(s) except Exception: # 同步模式下此类异常会直接使 start() 崩溃；线程池模式无法传播， # 打印 traceback 并移除该连接，保证服务整体可用 traceback.print_exc() with self._lock: self._failed.add(s) finally: with self._lock: self._busy.discard(s) def _reap_failed(self, socketlist): """主循环每轮统一移除回调失败的连接。""" with self._lock: failed, self._failed = self._failed, set() for s in failed: self._remove_client(s, socketlist) def _remove_client(self, s, socketlist): self._clients.pop(s, None) self._send_locks.pop(s, None) if s in socketlist: socketlist.remove(s) try: s.close() except OSError: pass def _shutdown(self, servsock, socketlist): for s in socketlist: try: s.close() except OSError: pass self._clients.clear() self._send_locks.clear() # ── JSON RPC 约定层（依赖上面的 MsgClient / MsgServer，故放在末尾导入）── from .rpc import RpcClient, RpcError, RpcServer # noqa: E402 from .async_queue import AsyncRpcQueue # noqa: E402 from .logc import LogClient, set_log_level # noqa: E402 __all__ += ["RpcClient", "RpcServer", "RpcError", "AsyncRpcQueue", "LogClient", "set_log_level"]