feat: 修复 WebSocket 越权与 AI 流 SSRF

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")