xiaoxia-saas/apps/api/app/api/routes/gpu_relay.py

"""GPU 编码回传 relay 端点。

P4000 编码完成后通过 HTTP PUT 把结果 mp4 写到这里；Worker 在发起 GPU 请求时携带
带签名（token + 随机 key）的 URL，等待 P4000 写入后用同 URL 把文件 GET 回本地。

安全：
- 生产环境必须配置 GPU_ENCODE_RELAY_SECRET；token=xxx 查询参数必须匹配。
- key 为随机 hex，无法被枚举。
- 写入/读取后 worker 会调用 DELETE 主动清理；文件落地在 generated-files/gpu_relay/，
  跟 generated-files 同卷，nginx 已对 generated-files 做静态挂载，但 gpu_relay/ 子目录
  通过本接口走鉴权，不直接暴露为静态目录（文件名随机 + token 保护双重保险）。
"""
from __future__ import annotations

import logging
import os
import secrets
import time
import uuid
from pathlib import Path
from typing import Optional

from fastapi import APIRouter, HTTPException, Query, Request
from fastapi.responses import FileResponse, Response

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/internal/gpu-relay", tags=["Internal-GpuRelay"])

_DEFAULT_SECRET_LOGGED = False


def _relay_dir() -> Path:
    base = os.getenv("GENERATED_FILES_DIR", "/app/generated")
    sub = os.getenv("GPU_ENCODE_RELAY_DIR", "gpu_relay")
    p = Path(base) / sub
    p.mkdir(parents=True, exist_ok=True)
    return p


def _secret() -> str:
    global _DEFAULT_SECRET_LOGGED
    secret = (os.getenv("GPU_ENCODE_RELAY_SECRET", "") or "").strip()
    if not secret:
        env = (os.getenv("APP_ENV", os.getenv("ENV", "development"))).lower()
        if env in ("production", "prod"):
            # Production: raise so deployment fails fast
            raise RuntimeError("GPU_ENCODE_RELAY_SECRET must be set in production")
        # Dev: ephemeral random secret, log once
        secret = os.environ.setdefault("GPU_ENCODE_RELAY_SECRET", secrets.token_urlsafe(32))
        if not _DEFAULT_SECRET_LOGGED:
            logger.warning(
                "[gpu-relay] GPU_ENCODE_RELAY_SECRET not set; using ephemeral dev token (%s...)",
                secret[:8],
            )
            _DEFAULT_SECRET_LOGGED = True
    return secret


def _safe_key(key: str) -> str:
    """只允许合法文件名字符，防 path traversal。"""
    k = key.strip()
    if not k or "/" in k or "\\" in k or k in (".", "..") or not all(
        c.isalnum() or c in "-_" for c in k
    ):
        raise HTTPException(status_code=400, detail="invalid key")
    return k


def _check_token(tok: Optional[str]) -> None:
    if not tok or tok != _secret():
        raise HTTPException(status_code=401, detail="unauthorized")


# ── Worker 侧：生成一个一次性 PUT URL ───────────────────────────────────
def build_relay_put_url(base_url: str, key: str, secret: str) -> str:
    """给 P4000 用的 PUT URL（含 token）。"""
    return f"{base_url.rstrip('/')}/api/v1/internal/gpu-relay/{key}?token={secret}"


def build_relay_get_url(base_url: str, key: str, secret: str) -> str:
    """Worker 取回结果用的 GET URL。"""
    return build_relay_put_url(base_url, key, secret)


def generate_key() -> str:
    return uuid.uuid4().hex


# ── HTTP endpoints ──────────────────────────────────────────────────────


@router.put("/{key}")
async def put_object(
    key: str,
    request: Request,
    token: Optional[str] = Query(None),
):
    _check_token(token)
    safe = _safe_key(key)
    dst = _relay_dir() / safe
    tmp = dst.with_suffix(dst.suffix + ".part")
    size = 0
    t0 = time.time()
    try:
        with open(tmp, "wb") as f:
            async for chunk in request.stream():
                f.write(chunk)
                size += len(chunk)
        os.replace(tmp, dst)
    except Exception as e:  # noqa: BLE001
        if tmp.exists():
            try:
                tmp.unlink()
            except OSError:
                pass
        logger.exception("[gpu-relay] PUT failed key=%s", safe)
        raise HTTPException(status_code=500, detail=f"write failed: {e}") from e
    logger.info(
        "[gpu-relay] PUT key=%s size=%d took=%.2fs",
        safe, size, time.time() - t0,
    )
    return {"ok": True, "key": safe, "size": size}


@router.get("/{key}")
async def get_object(
    key: str,
    token: Optional[str] = Query(None),
):
    _check_token(token)
    safe = _safe_key(key)
    path = _relay_dir() / safe
    if not path.exists():
        raise HTTPException(status_code=404, detail="not found")
    return FileResponse(
        path=path,
        media_type="video/mp4",
        filename=f"{safe}.mp4",
    )


@router.head("/{key}")
async def head_object(
    key: str,
    token: Optional[str] = Query(None),
):
    _check_token(token)
    safe = _safe_key(key)
    path = _relay_dir() / safe
    if not path.exists():
        return Response(status_code=404)
    return Response(
        status_code=200,
        media_type="video/mp4",
        headers={"Content-Length": str(path.stat().st_size)},
    )


@router.delete("/{key}")
async def delete_object(
    key: str,
    token: Optional[str] = Query(None),
):
    _check_token(token)
    safe = _safe_key(key)
    path = _relay_dir() / safe
    try:
        if path.exists():
            path.unlink()
    except OSError as e:
        raise HTTPException(status_code=500, detail=f"delete failed: {e}") from e
    return {"ok": True, "key": safe}