mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 23:37:28 +08:00
4e75db9c23
The endpoint was reachable via GET with all parameters taken from the query string, so any web page in the operator's browser could trigger a profiling run (CSRF via <img>) with an arbitrary client-controlled output_dir, making the engine create directories and write traces anywhere the process has access to. - Accept POST + JSON body only; query parameters are no longer read and GET returns 405, which closes the browser-forgery vectors - Constrain output_dir to the DEEPINFER_TORCH_PROFILER_DIR subtree (resolved with realpath, so ../ and symlink redirections are rejected) - Bind to 127.0.0.1 by default; multi-node collection now requires an explicit DEEPINFER_PROFILE_HTTP_BIND=0.0.0.0 curl usage is unchanged: POST with a JSON body, as documented in the module docstring and the startup log line. Co-Authored-By: Kscc <noreply@owtffssent.com>
99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
import json
|
||
import threading
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from typing import Any
|
||
|
||
import pytest
|
||
|
||
from deepinfer.engine.profile_server import ProfileManager
|
||
|
||
pytest.importorskip("fastapi", reason="profile http server tests need the fastapi/uvicorn dev dependencies")
|
||
|
||
|
||
@pytest.fixture
|
||
def manager(tmp_path: Any, monkeypatch: Any) -> ProfileManager:
|
||
"""Bare ProfileManager whose request slot is already occupied.
|
||
|
||
Occupying the slot (via _req_ready) makes every valid request return 409
|
||
immediately, so tests never wait on the engine-side pickup timeout.
|
||
"""
|
||
monkeypatch.setenv("DEEPINFER_TORCH_PROFILER_DIR", str(tmp_path / "traces"))
|
||
from deepinfer import environ
|
||
|
||
environ.rearm_all()
|
||
|
||
manager = object.__new__(ProfileManager)
|
||
manager._pending_req = None
|
||
manager._req_ready = threading.Event()
|
||
manager._req_ready.set()
|
||
manager._io_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="profile-io-test")
|
||
return manager
|
||
|
||
|
||
@pytest.fixture
|
||
def client(manager: ProfileManager) -> Any:
|
||
from fastapi.testclient import TestClient
|
||
|
||
with TestClient(manager._build_app()) as c:
|
||
yield c
|
||
|
||
|
||
def _post(client: Any, body: dict[str, Any] | None = None, **kwargs: Any) -> Any:
|
||
content = json.dumps(body) if body is not None else None
|
||
return client.post("/profile", content=content, **kwargs)
|
||
|
||
|
||
@pytest.mark.ci
|
||
def test_get_is_rejected(client: Any) -> None:
|
||
# <img src="http://engine:port/profile?..."> 是 GET：必须 405，浏览器无法借此触发采集
|
||
response = client.get("/profile", params={"num_steps": 5})
|
||
assert response.status_code == 405
|
||
|
||
|
||
@pytest.mark.ci
|
||
def test_query_params_are_not_honored(client: Any, tmp_path: Any) -> None:
|
||
# form POST 可以把参数塞进 URL，但 query 已不被读取：output_dir 越界参数不生效，
|
||
# 请求按默认值（合法）处理并落入 409
|
||
response = client.post("/profile", params={"output_dir": str(tmp_path / "outside")})
|
||
assert response.status_code == 409
|
||
|
||
|
||
@pytest.mark.ci
|
||
def test_output_dir_outside_profiler_dir_is_rejected(client: Any, tmp_path: Any) -> None:
|
||
response = _post(client, {"output_dir": str(tmp_path / "outside")})
|
||
assert response.status_code == 400
|
||
assert "DEEPINFER_TORCH_PROFILER_DIR" in response.json()["message"]
|
||
|
||
|
||
@pytest.mark.ci
|
||
def test_output_dir_traversal_is_rejected(client: Any, tmp_path: Any) -> None:
|
||
# ../ 越界在 realpath 之后同样落到子树外
|
||
base = tmp_path / "traces"
|
||
response = _post(client, {"output_dir": f"{base}/../evil"})
|
||
assert response.status_code == 400
|
||
|
||
|
||
@pytest.mark.ci
|
||
def test_output_dir_inside_profiler_dir_is_accepted(client: Any, tmp_path: Any) -> None:
|
||
response = _post(client, {"output_dir": str(tmp_path / "traces" / "run1")})
|
||
assert response.status_code == 409 # slot occupied; request itself parsed fine
|
||
|
||
|
||
@pytest.mark.ci
|
||
def test_default_output_dir_is_profiler_dir(client: Any, tmp_path: Any) -> None:
|
||
response = _post(client, {"num_steps": 2})
|
||
assert response.status_code == 409
|
||
assert (tmp_path / "traces").is_dir() # makedirs ran on the configured base dir
|
||
|
||
|
||
@pytest.mark.ci
|
||
def test_non_object_body_is_rejected(client: Any) -> None:
|
||
response = client.post("/profile", content=json.dumps([1, 2]))
|
||
assert response.status_code == 400
|
||
|
||
|
||
@pytest.mark.ci
|
||
def test_unknown_fields_are_rejected(client: Any) -> None:
|
||
response = _post(client, {"num_steps": 2, "evil": True})
|
||
assert response.status_code == 400
|