da7d2fb4b4
安全审计发现： - WS receive_text 无应用层大小限制，恶意客户端可发超大帧（uvicorn 默认 16MB 才拦截）占内存。加 512KB 单帧检查（与 content 256KB 限制对齐留余量）。 - 单 board 无连接数上限，恶意脚本可开海量连接耗尽资源。config 加 max_connections_per_board（默认 50），hub.register 超限返回 False， controller 回 error 帧并关闭连接。 审计结论（无需修复）： - SQL：全部 SQLAlchemy ORM 参数化，无注入。 - 路径穿越：storage_path 服务端生成（uuid+basename(ext)），用户不可控分隔符。 - 鉴权：管理类操作均有 require_docs_auth，公开写接口符合设计。 - 下载 filename CRLF：Starlette FileResponse 用 quote() 编码，无响应拆分。 - XSS：前端 el() 用 createTextNode，textarea 纯文本，innerHTML 仅用于静态文案。
183 lines
7.0 KiB
Python
183 lines
7.0 KiB
Python
"""白板 WebSocket 连接管理器（实时同步 + 心跳 + 失活清理）。
|
||
|
||
设计要点：
|
||
- 进程内单例 ``WhiteboardHub``，维护 ``{board_id: set[Connection]}``。
|
||
- 每个 Connection 封装 websocket / board_id / client_id / last_heartbeat。
|
||
- 心跳：客户端每 ``heartbeat_interval_seconds``（默认 3s）发一次 ping，服务端回 pong 并
|
||
刷新 last_heartbeat。reaper 每秒扫描，超过 ``interval * threshold``（默认 15s）未心跳
|
||
的连接判为失活，关闭并从 hub 移除。
|
||
- 内存安全：disconnect 幂等；空 set 从 dict 删除；broadcast 对单连接异常立即 disconnect；
|
||
close_board 关闭并清理整个 board 的连接集合。
|
||
- 并发：用一个 asyncio.Lock 保护 ``_boards`` 的结构变更（add/remove board key），
|
||
集合内的连接增删用 set 原子操作（Python 单线程事件循环下安全）。
|
||
|
||
注意：hub 是进程内存，多 worker 下不互通。生产部署需单 worker 或后续接 Redis pub/sub。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
|
||
from fastapi import WebSocket
|
||
|
||
from ..config import get_settings
|
||
|
||
logger = logging.getLogger("zikai.whiteboard")
|
||
|
||
|
||
@dataclass(eq=False)
|
||
class Connection:
|
||
"""一个白板在线连接。eq=False 使其按对象 identity 哈希/比较，可放入 set。"""
|
||
|
||
websocket: WebSocket
|
||
board_id: str
|
||
client_id: str
|
||
last_heartbeat: float = field(default_factory=time.monotonic)
|
||
|
||
async def send_json(self, msg: dict) -> bool:
|
||
"""发送一条消息；失败返回 False（调用方据此 disconnect）。"""
|
||
try:
|
||
await self.websocket.send_json(msg)
|
||
return True
|
||
except Exception as exc: # WebSocketDisconnect / 已关闭 / 编码失败
|
||
logger.debug("发送失败 board=%s client=%s: %s", self.board_id, self.client_id, exc)
|
||
return False
|
||
|
||
def touch(self) -> None:
|
||
self.last_heartbeat = time.monotonic()
|
||
|
||
|
||
class WhiteboardHub:
|
||
"""白板连接管理器单例。"""
|
||
|
||
def __init__(self) -> None:
|
||
cfg = get_settings().whiteboard
|
||
self.heartbeat_interval = cfg.heartbeat_interval_seconds
|
||
self.heartbeat_miss_threshold = cfg.heartbeat_miss_threshold
|
||
self.timeout_seconds = self.heartbeat_interval * self.heartbeat_miss_threshold
|
||
self.max_connections_per_board = cfg.max_connections_per_board
|
||
# {board_id: set[Connection]}
|
||
self._boards: dict[str, set[Connection]] = {}
|
||
self._lock = asyncio.Lock()
|
||
|
||
# ---------------- 连接生命周期 ----------------
|
||
|
||
async def register(self, conn: Connection) -> bool:
|
||
"""把已 accept 的连接加入 board 集合。
|
||
|
||
返回 False 表示该 board 连接数已达上限（调用方应关闭连接）。
|
||
"""
|
||
async with self._lock:
|
||
conns = self._boards.setdefault(conn.board_id, set())
|
||
if len(conns) >= self.max_connections_per_board:
|
||
return False
|
||
conns.add(conn)
|
||
logger.info("连接接入 board=%s client=%s（当前 %d 人）",
|
||
conn.board_id, conn.client_id, self.connection_count(conn.board_id))
|
||
return True
|
||
|
||
async def disconnect(self, conn: Connection) -> None:
|
||
"""幂等移除连接；空 set 从 dict 删除以防内存泄漏。"""
|
||
async with self._lock:
|
||
conns = self._boards.get(conn.board_id)
|
||
if conns is None:
|
||
return
|
||
conns.discard(conn)
|
||
if not conns:
|
||
self._boards.pop(conn.board_id, None)
|
||
# 尽力关闭 websocket（可能已关闭）
|
||
try:
|
||
await conn.websocket.close()
|
||
except Exception: # pragma: no cover
|
||
pass
|
||
logger.info("连接移除 board=%s client=%s（剩余 %d 人）",
|
||
conn.board_id, conn.client_id, self.connection_count(conn.board_id))
|
||
|
||
def connection_count(self, board_id: str) -> int:
|
||
"""调试/监控用：某 board 当前在线人数。"""
|
||
return len(self._boards.get(board_id, ()))
|
||
|
||
# ---------------- 广播 ----------------
|
||
|
||
async def broadcast(self, board_id: str, msg: dict, exclude: Connection | None = None) -> None:
|
||
"""把 msg 发给 board 内所有在线连接（可排除发送者）。单连接失败不影响其他。"""
|
||
async with self._lock:
|
||
conns = list(self._boards.get(board_id, ()))
|
||
dead: list[Connection] = []
|
||
for conn in conns:
|
||
if exclude is not None and conn is exclude:
|
||
continue
|
||
ok = await conn.send_json(msg)
|
||
if not ok:
|
||
dead.append(conn)
|
||
# 发送失败的连接统一清理
|
||
for conn in dead:
|
||
await self.disconnect(conn)
|
||
|
||
# ---------------- 心跳 reaper ----------------
|
||
|
||
async def reap_loop(self, stop: asyncio.Event) -> None:
|
||
"""后台循环：扫描失活连接。每秒一次，粒度细于 timeout。"""
|
||
logger.info("白板心跳 reaper 已启动（间隔 1s，超时 %ds）", self.timeout_seconds)
|
||
while not stop.is_set():
|
||
try:
|
||
await self._reap_once()
|
||
except Exception as exc: # pragma: no cover
|
||
logger.warning("reaper 循环异常：%s", exc)
|
||
try:
|
||
await asyncio.wait_for(stop.wait(), timeout=1.0)
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
|
||
async def _reap_once(self) -> None:
|
||
now = time.monotonic()
|
||
async with self._lock:
|
||
# 快照待检查连接，避免持锁时 await
|
||
stale: list[Connection] = []
|
||
for board_id, conns in self._boards.items():
|
||
for conn in conns:
|
||
if now - conn.last_heartbeat > self.timeout_seconds:
|
||
stale.append(conn)
|
||
for conn in stale:
|
||
logger.warning("心跳失活，移除 board=%s client=%s（静默 %ds）",
|
||
conn.board_id, conn.client_id,
|
||
int(now - conn.last_heartbeat))
|
||
await self.disconnect(conn)
|
||
|
||
async def close_board(self, board_id: str) -> None:
|
||
"""关闭并清理某 board 的所有连接（删除白板时调用）。"""
|
||
async with self._lock:
|
||
conns = self._boards.pop(board_id, None)
|
||
if not conns:
|
||
return
|
||
await asyncio.gather(
|
||
*(c.send_json({"type": "error", "msg": "白板已被删除"}) for c in conns),
|
||
return_exceptions=True,
|
||
)
|
||
for conn in conns:
|
||
try:
|
||
await conn.websocket.close()
|
||
except Exception: # pragma: no cover
|
||
pass
|
||
logger.info("关闭白板 board=%s，踢出 %d 个连接", board_id, len(conns))
|
||
|
||
|
||
# 进程内单例（由 main.py lifespan / controller 共享）
|
||
_hub: WhiteboardHub | None = None
|
||
|
||
|
||
def get_hub() -> WhiteboardHub:
|
||
global _hub
|
||
if _hub is None:
|
||
_hub = WhiteboardHub()
|
||
return _hub
|
||
|
||
|
||
def reset_hub() -> None:
|
||
"""测试用：重置单例。"""
|
||
global _hub
|
||
_hub = None
|