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 仅用于静态文案。
238 lines
8.7 KiB
Python
238 lines
8.7 KiB
Python
"""白板接口：REST（访问/管理）+ WebSocket（实时同步）。
|
||
|
||
路由：
|
||
GET /api/wb/{board_id} 公开：访问记事本元数据，不存在则新建（前端 init 用）
|
||
WS /ws/wb/{board_id} 公开：实时协作 + 心跳
|
||
GET /api/admin/wb Basic Auth：管理页列表
|
||
DELETE /api/admin/wb/{board_id} Basic Auth：删除记事本
|
||
|
||
HTML 页面 /wb/{id} 与管理页 /wb-admin 由 main.py 直接返回静态文件，
|
||
不在此 controller 注册，避免与 REST 同路径冲突。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import uuid
|
||
|
||
from fastapi import APIRouter, Depends, HTTPException, Query, WebSocket, WebSocketDisconnect
|
||
from sqlalchemy.orm import Session
|
||
|
||
from ..database import get_db
|
||
from ..dao.whiteboard_dao import WhiteboardDAO
|
||
from ..schemas.whiteboard import WhiteboardListResponse, WhiteboardOut
|
||
from ..security import require_docs_auth
|
||
from ..services.whiteboard_hub import Connection, get_hub
|
||
from ..services.whiteboard_service import WhiteboardService
|
||
|
||
logger = logging.getLogger("zikai.whiteboard")
|
||
|
||
router = APIRouter(tags=["whiteboard"])
|
||
|
||
|
||
def _service(db: Session = Depends(get_db)) -> WhiteboardService:
|
||
"""REST 路径的 service：注入 hub 以便删除时踢出连接。"""
|
||
return WhiteboardService(WhiteboardDAO(db), hub=get_hub())
|
||
|
||
|
||
# ---------------- 公开 REST ----------------
|
||
|
||
@router.get(
|
||
"/api/wb/{board_id}",
|
||
response_model=WhiteboardOut,
|
||
summary="访问记事本元数据（不存在则新建）",
|
||
description="前端打开 /wb/{id} 页面后调本接口拉取初始文本；不存在时自动创建空板。",
|
||
)
|
||
def get_whiteboard(board_id: str, service: WhiteboardService = Depends(_service)) -> WhiteboardOut:
|
||
return service.get_or_create(board_id)
|
||
|
||
|
||
# ---------------- 管理 REST（Basic Auth） ----------------
|
||
|
||
@router.get(
|
||
"/api/admin/wb",
|
||
response_model=WhiteboardListResponse,
|
||
summary="列出所有记事本（需鉴权）",
|
||
description="供记事本管理页使用：board_id / 创建时间 / 编辑次数 / 上次修改时间。",
|
||
)
|
||
def list_whiteboards(
|
||
limit: int = Query(100, ge=1, le=500),
|
||
offset: int = Query(0, ge=0),
|
||
service: WhiteboardService = Depends(_service),
|
||
_: str = Depends(require_docs_auth),
|
||
) -> WhiteboardListResponse:
|
||
total, items = service.list_all(limit=limit, offset=offset)
|
||
return WhiteboardListResponse(total=total, items=items)
|
||
|
||
|
||
@router.delete(
|
||
"/api/admin/wb/{board_id}",
|
||
summary="删除记事本（需鉴权）",
|
||
description="删 DB 行，并关闭该 board 的所有在线 WebSocket 连接。",
|
||
)
|
||
def delete_whiteboard(
|
||
board_id: str,
|
||
service: WhiteboardService = Depends(_service),
|
||
_: str = Depends(require_docs_auth),
|
||
) -> dict:
|
||
ok = service.delete(board_id)
|
||
if not ok:
|
||
raise HTTPException(404, "白板不存在")
|
||
return {"deleted": True}
|
||
|
||
|
||
# ---------------- WebSocket（公开，实时同步 + 心跳） ----------------
|
||
|
||
@router.websocket("/ws/wb/{board_id}")
|
||
async def whiteboard_ws(websocket: WebSocket, board_id: str) -> None:
|
||
"""白板实时协作端点（文本记事本）。
|
||
|
||
协议（JSON 文本帧）：
|
||
client -> server:
|
||
{"type":"hello","client_id":"..."} 首帧（可选，未带则服务端生成）
|
||
{"type":"ping"} 心跳，服务端回 pong 并刷新计时
|
||
{"type":"edit","content":"..."} debounce 后发完整文本，持久化并广播给他人
|
||
{"type":"clear"} 清空，持久化并广播给所有人
|
||
server -> client:
|
||
{"type":"init","content":"...","version":n,"edit_count":m}
|
||
{"type":"pong"}
|
||
{"type":"update","content":"...","version":n,"client_id":"..."} 文本变更
|
||
{"type":"cleared","client_id":"..."}
|
||
{"type":"error","msg":"..."}
|
||
"""
|
||
# 路径层只做最基本校验，详细校验交给 service（service 会查表）
|
||
hub = get_hub()
|
||
# 先 accept，便于对非法 board_id 也回一条 error 再关闭
|
||
await websocket.accept()
|
||
|
||
# 读取首帧 hello（或任意帧）拿 client_id
|
||
try:
|
||
first = await websocket.receive_text()
|
||
except WebSocketDisconnect:
|
||
return
|
||
|
||
client_id = _extract_client_id(first) or uuid.uuid4().hex[:12]
|
||
|
||
# 校验 board_id 并加载白板（不存在则新建）
|
||
try:
|
||
board = _with_db(lambda db: WhiteboardService(WhiteboardDAO(db), hub=hub).get_or_create(board_id))
|
||
except HTTPException as exc:
|
||
await _safe_send(websocket, {"type": "error", "msg": exc.detail})
|
||
await _safe_close(websocket)
|
||
return
|
||
|
||
# 注册连接并下发 init
|
||
conn = Connection(websocket=websocket, board_id=board_id, client_id=client_id)
|
||
ok = await hub.register(conn)
|
||
if not ok:
|
||
await _safe_send(websocket, {"type": "error", "msg": "该记事本在线人数已满"})
|
||
await _safe_close(websocket)
|
||
return
|
||
await _safe_send(websocket, {
|
||
"type": "init",
|
||
"content": board.content,
|
||
"version": board.version,
|
||
"edit_count": board.edit_count,
|
||
})
|
||
|
||
# 主循环：收消息 -> 处理 -> 广播
|
||
# 单帧大小上限：与 content 限制对齐（256KB 文本 + JSON 开销，留余量到 512KB）
|
||
MAX_FRAME = 512 * 1024
|
||
try:
|
||
while True:
|
||
raw = await websocket.receive_text()
|
||
if len(raw) > MAX_FRAME:
|
||
await _safe_send(websocket, {"type": "error", "msg": "消息过大"})
|
||
continue
|
||
msg = _parse(raw)
|
||
if msg is None:
|
||
continue
|
||
mtype = msg.get("type")
|
||
if mtype == "ping":
|
||
conn.touch()
|
||
await _safe_send(websocket, {"type": "pong"})
|
||
continue
|
||
# 任何有效业务帧都视为活性证据
|
||
conn.touch()
|
||
if mtype == "edit":
|
||
content = msg.get("content")
|
||
if not isinstance(content, str):
|
||
await _safe_send(websocket, {"type": "error", "msg": "content 必须是字符串"})
|
||
continue
|
||
try:
|
||
out = _with_db(
|
||
lambda db: WhiteboardService(WhiteboardDAO(db), hub=hub).update_content(board_id, content)
|
||
)
|
||
except HTTPException as exc:
|
||
await _safe_send(websocket, {"type": "error", "msg": exc.detail})
|
||
continue
|
||
# 广播给他人（发送者本地已更新，不回推）
|
||
await hub.broadcast(
|
||
board_id,
|
||
{"type": "update", "content": out.content, "version": out.version, "client_id": client_id},
|
||
exclude=conn,
|
||
)
|
||
elif mtype == "clear":
|
||
try:
|
||
_with_db(lambda db: WhiteboardService(WhiteboardDAO(db), hub=hub).clear(board_id))
|
||
except HTTPException as exc:
|
||
await _safe_send(websocket, {"type": "error", "msg": exc.detail})
|
||
continue
|
||
# clear 广播给所有人（含发送者，用于确认）
|
||
await hub.broadcast(
|
||
board_id, {"type": "cleared", "client_id": client_id}
|
||
)
|
||
else:
|
||
await _safe_send(websocket, {"type": "error", "msg": f"未知消息类型 {mtype}"})
|
||
except WebSocketDisconnect:
|
||
pass
|
||
except Exception as exc: # pragma: no cover
|
||
logger.warning("白板 WS 异常 board=%s client=%s: %s", board_id, client_id, exc)
|
||
finally:
|
||
await hub.disconnect(conn)
|
||
|
||
|
||
# ---------------- helpers ----------------
|
||
|
||
def _with_db(fn):
|
||
"""在独立 Session 中执行 fn 并返回结果；用完即关。供 WS 路径每帧独立事务使用。"""
|
||
from ..database import get_session_local
|
||
db = get_session_local()()
|
||
try:
|
||
return fn(db)
|
||
finally:
|
||
db.close()
|
||
|
||
def _extract_client_id(raw: str) -> str | None:
|
||
try:
|
||
data = json.loads(raw)
|
||
except (json.JSONDecodeError, TypeError):
|
||
return None
|
||
cid = data.get("client_id") if isinstance(data, dict) else None
|
||
if isinstance(cid, str) and cid:
|
||
return cid
|
||
return None
|
||
|
||
|
||
def _parse(raw: str) -> dict | None:
|
||
try:
|
||
data = json.loads(raw)
|
||
except (json.JSONDecodeError, TypeError):
|
||
return None
|
||
return data if isinstance(data, dict) else None
|
||
|
||
|
||
async def _safe_send(ws: WebSocket, msg: dict) -> None:
|
||
try:
|
||
await ws.send_json(msg)
|
||
except Exception: # pragma: no cover
|
||
pass
|
||
|
||
|
||
async def _safe_close(ws: WebSocket) -> None:
|
||
try:
|
||
await ws.close()
|
||
except Exception: # pragma: no cover
|
||
pass
|