64 lines
2.7 KiB
Python
64 lines
2.7 KiB
Python
|
"""compute ``download_file`` 端点安全回归测试。
|
|||
|
|
|
|||
|
|
修复前 ``file_id`` 直接拼进 glob 模式且不校验路径包含关系，可通过 ``../``
|
|||
|
|
穿越出 upload 目录，并在 Linux 上跟随符号链接读取任意文件。
|
|||
|
|
修复后：file_id 仅允许字母/数字/下划线/连字符，返回前对解析后的路径
|
|||
|
|
做 upload 根目录包含性校验。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import pytest
|
|||
|
|
from fastapi.testclient import TestClient
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _make_client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> TestClient:
|
|||
|
|
monkeypatch.setenv("TRAINING_LOG_ROOT", str(tmp_path / "logs"))
|
|||
|
|
monkeypatch.setenv("YG_FT_DATA_ROOT", str(tmp_path / "data"))
|
|||
|
|
monkeypatch.setenv("COMPUTE_EXECUTION_MODE", "simulator")
|
|||
|
|
monkeypatch.setenv("COMPUTE_AUTH_ENABLED", "false")
|
|||
|
|
monkeypatch.delenv("ENABLE_DOCS", raising=False)
|
|||
|
|
from compute.api.main import create_app
|
|||
|
|
|
|||
|
|
return TestClient(create_app())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _upload_root(tmp_path: Path) -> Path:
|
|||
|
|
root = tmp_path / "data" / "uploads"
|
|||
|
|
root.mkdir(parents=True, exist_ok=True)
|
|||
|
|
return root
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_download_legit_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
|
|
client = _make_client(tmp_path, monkeypatch)
|
|||
|
|
(_upload_root(tmp_path) / "file_123456_hello.txt").write_text("HELLO-DOWNLOAD", encoding="utf-8")
|
|||
|
|
response = client.get("/modelTF/compute/files/file_123456/download")
|
|||
|
|
assert response.status_code == 200
|
|||
|
|
assert response.content == b"HELLO-DOWNLOAD"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_download_rejects_traversal_file_id(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
|
|
client = _make_client(tmp_path, monkeypatch)
|
|||
|
|
outside = tmp_path / "secret" / "passwd_1.txt"
|
|||
|
|
outside.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
outside.write_text("TOP-SECRET", encoding="utf-8")
|
|||
|
|
for file_id in ["..", "file.123", "..%2F..%2Fsecret%2Fpasswd", "file%20name"]:
|
|||
|
|
response = client.get(f"/modelTF/compute/files/{file_id}/download")
|
|||
|
|
assert response.status_code in (400, 404), f"file_id={file_id!r} -> {response.status_code}"
|
|||
|
|
assert b"TOP-SECRET" not in response.content
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_download_blocks_symlink_escape(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
|
|
client = _make_client(tmp_path, monkeypatch)
|
|||
|
|
upload_root = _upload_root(tmp_path)
|
|||
|
|
outside = tmp_path / "secret.txt"
|
|||
|
|
outside.write_text("TOP-SECRET", encoding="utf-8")
|
|||
|
|
try:
|
|||
|
|
(upload_root / "file_999999_link.txt").symlink_to(outside)
|
|||
|
|
except OSError:
|
|||
|
|
pytest.skip("symlink creation not permitted on this platform")
|
|||
|
|
response = client.get("/modelTF/compute/files/file_999999/download")
|
|||
|
|
assert response.status_code == 404
|
|||
|
|
assert b"TOP-SECRET" not in response.content
|