feat: 修复 WebSocket 越权与 AI 流 SSRF

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