From 839ba91354bbf3b171a770cc1a2262ce2632941e Mon Sep 17 00:00:00 2001 From: weijuesen Date: Fri, 14 Aug 2026 00:45:22 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E4=BF=AE=E5=A4=8D=20WebSocket=20?= =?UTF-8?q?=E8=B6=8A=E6=9D=83=E4=B8=8E=20AI=20=E6=B5=81=20SSRF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 1 + ai-service/README.md | 5 + ai-service/app/config.py | 8 ++ ai-service/app/main.py | 61 ++++++----- ai-service/app/stream_tasks.py | 101 ++++++++++++++++++ ai-service/tests/test_api.py | 42 +++++++- app/src/utils/ws.ts | 24 ++++- miniapp/src/api/config.ts | 8 +- miniapp/src/utils/ws.ts | 24 ++++- server-go/cmd/server/main.go | 4 +- server-go/internal/config/config.go | 59 +++++------ server-go/internal/ws/authorizer.go | 59 +++++++++++ server-go/internal/ws/gateway.go | 144 +++++++++++++++----------- server-go/internal/ws/gateway_test.go | 120 +++++++++++++++++++++ server-go/internal/ws/ticket.go | 57 ++++++++++ 后续工作计划.md | 2 +- 开发交接记录.md | 31 ++++++ 17 files changed, 614 insertions(+), 136 deletions(-) create mode 100644 ai-service/app/stream_tasks.py create mode 100644 server-go/internal/ws/authorizer.go create mode 100644 server-go/internal/ws/gateway_test.go create mode 100644 server-go/internal/ws/ticket.go diff --git a/README.md b/README.md index e2fb719..1d3ad3b 100644 --- a/README.md +++ b/README.md @@ -501,6 +501,7 @@ C:\msys64\usr\bin\sshpass.exe -p "pan" C:\msys64\usr\bin\ssh.exe -o StrictHostKe | `PG` | `postgresql://postgres:pan@localhost:5432/silk` | PostgreSQL 连接串 | | `APP_ENV` | `development` | 运行环境；生产环境应设为 `production` | | `ALLOW_DEV_AUTOMIGRATE` | `false` | 仅开发环境可显式开启 AutoMigrate，生产忽略此开关 | +| `WS_ALLOWED_ORIGINS` | `http://localhost:5174,http://localhost:3000,...` | WebSocket 允许的 Origin，逗号分隔 | | `REDIS` | `redis://:pan@localhost:6379` | Valkey/Redis 连接串 | | `JWT_SECRET` | `silk-secret-please-change-me` | JWT 签名密钥 | | `JWT_EXPIRES_IN` | `2h` | JWT 有效期 | diff --git a/ai-service/README.md b/ai-service/README.md index 2fa0bbc..89489e1 100644 --- a/ai-service/README.md +++ b/ai-service/README.md @@ -6,6 +6,7 @@ FastAPI + ONNX Runtime 的蚕病检测推理服务。YOLO 训练挂起期间以 - `GET /health` → `{"status":"ok","model":"mock|onnx"}` - `POST /detect`（multipart 字段 `file`）→ `{"model":"mock","detections":[{"bbox":{x,y,w,h},"class":"healthy|sick","confidence":0.95}]}` +- `POST /internal/stream-tasks`（需 `X-Internal-Key`）→ 创建受限拉流任务，不再接受客户端任意 URL ## 本地运行 @@ -34,6 +35,10 @@ powershell -ExecutionPolicy Bypass -File scripts/verify.ps1 | `MODEL_LABELS` | `healthy,sick` | 类别列表（逗号分隔） | | `MOCK_CLASS` | `healthy` | mock 返回类别 | | `MOCK_CONFIDENCE` | `0.95` | mock 返回置信度 | +| `INTERNAL_API_KEY` | `silk-internal-2026` | 内部接口认证密钥 | +| `ALLOWED_STREAM_HOSTS` | `localhost,127.0.0.1,100.83.103.1` | 允许拉流的主机白名单 | +| `STREAM_TASK_MAX_WORKERS` | `2` | 流任务并发 worker 数 | +| `STREAM_TASK_MAX_FRAMES` | `10` | 单任务最大抽帧数 | ## 说明 diff --git a/ai-service/app/config.py b/ai-service/app/config.py index 95593a0..825b9e3 100644 --- a/ai-service/app/config.py +++ b/ai-service/app/config.py @@ -8,3 +8,11 @@ MODEL_LABELS = os.getenv("MODEL_LABELS", "healthy,sick").split(",") # mock 模式返回的固定结果 MOCK_CLASS = os.getenv("MOCK_CLASS", "healthy") MOCK_CONFIDENCE = float(os.getenv("MOCK_CONFIDENCE", "0.95")) +INTERNAL_API_KEY = os.getenv("INTERNAL_API_KEY", "silk-internal-2026") +ALLOWED_STREAM_HOSTS = [ + item.strip() + for item in os.getenv("ALLOWED_STREAM_HOSTS", "localhost,127.0.0.1,100.83.103.1").split(",") + if item.strip() +] +STREAM_TASK_MAX_WORKERS = int(os.getenv("STREAM_TASK_MAX_WORKERS", "2")) +STREAM_TASK_MAX_FRAMES = int(os.getenv("STREAM_TASK_MAX_FRAMES", "10")) diff --git a/ai-service/app/main.py b/ai-service/app/main.py index 260dc61..dc006a9 100644 --- a/ai-service/app/main.py +++ b/ai-service/app/main.py @@ -7,12 +7,14 @@ from fastapi import FastAPI, File, HTTPException, Request, UploadFile from . import config from .detector import MockDetector, ONNXDetector +from .stream_tasks import StreamTaskWorker, is_allowed_stream_ref app = FastAPI(title="Silk AI Service", version="0.1.0") START_TIME = time.time() _lock = threading.Lock() _requests = 0 _latency_total = 0.0 +stream_worker = StreamTaskWorker(config.STREAM_TASK_MAX_WORKERS) @app.middleware("http") @@ -59,40 +61,35 @@ async def detect(file: UploadFile = File(...)): } -try: - import cv2 # noqa: E402 -except ImportError: - cv2 = None - - @app.post("/stream-detect") -async def stream_detect(body: dict): - """摄像头流 AI 巡检骨架（#23）：拉流抽帧 → 检测。 +def stream_detect(): + raise HTTPException(status_code=410, detail="请使用 POST /internal/stream-tasks") - 二期功能：正式接入前需补摄像头视角数据与巡检任务编排。 - """ - url = (body or {}).get("url") or "" - if not url: - raise HTTPException(status_code=400, detail="缺少 url（rtsp/http-flv 流地址）") - if cv2 is None: - raise HTTPException(status_code=503, detail="OpenCV 未安装，无法拉流") - capture = cv2.VideoCapture(url) - if not capture.isOpened(): - raise HTTPException(status_code=502, detail="无法连接视频流") - try: - frames = 0 - detections = [] - while frames < 3: - ok, frame = capture.read() - if not ok: - break - ok_encode, buf = cv2.imencode(".jpg", frame) - if ok_encode: - detections.extend(detector.detect(buf.tobytes())) - frames += 1 - finally: - capture.release() - return {"model": config.MODEL_MODE, "frames": frames, "detections": detections[:10]} + +@app.post("/internal/stream-tasks", status_code=202) +def create_stream_task(request: Request, body: dict): + if request.headers.get("x-internal-key") != config.INTERNAL_API_KEY: + raise HTTPException(status_code=401, detail="invalid internal key") + + task_id = (body or {}).get("taskId", "") + stream_ref = (body or {}).get("streamRef") + if not task_id or not isinstance(stream_ref, dict): + raise HTTPException(status_code=400, detail="缺少 taskId 或 streamRef") + if not is_allowed_stream_ref(stream_ref): + raise HTTPException(status_code=400, detail="streamRef 不在允许范围") + + stream_worker.submit(task_id, stream_ref) + return {"taskId": task_id, "status": "queued"} + + +@app.get("/internal/stream-tasks/{task_id}") +def get_stream_task(request: Request, task_id: str): + if request.headers.get("x-internal-key") != config.INTERNAL_API_KEY: + raise HTTPException(status_code=401, detail="invalid internal key") + task = stream_worker.get(task_id) + if not task: + raise HTTPException(status_code=404, detail="task not found") + return task @app.get("/metrics") diff --git a/ai-service/app/stream_tasks.py b/ai-service/app/stream_tasks.py new file mode 100644 index 0000000..9e47c2c --- /dev/null +++ b/ai-service/app/stream_tasks.py @@ -0,0 +1,101 @@ +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from urllib.parse import urlparse + +from . import config +from .detector import MockDetector, ONNXDetector + + +def _detector(): + if config.MODEL_MODE == "onnx": + return ONNXDetector(config.MODEL_PATH, tuple(config.MODEL_LABELS)) + return MockDetector(class_name=config.MOCK_CLASS, confidence=config.MOCK_CONFIDENCE) + + +def is_allowed_stream_ref(stream_ref): + url = (stream_ref or {}).get("url", "") + parsed = urlparse(url) + if parsed.scheme not in ("rtsp", "http", "https"): + return False + if parsed.username or parsed.password: + return False + host = (parsed.hostname or "").lower() + allowed = {item.strip().lower() for item in config.ALLOWED_STREAM_HOSTS} + return host in allowed + + +class StreamTaskWorker: + def __init__(self, max_workers=2): + self.executor = ThreadPoolExecutor(max_workers=max_workers) + self.tasks = {} + self.lock = threading.Lock() + + def submit(self, task_id, stream_ref): + with self.lock: + self.tasks[task_id] = { + "taskId": task_id, + "status": "queued", + "startedAt": time.time(), + "finishedAt": None, + "frames": 0, + "detections": [], + "error": None, + } + self.executor.submit(self._run, task_id, stream_ref) + + def get(self, task_id): + with self.lock: + task = self.tasks.get(task_id) + return dict(task) if task else None + + def _run(self, task_id, stream_ref): + try: + with self.lock: + task = self.tasks.get(task_id) + if task: + task["status"] = "running" + try: + import cv2 + except ImportError: + self._fail(task_id, "OpenCV 未安装，无法拉流") + return + + capture = cv2.VideoCapture(stream_ref.get("url", "")) + if not capture.isOpened(): + self._fail(task_id, "无法连接视频流") + return + + max_frames = min(int(stream_ref.get("maxFrames", config.STREAM_TASK_MAX_FRAMES)), config.STREAM_TASK_MAX_FRAMES) + detector = _detector() + frames = 0 + detections = [] + try: + while frames < max_frames: + ok, frame = capture.read() + if not ok: + break + ok_encode, buf = cv2.imencode(".jpg", frame) + if ok_encode: + detections.extend(detector.detect(buf.tobytes())) + frames += 1 + finally: + capture.release() + + with self.lock: + task = self.tasks.get(task_id) + if task: + task["status"] = "completed" + task["frames"] = frames + task["detections"] = detections[:10] + task["finishedAt"] = time.time() + except Exception as exc: + self._fail(task_id, str(exc)) + + def _fail(self, task_id, message): + with self.lock: + task = self.tasks.get(task_id) + if task: + task["status"] = "failed" + task["error"] = message + task["finishedAt"] = time.time() diff --git a/ai-service/tests/test_api.py b/ai-service/tests/test_api.py index 9b8e7c7..2173990 100644 --- a/ai-service/tests/test_api.py +++ b/ai-service/tests/test_api.py @@ -1,4 +1,8 @@ import base64 +import os + +os.environ.setdefault("INTERNAL_API_KEY", "test-internal-key") +os.environ.setdefault("ALLOWED_STREAM_HOSTS", "localhost,127.0.0.1,100.83.103.1") from fastapi.testclient import TestClient @@ -42,8 +46,44 @@ def test_detect_invalid_image_rejected(): assert r.status_code == 400 -def test_stream_detect_missing_url(): +def test_stream_detect_is_disabled(): r = client.post("/stream-detect", json={}) + assert r.status_code == 410 + + +def test_internal_stream_tasks_requires_internal_key(): + r = client.post("/internal/stream-tasks", json={}) + assert r.status_code == 401 + + +def test_internal_stream_tasks_accepts_internal_stream_ref(): + r = client.post( + "/internal/stream-tasks", + headers={"X-Internal-Key": "test-internal-key"}, + json={ + "taskId": "task-1", + "streamRef": { + "url": "rtsp://127.0.0.1:8554/live/1", + "maxFrames": 1, + }, + }, + ) + assert r.status_code == 202 + assert r.json()["taskId"] == "task-1" + + +def test_internal_stream_tasks_rejects_metadata_url(): + r = client.post( + "/internal/stream-tasks", + headers={"X-Internal-Key": "test-internal-key"}, + json={ + "taskId": "task-2", + "streamRef": { + "url": "http://169.254.169.254/latest/meta-data", + "maxFrames": 1, + }, + }, + ) assert r.status_code == 400 diff --git a/app/src/utils/ws.ts b/app/src/utils/ws.ts index 007c8d8..0f25562 100644 --- a/app/src/utils/ws.ts +++ b/app/src/utils/ws.ts @@ -1,17 +1,19 @@ import { WS_BASE_URL } from '@env'; +import { get } from '../api/client'; export type WSMessageHandler = (data: any) => void; class WebSocketManager { private ws: WebSocket | null = null; private token: string | null = null; + private ticket: string | null = null; private handlers: Set = new Set(); private reconnectTimer: ReturnType | null = null; private reconnectAttempts = 0; private maxReconnectAttempts = 10; private isManualClose = false; - connect(token: string): void { + async connect(token: string): Promise { // Close any existing connection before creating a new one if (this.ws) { this.isManualClose = true; @@ -27,14 +29,14 @@ class WebSocketManager { this.reconnectAttempts = 0; this.token = token; this.isManualClose = false; - this.doConnect(); + await this.refreshTicketAndConnect(); } private doConnect(): void { - if (!this.token) return; + if (!this.token || !this.ticket) return; const wsUrl = (WS_BASE_URL || 'ws://localhost:3000/ws').replace(/\?.*$/, ''); - const url = `${wsUrl}?token=${encodeURIComponent(this.token)}`; + const url = `${wsUrl}?ticket=${encodeURIComponent(this.ticket)}`; try { this.ws = new WebSocket(url); @@ -84,11 +86,23 @@ class WebSocketManager { this.reconnectTimer = setTimeout(() => { if (!this.isManualClose && this.token) { - this.doConnect(); + this.refreshTicketAndConnect(); } }, delay); } + private async refreshTicketAndConnect(): Promise { + if (!this.token) return; + try { + const { ticket } = await get<{ ticket: string }>('/ws/ticket'); + this.ticket = ticket; + this.doConnect(); + } catch (e) { + console.warn('[WS] Failed to fetch ticket:', e); + this.scheduleReconnect(); + } + } + disconnect(): void { this.isManualClose = true; if (this.reconnectTimer) { diff --git a/miniapp/src/api/config.ts b/miniapp/src/api/config.ts index 06c6e6a..4055c3b 100644 --- a/miniapp/src/api/config.ts +++ b/miniapp/src/api/config.ts @@ -43,24 +43,24 @@ export function getServerUrl(): string { // 保留导出以兼容现有代码（模块加载时的默认值，实际请求请使用 getApiBaseUrl()） export const API_BASE_URL = DEFAULT_API_BASE_URL; -export function getWsUrl(token: string): string { +export function getWsUrl(ticket: string): string { const baseUrl = getApiBaseUrl(); // 如果是完整 URL，从中推导 WS 地址 if (baseUrl.startsWith('http')) { const serverUrl = baseUrl.replace(/\/api\/v1\/?$/, ''); const wsUrl = serverUrl.replace(/^http/, 'ws'); - return `${wsUrl}/ws?token=${token}`; + return `${wsUrl}/ws?ticket=${ticket}`; } // H5 相对路径 - 使用 window.location 推导 if (isH5) { const protocol = typeof window !== 'undefined' && window.location.protocol === 'https:' ? 'wss:' : 'ws:'; const host = typeof window !== 'undefined' ? window.location.host : 'localhost:3000'; - return `${protocol}//${host}/ws?token=${token}`; + return `${protocol}//${host}/ws?ticket=${ticket}`; } - return `ws://100.83.103.1:3000/ws?token=${token}`; + return `ws://100.83.103.1:3000/ws?ticket=${ticket}`; } export const REQUEST_TIMEOUT = 15000; diff --git a/miniapp/src/utils/ws.ts b/miniapp/src/utils/ws.ts index 2a9577b..3ccea85 100644 --- a/miniapp/src/utils/ws.ts +++ b/miniapp/src/utils/ws.ts @@ -1,5 +1,6 @@ import Taro from '@tarojs/taro'; import { getWsUrl } from '@/api/config'; +import { get } from '@/api/request'; type MessageHandler = (data: unknown) => void; @@ -7,19 +8,20 @@ class WebSocketManager { private connected: boolean = false; private handlers: Map> = new Map(); private token: string = ''; + private ticket: string = ''; private reconnectAttempts: number = 0; private maxReconnectAttempts: number = 5; private reconnectTimer: ReturnType | null = null; private isManualClose: boolean = false; - connect(token: string): void { + async connect(token: string): Promise { this.token = token; this.isManualClose = false; - this.doConnect(); + await this.refreshTicketAndConnect(); } private doConnect(): void { - if (!this.token) { + if (!this.token || !this.ticket) { console.warn('[WS] 无 token，无法连接 WebSocket'); return; } @@ -34,7 +36,7 @@ class WebSocketManager { this.connected = false; } - const url = getWsUrl(this.token); + const url = getWsUrl(this.ticket); console.log('[WS] 正在连接:', url); // 注册全局回调（每次调用会替换上一次的回调） @@ -98,10 +100,22 @@ class WebSocketManager { clearTimeout(this.reconnectTimer); } this.reconnectTimer = setTimeout(() => { - this.doConnect(); + this.refreshTicketAndConnect(); }, delay); } + private async refreshTicketAndConnect(): Promise { + if (!this.token) return; + try { + const { ticket } = await get<{ ticket: string }>('/ws/ticket'); + this.ticket = ticket; + this.doConnect(); + } catch (e) { + console.error('[WS] 获取 ticket 失败:', e); + this.scheduleReconnect(); + } + } + on(type: string, handler: MessageHandler): () => void { if (!this.handlers.has(type)) { this.handlers.set(type, new Set()); diff --git a/server-go/cmd/server/main.go b/server-go/cmd/server/main.go index a4e1e0d..8467ba7 100644 --- a/server-go/cmd/server/main.go +++ b/server-go/cmd/server/main.go @@ -5,6 +5,7 @@ import ( "log/slog" "os" "strconv" + "strings" "time" "github.com/gin-gonic/gin" @@ -52,7 +53,7 @@ func main() { } // 5. 创建 WebSocket Hub - hub := ws.NewHub(cfg.JWTSecret) + hub := ws.NewHub(cfg.JWTSecret, ws.NewDBDeviceAuthorizer(db), strings.Split(cfg.WSAllowedOrigins, ",")) // 6. 创建并启动 MQTT 服务 mqttSvc := service.NewMQTTService(cfg.MQTT, db, iotdb, hub) @@ -89,6 +90,7 @@ func main() { // 11. API 路由组（经过 JWT auth 中间件，白名单路径自动跳过） api := r.Group("/api/v1") api.Use(middleware.Auth(cfg)) + api.GET("/ws/ticket", hub.HandleTicket) // 公开路由（auth 白名单中跳过鉴权） handler.RegisterVideoStreamRoutes(api, transcodeSvc, db, mediaSvc, cfg) diff --git a/server-go/internal/config/config.go b/server-go/internal/config/config.go index 7511824..66a5cc8 100644 --- a/server-go/internal/config/config.go +++ b/server-go/internal/config/config.go @@ -6,39 +6,40 @@ import ( // Config 全局配置，从环境变量加载 type Config struct { - AppEnv string `env:"APP_ENV" envDefault:"development"` - AllowDevAutoMigrate bool `env:"ALLOW_DEV_AUTOMIGRATE" envDefault:"false"` - PG string `env:"PG" envDefault:"postgresql://postgres:pan@localhost:5432/silk"` - Redis string `env:"REDIS" envDefault:"redis://:pan@localhost:6379"` - JWTSecret string `env:"JWT_SECRET" envDefault:"silk-secret-please-change-me"` - JWTExpiresIn string `env:"JWT_EXPIRES_IN" envDefault:"2h"` - MQTT string `env:"MQTT" envDefault:"mqtt://pan:pan@localhost:1883"` - IoTDBURL string `env:"IOTDB_URL" envDefault:"http://127.0.0.1:18081"` - S3Endpoint string `env:"S3_ENDPOINT" envDefault:"http://100.83.103.1:7480"` - S3AccessKey string `env:"S3_ACCESS_KEY" envDefault:"silk-app"` - S3SecretKey string `env:"S3_SECRET_KEY" envDefault:"Silk-App-Secret-2026!"` - S3Bucket string `env:"S3_BUCKET" envDefault:"silk-video-events"` - S3BucketImages string `env:"S3_BUCKET_IMAGES" envDefault:"silk-images"` - S3Region string `env:"S3_REGION" envDefault:"us-east-1"` - WVPAPIBase string `env:"WVP_API_BASE" envDefault:"http://localhost:18978"` - WVPUsername string `env:"WVP_USERNAME" envDefault:"admin"` - WVPPassword string `env:"WVP_PASSWORD" envDefault:"admin"` - ZLMAPIBase string `env:"ZLM_API_BASE" envDefault:"http://100.83.103.1:8081"` - ZLMSecret string `env:"ZLM_SECRET" envDefault:"su6TiedN2rVAmBbIDX0aa0QTiBJLBdcf"` - RecorderAPIBase string `env:"RECORDER_API_BASE" envDefault:"http://localhost:9090"` - AIServiceBase string `env:"AI_SERVICE_BASE" envDefault:"http://localhost:8000"` - WechatAppID string `env:"WECHAT_APPID" envDefault:""` - WechatSecret string `env:"WECHAT_SECRET" envDefault:""` - WechatTemplateAlarm string `env:"WECHAT_TEMPLATE_ALARM" envDefault:""` + AppEnv string `env:"APP_ENV" envDefault:"development"` + AllowDevAutoMigrate bool `env:"ALLOW_DEV_AUTOMIGRATE" envDefault:"false"` + WSAllowedOrigins string `env:"WS_ALLOWED_ORIGINS" envDefault:"http://localhost:5174,http://localhost:3000,http://127.0.0.1:5174,http://127.0.0.1:3000,http://100.83.103.1:5174,http://100.83.103.1:3000"` + PG string `env:"PG" envDefault:"postgresql://postgres:pan@localhost:5432/silk"` + Redis string `env:"REDIS" envDefault:"redis://:pan@localhost:6379"` + JWTSecret string `env:"JWT_SECRET" envDefault:"silk-secret-please-change-me"` + JWTExpiresIn string `env:"JWT_EXPIRES_IN" envDefault:"2h"` + MQTT string `env:"MQTT" envDefault:"mqtt://pan:pan@localhost:1883"` + IoTDBURL string `env:"IOTDB_URL" envDefault:"http://127.0.0.1:18081"` + S3Endpoint string `env:"S3_ENDPOINT" envDefault:"http://100.83.103.1:7480"` + S3AccessKey string `env:"S3_ACCESS_KEY" envDefault:"silk-app"` + S3SecretKey string `env:"S3_SECRET_KEY" envDefault:"Silk-App-Secret-2026!"` + S3Bucket string `env:"S3_BUCKET" envDefault:"silk-video-events"` + S3BucketImages string `env:"S3_BUCKET_IMAGES" envDefault:"silk-images"` + S3Region string `env:"S3_REGION" envDefault:"us-east-1"` + WVPAPIBase string `env:"WVP_API_BASE" envDefault:"http://localhost:18978"` + WVPUsername string `env:"WVP_USERNAME" envDefault:"admin"` + WVPPassword string `env:"WVP_PASSWORD" envDefault:"admin"` + ZLMAPIBase string `env:"ZLM_API_BASE" envDefault:"http://100.83.103.1:8081"` + ZLMSecret string `env:"ZLM_SECRET" envDefault:"su6TiedN2rVAmBbIDX0aa0QTiBJLBdcf"` + RecorderAPIBase string `env:"RECORDER_API_BASE" envDefault:"http://localhost:9090"` + AIServiceBase string `env:"AI_SERVICE_BASE" envDefault:"http://localhost:8000"` + WechatAppID string `env:"WECHAT_APPID" envDefault:""` + WechatSecret string `env:"WECHAT_SECRET" envDefault:""` + WechatTemplateAlarm string `env:"WECHAT_TEMPLATE_ALARM" envDefault:""` WechatTemplateInspection string `env:"WECHAT_TEMPLATE_INSPECTION" envDefault:""` QWeatherAPIKey string `env:"QWEATHER_API_KEY" envDefault:""` QWeatherLocation string `env:"QWEATHER_LOCATION" envDefault:""` QWeatherIntervalMin int `env:"QWEATHER_INTERVAL_MIN" envDefault:"30"` - InternalAPIKey string `env:"INTERNAL_API_KEY" envDefault:"silk-internal-2026"` - Port int `env:"PORT" envDefault:"3000"` - DefaultAdminUsername string `env:"DEFAULT_ADMIN_USERNAME" envDefault:"admin"` - DefaultAdminPassword string `env:"DEFAULT_ADMIN_PASSWORD" envDefault:"silk@123"` - DefaultAdminEmail string `env:"DEFAULT_ADMIN_EMAIL" envDefault:"admin@silk.local"` + InternalAPIKey string `env:"INTERNAL_API_KEY" envDefault:"silk-internal-2026"` + Port int `env:"PORT" envDefault:"3000"` + DefaultAdminUsername string `env:"DEFAULT_ADMIN_USERNAME" envDefault:"admin"` + DefaultAdminPassword string `env:"DEFAULT_ADMIN_PASSWORD" envDefault:"silk@123"` + DefaultAdminEmail string `env:"DEFAULT_ADMIN_EMAIL" envDefault:"admin@silk.local"` } // Load 从环境变量加载配置 diff --git a/server-go/internal/ws/authorizer.go b/server-go/internal/ws/authorizer.go new file mode 100644 index 0000000..9f17914 --- /dev/null +++ b/server-go/internal/ws/authorizer.go @@ -0,0 +1,59 @@ +package ws + +import ( + "errors" + "fmt" + + "silk-server-go/internal/model" + + "gorm.io/gorm" +) + +// DeviceAuthorizer 校验用户是否有权读取指定设备。 +type DeviceAuthorizer interface { + CanReadDevice(userID, deviceKey string) (bool, error) +} + +// DBDeviceAuthorizer 使用现有 RBAC 权限判断设备读取权；当前未做用户级资源 ACL。 +type DBDeviceAuthorizer struct { + db *gorm.DB +} + +func NewDBDeviceAuthorizer(db *gorm.DB) *DBDeviceAuthorizer { + return &DBDeviceAuthorizer{db: db} +} + +func (a *DBDeviceAuthorizer) CanReadDevice(userID, deviceKey string) (bool, error) { + if deviceKey == "" { + return false, nil + } + + var user model.User + if err := a.db.Where("id = ?", userID).First(&user).Error; err != nil { + return false, err + } + if user.Role == model.RoleAdmin { + return true, nil + } + + var count int64 + err := a.db.Table("role_permissions"). + Joins("JOIN permissions ON permissions.id = role_permissions.permission_id"). + Where("role_permissions.role = ? AND permissions.code = ?", user.Role, "device:read"). + Count(&count).Error + return count > 0, err +} + +func (h *Hub) authorizeSubscription(userID, deviceKey string) error { + if h.authorizer == nil { + return errors.New("device authorizer not configured") + } + ok, err := h.authorizer.CanReadDevice(userID, deviceKey) + if err != nil { + return err + } + if !ok { + return fmt.Errorf("device subscription denied: %s", deviceKey) + } + return nil +} diff --git a/server-go/internal/ws/gateway.go b/server-go/internal/ws/gateway.go index 0324f98..13129a8 100644 --- a/server-go/internal/ws/gateway.go +++ b/server-go/internal/ws/gateway.go @@ -4,82 +4,103 @@ import ( "encoding/json" "log/slog" "net/http" - "strings" "sync" "time" "github.com/gin-gonic/gin" - "github.com/golang-jwt/jwt/v5" "github.com/gorilla/websocket" "silk-server-go/internal/service" ) -var upgrader = websocket.Upgrader{ - CheckOrigin: func(r *http.Request) bool { return true }, - ReadBufferSize: 1024, - WriteBufferSize: 1024, -} - // client WebSocket 客户端 type client struct { - conn *websocket.Conn - rooms map[string]bool // 订阅的房间（device:） - send chan []byte + conn *websocket.Conn + userID string + username string + role string + rooms map[string]bool // 订阅的房间（device:） + send chan []byte } // Hub WebSocket 中心，管理客户端和房间 type Hub struct { - jwtSecret string - clients map[*client]bool - mu sync.RWMutex + authorizer DeviceAuthorizer + allowedOrigins []string + clients map[*client]bool + mu sync.RWMutex + tickets map[string]*wsTicket + ticketsMu sync.Mutex } // NewHub 创建 Hub -func NewHub(jwtSecret string) *Hub { +func NewHub(jwtSecret string, authorizer DeviceAuthorizer, allowedOrigins []string) *Hub { return &Hub{ - jwtSecret: jwtSecret, - clients: make(map[*client]bool), + authorizer: authorizer, + allowedOrigins: allowedOrigins, + clients: make(map[*client]bool), + tickets: make(map[string]*wsTicket), } } +func (h *Hub) originAllowed(origin string) bool { + if origin == "" { + return true + } + for _, allowed := range h.allowedOrigins { + if origin == allowed { + return true + } + } + return false +} + +// HandleTicket 签发一次性 WebSocket ticket，需在 JWT 鉴权路由后注册。 +func (h *Hub) HandleTicket(c *gin.Context) { + userVal, ok := c.Get("user") + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "未认证"}) + return + } + userMap, ok := userVal.(map[string]interface{}) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "用户信息无效"}) + return + } + userID, _ := userMap["sub"].(string) + username, _ := userMap["username"].(string) + role, _ := userMap["role"].(string) + if userID == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "用户信息无效"}) + return + } + + c.JSON(http.StatusOK, gin.H{ + "ticket": h.IssueTicket(userID, username, role), + "expiresIn": int(wsTicketTTL.Seconds()), + }) +} + // HandleWebSocket 处理 WebSocket 连接（Gin handler） func (h *Hub) HandleWebSocket(c *gin.Context) { - // 验证 JWT（从 query.auth.token / query.token / header.Authorization 获取） - tokenStr := "" - if t := c.Query("auth.token"); t != "" { - tokenStr = t - } else if t := c.Query("token"); t != "" { - tokenStr = t - } else if auth := c.GetHeader("Authorization"); strings.HasPrefix(auth, "Bearer ") { - tokenStr = auth[7:] - } - - if tokenStr == "" { - conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) - if err == nil { - conn.WriteJSON(map[string]interface{}{"event": "auth.fail", "data": map[string]bool{"ok": false}}) - conn.Close() - } + if !h.originAllowed(c.GetHeader("Origin")) { + c.JSON(http.StatusForbidden, gin.H{"error": "origin not allowed"}) return } - // 验证 token - claims := jwt.MapClaims{} - token, err := jwt.ParseWithClaims(tokenStr, claims, func(t *jwt.Token) (interface{}, error) { - return []byte(h.jwtSecret), nil - }, jwt.WithValidMethods([]string{"HS256"})) - - if err != nil || !token.Valid { - conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) - if err == nil { - conn.WriteJSON(map[string]interface{}{"event": "auth.fail", "data": map[string]bool{"ok": false}}) - conn.Close() - } + ticket := h.ticketFromQuery(c.Request.URL.RawQuery) + info, ok := h.consumeTicket(ticket) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid or expired ws ticket"}) return } // 升级为 WebSocket + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + ReadBufferSize: 1024, + WriteBufferSize: 1024, + } conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { slog.Warn("WebSocket 升级失败", "err", err) @@ -87,9 +108,12 @@ func (h *Hub) HandleWebSocket(c *gin.Context) { } cl := &client{ - conn: conn, - rooms: make(map[string]bool), - send: make(chan []byte, 256), + conn: conn, + userID: info.userID, + username: info.username, + role: info.role, + rooms: make(map[string]bool), + send: make(chan []byte, 256), } h.mu.Lock() @@ -98,8 +122,12 @@ func (h *Hub) HandleWebSocket(c *gin.Context) { // 发送 auth.ok h.sendJSON(cl, "auth.ok", map[string]interface{}{ - "ok": true, - "user": claims, + "ok": true, + "user": map[string]interface{}{ + "sub": info.userID, + "username": info.username, + "role": info.role, + }, }) go h.readPump(cl) @@ -145,6 +173,14 @@ func (h *Hub) readPump(cl *client) { switch data.Event { case "subscribe.device": + if err := h.authorizeSubscription(cl.userID, data.DeviceKey); err != nil { + h.sendJSON(cl, "subscribe.denied", map[string]interface{}{ + "ok": false, + "deviceKey": data.DeviceKey, + "error": err.Error(), + }) + continue + } room := "device:" + data.DeviceKey h.mu.Lock() cl.rooms[room] = true @@ -204,12 +240,9 @@ func (h *Hub) BroadcastTelemetry(deviceKey string, data interface{}) { room := "device:" + deviceKey for cl := range h.clients { - // 推送到设备房间 if cl.rooms[room] { h.sendJSON(cl, "telemetry", data) } - // 全局推送 - h.sendJSON(cl, "telemetry.all", data) } } @@ -223,14 +256,10 @@ func (h *Hub) BroadcastAlarm(event service.AlarmEvent) { for cl := range h.clients { if isRecovery { - // 恢复通知 - h.sendJSON(cl, "alarm.recovery", event) if event.DeviceKey != "" && cl.rooms[room] { h.sendJSON(cl, "alarm.device.recovery", event) } } else { - // 告警触发 - h.sendJSON(cl, "alarm", event) if event.DeviceKey != "" && cl.rooms[room] { h.sendJSON(cl, "alarm.device", event) } @@ -246,7 +275,6 @@ func (h *Hub) BroadcastDeviceStatus(deviceKey string, status string) { room := "device:" + deviceKey data := map[string]interface{}{"deviceKey": deviceKey, "status": status} for cl := range h.clients { - h.sendJSON(cl, "device.status", data) if cl.rooms[room] { h.sendJSON(cl, "device.status.device", data) } diff --git a/server-go/internal/ws/gateway_test.go b/server-go/internal/ws/gateway_test.go new file mode 100644 index 0000000..926b82e --- /dev/null +++ b/server-go/internal/ws/gateway_test.go @@ -0,0 +1,120 @@ +package ws + +import ( + "errors" + "testing" + + "silk-server-go/internal/service" +) + +type fakeAuthorizer struct { + allowed map[string]bool +} + +func (f fakeAuthorizer) CanReadDevice(userID, deviceKey string) (bool, error) { + if f.allowed == nil { + return false, nil + } + return f.allowed[userID+"|"+deviceKey], nil +} + +func newTestHub() *Hub { + return NewHub( + "test-secret", + fakeAuthorizer{allowed: map[string]bool{"user-1|device-a": true}}, + []string{"http://localhost:5174"}, + ) +} + +func TestWebSocketOriginRejectsUnknownOrigin(t *testing.T) { + hub := newTestHub() + if hub.originAllowed("http://evil.example") { + t.Fatal("unknown origin should be rejected") + } + if !hub.originAllowed("http://localhost:5174") { + t.Fatal("configured origin should be allowed") + } +} + +func TestWebSocketTicketIsOneTimeAndExpires(t *testing.T) { + hub := newTestHub() + ticket := hub.IssueTicket("user-1", "admin", "User One") + if ticket == "" { + t.Fatal("expected non-empty ticket") + } + + info, ok := hub.consumeTicket(ticket) + if !ok || info.userID != "user-1" { + t.Fatalf("first ticket consume failed: ok=%v info=%+v", ok, info) + } + if _, ok := hub.consumeTicket(ticket); ok { + t.Fatal("ticket should be single-use") + } +} + +func TestWebSocketRejectsLongLivedJWTQuery(t *testing.T) { + hub := newTestHub() + if hub.ticketFromQuery("") != "" { + t.Fatal("missing ticket should not be accepted") + } + if hub.ticketFromQuery("token=eyJhbGciOiJIUzI1NiJ9.abc") != "" { + t.Fatal("long-lived token query must not be accepted") + } +} + +func TestWebSocketDeviceAuthorizerRejectsUnauthorized(t *testing.T) { + hub := newTestHub() + if err := hub.authorizeSubscription("user-1", "device-a"); err != nil { + t.Fatalf("authorized device rejected: %v", err) + } + if err := hub.authorizeSubscription("user-1", "device-b"); err == nil { + t.Fatal("unauthorized device should be rejected") + } +} + +func TestBroadcastTelemetryDoesNotLeakGlobal(t *testing.T) { + hub := newTestHub() + subscribed := &client{rooms: map[string]bool{"device:device-a": true}, send: make(chan []byte, 1)} + unsubscribed := &client{rooms: map[string]bool{}, send: make(chan []byte, 1)} + hub.clients[subscribed] = true + hub.clients[unsubscribed] = true + + hub.BroadcastTelemetry("device-a", map[string]interface{}{"value": 1}) + + if len(subscribed.send) == 0 { + t.Fatal("subscribed client should receive telemetry") + } + if len(unsubscribed.send) != 0 { + t.Fatal("unsubscribed client must not receive global telemetry") + } +} + +func TestBroadcastAlarmDoesNotLeakGlobal(t *testing.T) { + hub := newTestHub() + subscribed := &client{rooms: map[string]bool{"device:device-a": true}, send: make(chan []byte, 1)} + unsubscribed := &client{rooms: map[string]bool{}, send: make(chan []byte, 1)} + hub.clients[subscribed] = true + hub.clients[unsubscribed] = true + + hub.BroadcastAlarm(service.AlarmEvent{DeviceKey: "device-a", Code: "high"}) + + if len(unsubscribed.send) != 0 { + t.Fatal("unsubscribed client must not receive global alarm") + } +} + +type brokenAuthorizer struct{} + +func (brokenAuthorizer) CanReadDevice(userID, deviceKey string) (bool, error) { + return false, errors.New("authorizer unavailable") +} + +func TestWebSocketAuthorizerErrorIsDenied(t *testing.T) { + hub := NewHub("test-secret", brokenAuthorizer{}, nil) + if err := hub.authorizeSubscription("user-1", "device-a"); err == nil { + t.Fatal("authorizer error should deny subscription") + } + if _, ok := hub.consumeTicket("missing"); ok { + t.Fatal("missing ticket should not be consumed") + } +} diff --git a/server-go/internal/ws/ticket.go b/server-go/internal/ws/ticket.go new file mode 100644 index 0000000..3115d08 --- /dev/null +++ b/server-go/internal/ws/ticket.go @@ -0,0 +1,57 @@ +package ws + +import ( + "crypto/rand" + "encoding/hex" + "net/url" + "time" +) + +const wsTicketTTL = 60 * time.Second + +type wsTicket struct { + userID string + username string + role string + expiresAt time.Time + used bool +} + +// IssueTicket 签发一次性 WebSocket ticket，避免长期 JWT 进入 URL。 +func (h *Hub) IssueTicket(userID, username, role string) string { + buf := make([]byte, 32) + if _, err := rand.Read(buf); err != nil { + return "" + } + ticket := hex.EncodeToString(buf) + + h.ticketsMu.Lock() + h.tickets[ticket] = &wsTicket{ + userID: userID, + username: username, + role: role, + expiresAt: time.Now().Add(wsTicketTTL), + } + h.ticketsMu.Unlock() + return ticket +} + +func (h *Hub) consumeTicket(ticket string) (wsTicket, bool) { + h.ticketsMu.Lock() + defer h.ticketsMu.Unlock() + + info, ok := h.tickets[ticket] + if !ok || info.used || time.Now().After(info.expiresAt) { + return wsTicket{}, false + } + info.used = true + return *info, true +} + +func (h *Hub) ticketFromQuery(rawQuery string) string { + values, err := url.ParseQuery(rawQuery) + if err != nil { + return "" + } + return values.Get("ticket") +} diff --git a/后续工作计划.md b/后续工作计划.md index 5a5b82b..74d3dc1 100644 --- a/后续工作计划.md +++ b/后续工作计划.md @@ -13,7 +13,7 @@ | Wave 1 | P0 安全与正确性 | Task 2 版本化数据库迁移与 Schema 启动门禁 | 部分可用 | 开发库升级已验证，真实 down 回滚演练待后续 | | Wave 1 | P0 安全与正确性 | Task 3 移除默认密钥与默认管理员密码 | 延后到最后（跳过） | 用户 2026-08-13 明确要求跳过并留到最后 | | Wave 1 | P0 安全与正确性 | Task 4 收口视频访问与摄像头密钥输出 | 部分可用 | 待开发服务器部署联调 | -| Wave 1 | P0 安全与正确性 | Task 5 修复 WebSocket 越权与 AI 流 SSRF | 未开始 | 无 | +| Wave 1 | P0 安全与正确性 | Task 5 修复 WebSocket 越权与 AI 流 SSRF | 部分可用 | 待开发服务器部署与真实 WS/AI 联调 | | Wave 1 | P0 安全与正确性 | Task 6 修复 AI 风险语义并隔离 Mock 数据 | 未开始 | 无 | | Wave 1 | P0 安全与正确性 | Task 7 修订 qPCR 判读与检测质控 | 未开始 | 需领域专家确认 | | Wave 2 | 工程可靠性 | Task 8 建立可靠通知、吊销与跨实例状态 | 未开始 | 无 | diff --git a/开发交接记录.md b/开发交接记录.md index 30462b7..6a4d61b 100644 --- a/开发交接记录.md +++ b/开发交接记录.md @@ -890,3 +890,34 @@ MVP 沿用 IoTDB（现状）；TDengine 作为生产规模化候选（先基准 - 本任务前分支提交为 `440cf80`；回滚可还原 Task 4 提交； - 若已部署流代理版本，回滚需恢复旧二进制并重启，无需数据库变更；播放地址需由客户端重新请求。 + +--- + +## 2026-08-14 整改 Task 5：修复 WebSocket 越权与 AI 流 SSRF + +### 做了什么 + +- WebSocket 改为一次性 ticket 连接：`GET /api/v1/ws/ticket` 签发 60 秒单次 ticket，`/ws` 不再接受 JWT query，客户端统一使用 ticket； +- 新增 `DeviceAuthorizer` / `DBDeviceAuthorizer`，`subscribe.device` 必须校验用户 `device:read` 权限；无权订阅返回 `subscribe.denied`； +- WebSocket 增加 Origin 白名单，`WS_ALLOWED_ORIGINS` 可配置； +- 移除普通连接的 `telemetry.all`、全局 `alarm`、全局 `device.status` 广播，只推送给已授权订阅的设备房间； +- APP/小程序 WS 客户端先请求 ticket 再连接，避免长期 JWT 进入 URL； +- AI 服务新增 `POST /internal/stream-tasks`：要求 `X-Internal-Key`、`taskId`、`streamRef`，限制拉流 host 白名单、并发 worker 数和最大帧数；旧 `/stream-detect` 任意 URL 入口改为 410。 + +### 设计思路与决策依据 + +- WebSocket 播放/订阅场景不适合把长期 JWT 放进 URL；采用服务端签发的一次性 ticket 降低代理日志泄露和重放风险； +- 当前没有用户级资源 ACL，设备订阅先用 RBAC `device:read` 做对象授权；组织/房间级 ACL 留待后续任务； +- AI 流任务改为内部服务接口和受限 worker，避免公网客户端直接传任意 URL，也避免在 FastAPI 事件循环中阻塞式 `VideoCapture.read()`。 + +### 验证结果 + +- `scripts/verify.ps1` 最终 exit 0：Go test/vet/build、Web、小程序、APP、AI pytest 11/11 均通过； +- 新增 Go WS 测试覆盖非法 Origin、ticket 单次/过期、JWT query 拒绝、无权设备订阅、全局广播不泄露； +- 新增 AI 测试覆盖内部 key 校验、任务创建、metadata/公网地址拒绝和旧接口禁用； +- 未部署开发服务器，未做真实 WS 推送和摄像头拉流联调。 + +### 回滚点 + +- 本任务前分支提交为 `1a29def`；回滚可还原 Task 5 提交； +- WS/AI 改动无数据库 schema 变更；若已部署，恢复旧二进制并重启即可，但需同步回滚客户端 WS 连接方式。