media_proxy_service.py

"""公开媒体读取服务；仅代理既有允许目录与扩展名，不处理工作流下载。"""
import logging
import mimetypes
import os
import time
import urllib.parse
import uuid
from fastapi import HTTPException
from fastapi.responses import JSONResponse
from proxy_download_work import ProxyDownloadWork

_public_media_work = ProxyDownloadWork(16)
public_media_endpoint = _public_media_work.endpoint
start_public_media = _public_media_work.start
stop_public_media = _public_media_work.stop
close_public_media = _public_media_work.close

logger = logging.getLogger("ComfyUI-Ranking")


def read_public_media(url=None, path=None, *, store, download, file_response, error_response):
    """公开媒体服务：使用服务凭据读取既有允许目录中的媒体"""

    # 兼容处理：如果本地发来的是已经被污染的老版本 HF 直链，我们自动将其转换为相对路径
    if url and url.startswith("https://huggingface.co/datasets/"):
        try:
            path = url.split("resolve/main/")[-1]
            path = urllib.parse.unquote(path)
        except Exception:
            raise HTTPException(status_code=400, detail="无效的 HF 原链接格式")

    if not path:
        raise HTTPException(status_code=400, detail="缺少路径参数")

    # 🛡️ 路径穿越防护：禁止包含 .. 和绝对路径
    if ".." in path or path.startswith("/"):
        raise HTTPException(status_code=400, detail="非法路径")

    # 🛡️ 绝对的安全红线：限制只能代理下载图片目录，严禁黑客通过此接口下载 users.json 或账本数据！
    allowed_dirs = ["uploads/", "avatars/", "covers/"]
    if not any(path.startswith(d) for d in allowed_dirs):
        raise HTTPException(status_code=403, detail="非法访问：该接口仅允许代理图片资源")

    media_extensions = {"jpg", "jpeg", "png", "gif", "webp", "mp4", "webm", "mov", "wav", "mp3", "m4a", "ogg"}
    if path.rsplit(".", 1)[-1].lower() not in media_extensions:
        raise HTTPException(status_code=403, detail="该接口不允许访问非媒体文件")

    hf_token = os.environ.get("HF_TOKEN")
    dataset_repo_id = store.DATASET_REPO_ID

    try:
        # hf_hub_download 会自动利用云端容器的缓存，只有第一次会去真实请求 Dataset
        cached_file_path = download(
            repo_id=dataset_repo_id,
            repo_type="dataset",
            filename=path,
            token=hf_token
        )
        # 智能识别文件类型 (image/jpeg, image/png 等)
        content_type, _ = mimetypes.guess_type(cached_file_path)
        return file_response(cached_file_path, media_type=content_type or "application/octet-stream")
    except Exception as e:
        error_id = f"ERR_{int(time.time())}_{uuid.uuid4().hex[:8]}"
        upstream_status = getattr(getattr(e, "response", None), "status_code", None)
        if not isinstance(upstream_status, int):
            upstream_status = getattr(e, "code", None)
        status_text = f" status={upstream_status}" if isinstance(upstream_status, int) and 100 <= upstream_status <= 599 else ""
        logger.warning(f"IMAGE_PROXY_FAIL | error_id={error_id} | error_type={type(e).__name__}{status_text}")
        return JSONResponse(
            status_code=502,
            content=error_response(502, "图片暂时无法获取，请稍后重试", "upstream_error", error_id=error_id),
        )