chore: initial baseline with P0-safety .gitignore
This commit is contained in:
Simon
@@ -0,0 +1,278 @@
|
||||
# =============================================================================
|
||||
# 企微IT智能服务台 — WebSocket 端点
|
||||
# =============================================================================
|
||||
# 说明：提供 WebSocket 端点，供坐席前端和H5用户端建立长连接，实现实时推送。
|
||||
# 核心功能：
|
||||
# 1. 接受坐席的 WebSocket 连接请求（含 token 认证）— /ws/{agent_id}
|
||||
# 2. 接受H5员工的 WebSocket 连接请求（含 token 认证）— /ws/h5/{employee_id}
|
||||
# 3. 维持连接，监听客户端消息（主要是心跳 ping）
|
||||
# 4. 连接断开时自动清理注册信息
|
||||
# 安全（WS-01）：
|
||||
# 握手时从 query param 取 token → 查 Redis 验证 → 不通过则 close(code=4001)
|
||||
# 防止未授权用户冒充坐席/员工建立 WS 连接
|
||||
#
|
||||
# 端点路径：
|
||||
# - 坐席端：/ws/{agent_id}?token=xxx
|
||||
# - H5员工端：/ws/h5/{employee_id}?token=xxx
|
||||
# 为什么不挂 /api 前缀：WebSocket 不是 REST API，不走 Vite 的 /api 代理配置
|
||||
# =============================================================================
|
||||
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
|
||||
|
||||
from app.services.ws_manager import manager as ws_manager
|
||||
from app.services.cache_service import cache_service
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# WebSocket 路由器（不挂 /api 前缀，直接注册在应用根路径）
|
||||
router = APIRouter()
|
||||
|
||||
# 认证失败时的 WebSocket 关闭码
|
||||
# 4001 = 自定义码，表示"未授权"（4000+ 为应用自定义范围）
|
||||
WS_CLOSE_UNAUTHORIZED = 4001
|
||||
|
||||
|
||||
@router.websocket("/ws/{agent_id}")
|
||||
async def websocket_endpoint(
|
||||
websocket: WebSocket,
|
||||
agent_id: str,
|
||||
token: str = Query(default="", description="登录 token，用于 WebSocket 认证"),
|
||||
) -> None:
|
||||
"""坐席 WebSocket 端点主循环（含 WS-01 token 认证）。
|
||||
|
||||
做什么：
|
||||
1. 验证 token 有效性（查 Redis）
|
||||
2. 验证 token 与 agent_id 一致性（防冒充）
|
||||
3. 认证通过后接受连接，注册到 ConnectionManager
|
||||
4. 进入消息接收循环，处理客户端发送的消息
|
||||
5. 连接断开时清理注册信息
|
||||
|
||||
为什么需要 token 认证（WS-01）：
|
||||
- 之前 /ws/{agent_id} 无任何认证，任何人知道 URL 即可冒充任意坐席
|
||||
- 攻击者可监听所有消息、发送伪造消息，是 P0 级安全漏洞
|
||||
- 修复后，必须提供与 agent_id 匹配的有效 token 才能建立连接
|
||||
|
||||
Args:
|
||||
websocket: FastAPI WebSocket 对象（框架自动注入）
|
||||
agent_id: 坐席ID（从 URL 路径参数获取）
|
||||
token: 登录 token（从 URL query parameter 获取）
|
||||
"""
|
||||
# ======================================================================
|
||||
# WS-01: Token 认证
|
||||
# ======================================================================
|
||||
|
||||
# 步骤1: 检查 token 是否为空
|
||||
if not token:
|
||||
# 先 accept 再 close，否则客户端收不到关闭帧
|
||||
await websocket.accept()
|
||||
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Missing token")
|
||||
logger.warning(f"WebSocket 拒绝连接: agent_id={agent_id}, 原因=缺少token")
|
||||
return
|
||||
|
||||
# 步骤2: 从 Redis 查询 token 对应的坐席信息
|
||||
# Redis 中存储格式: agent:token:{token} -> agent_user_id
|
||||
# （与坐席登录 API /api/agents/login 存储格式一致）
|
||||
try:
|
||||
stored_agent_id = await cache_service.get(f"agent:token:{token}")
|
||||
except Exception as e:
|
||||
# Redis 不可用时必须拒绝连接：token 验证依赖 Redis，无法验证身份
|
||||
# 如果降级放行，攻击者可在 Redis 故障时用任意 agent_id 冒充坐席
|
||||
logger.error(f"Redis 查询失败，拒绝 WS 连接: agent_id={agent_id}, error={e}")
|
||||
await websocket.accept()
|
||||
await websocket.close(
|
||||
code=WS_CLOSE_UNAUTHORIZED,
|
||||
reason="Authentication service unavailable"
|
||||
)
|
||||
return
|
||||
|
||||
# 步骤3: 验证 token 与 agent_id 一致性
|
||||
if not stored_agent_id:
|
||||
# token 不存在（已过期或伪造）
|
||||
await websocket.accept()
|
||||
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Invalid or expired token")
|
||||
logger.warning(f"WebSocket 拒绝连接: agent_id={agent_id}, 原因=token无效或已过期")
|
||||
return
|
||||
|
||||
if stored_agent_id != agent_id:
|
||||
# token 对应的坐席与请求的 agent_id 不匹配（冒充）
|
||||
await websocket.accept()
|
||||
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Token-agent mismatch")
|
||||
logger.warning(
|
||||
f"WebSocket 拒绝连接: agent_id={agent_id}, "
|
||||
f"原因=token对应坐席{stored_agent_id}与请求不匹配"
|
||||
)
|
||||
return
|
||||
|
||||
# ======================================================================
|
||||
# 认证通过，建立连接
|
||||
# ======================================================================
|
||||
|
||||
# 注册连接（内部会调用 websocket.accept()）
|
||||
await ws_manager.connect(agent_id, websocket)
|
||||
logger.info(f"坐席 WebSocket 连接已认证: agent_id={agent_id}")
|
||||
|
||||
try:
|
||||
# 消息接收循环
|
||||
# 保持连接打开，监听客户端发来的消息
|
||||
# 即使客户端不发消息，这个循环也必须保持，否则连接会关闭
|
||||
while True:
|
||||
# 等待接收客户端消息（阻塞等待）
|
||||
data = await websocket.receive_json()
|
||||
|
||||
# 处理心跳 ping
|
||||
# 前端每 30 秒发送一次 ping，后端回复 pong
|
||||
# 作用：检测连接是否存活，防止中间代理（如 Nginx）因超时断开连接
|
||||
if data.get("type") == "ping":
|
||||
await websocket.send_json({"type": "pong"})
|
||||
logger.debug(f"WebSocket 心跳: agent_id={agent_id}")
|
||||
|
||||
# 处理输入指示器 typing 事件
|
||||
# 前端在用户输入时发送 typing 事件，后端广播给同一会话的其他参与者
|
||||
elif data.get("type") == "typing":
|
||||
conversation_id = data.get("conversation_id")
|
||||
sender_name = data.get("sender_name", agent_id)
|
||||
if conversation_id:
|
||||
# 广播给所有坐席（包含 sender_type 和 sender_id，
|
||||
# 前端可据此过滤掉自己的 typing 事件）
|
||||
await ws_manager.broadcast({
|
||||
"type": "typing",
|
||||
"data": {
|
||||
"conversation_id": conversation_id,
|
||||
"sender_id": agent_id,
|
||||
"sender_name": sender_name,
|
||||
"sender_type": "agent",
|
||||
}
|
||||
})
|
||||
|
||||
else:
|
||||
# 未来可扩展处理其他类型的客户端消息
|
||||
logger.debug(
|
||||
f"WebSocket 收到未知消息: agent_id={agent_id}, "
|
||||
f"type={data.get('type', 'unknown')}"
|
||||
)
|
||||
|
||||
except WebSocketDisconnect:
|
||||
# 客户端主动断开连接（正常行为）
|
||||
# 清理 ConnectionManager 中的注册信息
|
||||
ws_manager.disconnect(agent_id)
|
||||
logger.info(f"坐席断开 WebSocket 连接: agent_id={agent_id}")
|
||||
|
||||
except Exception as e:
|
||||
# 其他异常（如网络错误、JSON 解析错误等）
|
||||
# 确保注册信息被清理
|
||||
ws_manager.disconnect(agent_id)
|
||||
logger.warning(f"WebSocket 异常断开: agent_id={agent_id}, error={e}")
|
||||
|
||||
|
||||
# ==========================================================================
|
||||
# H5员工 WebSocket 端点
|
||||
# ==========================================================================
|
||||
|
||||
@router.websocket("/ws/h5/{employee_id}")
|
||||
async def h5_websocket_endpoint(
|
||||
websocket: WebSocket,
|
||||
employee_id: str,
|
||||
token: str = Query(default="", description="H5员工登录 token，用于 WebSocket 认证"),
|
||||
) -> None:
|
||||
"""H5员工 WebSocket 端点主循环（含 token 认证）。
|
||||
|
||||
做什么：
|
||||
1. 验证 employee token 有效性（查 Redis）
|
||||
2. 验证 token 与 employee_id 一致性（防冒充）
|
||||
3. 认证通过后接受连接，注册到 ConnectionManager 的员工连接表
|
||||
4. 进入消息接收循环，处理心跳 ping
|
||||
5. 连接断开时清理注册信息
|
||||
|
||||
为什么需要 H5 WS 连接：
|
||||
- H5员工需要实时接收参与者变更事件（新参与者加入、有人退出等）
|
||||
- 当前仅通过 3 秒轮询获取更新，实时性不足
|
||||
- WS 推送 + 轮询降级，双通道保证消息可达
|
||||
|
||||
认证机制（与坐席端一致）：
|
||||
- Redis 中存储格式: employee:token:{token} -> employee_id
|
||||
- （与H5登录 API /api/h5/mock-login 存储格式一致）
|
||||
- token 缺失、无效、过期、与 employee_id 不匹配均拒绝连接
|
||||
|
||||
Args:
|
||||
websocket: FastAPI WebSocket 对象（框架自动注入）
|
||||
employee_id: 员工企微 UserID（从 URL 路径参数获取）
|
||||
token: H5员工登录 token（从 URL query parameter 获取）
|
||||
"""
|
||||
# ======================================================================
|
||||
# Token 认证
|
||||
# ======================================================================
|
||||
|
||||
# 步骤1: 检查 token 是否为空
|
||||
if not token:
|
||||
await websocket.accept()
|
||||
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Missing token")
|
||||
logger.warning(f"H5 WebSocket 拒绝连接: employee_id={employee_id}, 原因=缺少token")
|
||||
return
|
||||
|
||||
# 步骤2: 从 Redis 查询 token 对应的员工信息
|
||||
# Redis 中存储格式: employee:token:{token} -> employee_id
|
||||
# （与H5登录 API /api/h5/mock-login 存储格式一致）
|
||||
try:
|
||||
stored_employee_id = await cache_service.get(f"employee:token:{token}")
|
||||
except Exception as e:
|
||||
# Redis 不可用时必须拒绝连接（与坐席端一致的安全策略）
|
||||
logger.error(f"Redis 查询失败，拒绝 H5 WS 连接: employee_id={employee_id}, error={e}")
|
||||
await websocket.accept()
|
||||
await websocket.close(
|
||||
code=WS_CLOSE_UNAUTHORIZED,
|
||||
reason="Authentication service unavailable"
|
||||
)
|
||||
return
|
||||
|
||||
# 步骤3: 验证 token 与 employee_id 一致性
|
||||
if not stored_employee_id:
|
||||
await websocket.accept()
|
||||
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Invalid or expired token")
|
||||
logger.warning(f"H5 WebSocket 拒绝连接: employee_id={employee_id}, 原因=token无效或已过期")
|
||||
return
|
||||
|
||||
if stored_employee_id != employee_id:
|
||||
await websocket.accept()
|
||||
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Token-employee mismatch")
|
||||
logger.warning(
|
||||
f"H5 WebSocket 拒绝连接: employee_id={employee_id}, "
|
||||
f"原因=token对应员工{stored_employee_id}与请求不匹配"
|
||||
)
|
||||
return
|
||||
|
||||
# ======================================================================
|
||||
# 认证通过，建立连接
|
||||
# ======================================================================
|
||||
|
||||
# 注册员工连接（内部会调用 websocket.accept()）
|
||||
await ws_manager.connect_employee(employee_id, websocket)
|
||||
logger.info(f"H5员工 WebSocket 连接已认证: employee_id={employee_id}")
|
||||
|
||||
try:
|
||||
# 消息接收循环
|
||||
# H5员工端目前只发送心跳 ping，不需要发送 typing 等事件
|
||||
while True:
|
||||
data = await websocket.receive_json()
|
||||
|
||||
# 处理心跳 ping
|
||||
if data.get("type") == "ping":
|
||||
await websocket.send_json({"type": "pong"})
|
||||
logger.debug(f"H5 WebSocket 心跳: employee_id={employee_id}")
|
||||
|
||||
else:
|
||||
logger.debug(
|
||||
f"H5 WebSocket 收到未知消息: employee_id={employee_id}, "
|
||||
f"type={data.get('type', 'unknown')}"
|
||||
)
|
||||
|
||||
except WebSocketDisconnect:
|
||||
# 客户端主动断开连接
|
||||
ws_manager.disconnect_employee(employee_id)
|
||||
logger.info(f"H5员工断开 WebSocket 连接: employee_id={employee_id}")
|
||||
|
||||
except Exception as e:
|
||||
# 其他异常
|
||||
ws_manager.disconnect_employee(employee_id)
|
||||
logger.warning(f"H5 WebSocket 异常断开: employee_id={employee_id}, error={e}")
|
||||
Reference in New Issue
Block a user