v6ole/silk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 修复 WebSocket 越权与 AI 流 SSRF

Browse Source
This commit is contained in:
weijuesen
2026-08-14 00:45:22 +08:00
parent 1a29def480
commit 839ba91354
17 changed files with 614 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ai-service/README.md
+5
View File
@@ -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` | 单任务最大抽帧数 |
## 说明
ai-service/app/config.py
+8
View File
@@ -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"))
ai-service/app/main.py
+29 -32
View File
@@ -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")
ai-service/app/stream_tasks.py
+101
View File
@@ -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()
ai-service/tests/test_api.py
+41 -1
View File
@@ -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