feat: 修复 WebSocket 越权与 AI 流 SSRF

2026-08-14 00:45:22 +08:00
parent 1a29def480
commit 839ba91354
17 changed files with 614 additions and 136 deletions
@@ -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()