blender_runtime.py

"""安装可信 Blender 运行时，并在可取消的隔离子进程中执行固定渲染脚本。"""

from __future__ import annotations

import contextlib
import hashlib
import json
import os
import posixpath
import re
import shutil
import signal
import subprocess
import tarfile
import tempfile
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Iterator


_INSTALL_LOCK = threading.Lock()
_VERSION_RE = re.compile(r"^[0-9]+\.[0-9]+\.[0-9]+$")
_WORKER = Path(__file__).resolve().with_name("render_worker.py")
_MAX_ARCHIVE_BYTES = 4 * 1024**3
_MAX_EXTRACTED_BYTES = 12 * 1024**3


def _valid_executable(path: Path) -> bool:
    return path.is_file() and os.access(path, os.X_OK)


@contextlib.contextmanager
def _installation_file_lock(path: Path) -> Iterator[None]:
    """Linux 进程间锁；与线程锁一起避免重复下载和并发解压。"""
    if os.name != "posix":
        raise RuntimeError("自动安装 Blender 目前只支持 Linux；请设置 BLENDER_BIN。")
    import fcntl

    with path.open("a+b") as lock_file:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)


def _download(url: str, destination: Path, maximum: int) -> None:
    """HTTPS 流式下载；首次失败后最多重试两次，每次重写临时文件。"""
    last_error: Exception | None = None
    for attempt in range(3):
        try:
            request = urllib.request.Request(
                url, headers={"User-Agent": "Blender-Render-Space/1.0"}
            )
            started = time.monotonic()
            with urllib.request.urlopen(request, timeout=60) as response:
                if not response.geturl().startswith("https://"):
                    raise RuntimeError("下载发生了非 HTTPS 重定向，已拒绝。")
                content_length = response.headers.get("Content-Length")
                if content_length and int(content_length) > maximum:
                    raise RuntimeError("下载内容超过允许大小。")
                total = 0
                with destination.open("wb") as output:
                    while block := response.read(1024 * 1024):
                        total += len(block)
                        if total > maximum:
                            raise RuntimeError("下载内容超过允许大小。")
                        if time.monotonic() - started > 1800:
                            raise TimeoutError("单次下载超过 30 分钟。")
                        output.write(block)
            return
        except (OSError, urllib.error.URLError, ValueError, RuntimeError) as error:
            last_error = error
            destination.unlink(missing_ok=True)
            if attempt < 2:
                time.sleep(attempt + 1)
    raise RuntimeError(f"无法下载 Blender 官方文件：{url}；原因：{last_error}") from last_error


def _expected_digest(checksums: Path, filename: str) -> str:
    matches: set[str] = set()
    for line in checksums.read_text(encoding="utf-8").splitlines():
        match = re.fullmatch(r"([a-fA-F0-9]{64})\s+\*?(.+?)\s*", line)
        if match and match.group(2) == filename:
            matches.add(match.group(1).lower())
    if len(matches) != 1:
        raise RuntimeError(f"官方 SHA-256 清单中没有唯一匹配的 {filename}，安装已停止。")
    return matches.pop()


def _safe_extract(archive: Path, destination: Path, archive_root: str) -> Path:
    """只接受官方根目录中的普通文件、目录和不越界的内部链接。"""
    with tarfile.open(archive, "r:xz") as bundle:
        members = bundle.getmembers()
        if len(members) > 200_000:
            raise RuntimeError("Blender 压缩包文件数量异常。")
        expanded_size = 0
        seen: set[str] = set()
        for member in members:
            name = member.name.rstrip("/")
            parts = name.split("/")
            if (
                not name
                or "\\" in name
                or name.startswith("/")
                or any(part in {"", ".", ".."} for part in parts)
                or parts[0] != archive_root
            ):
                raise RuntimeError(f"Blender 压缩包含有不安全路径：{member.name}")
            if name in seen:
                raise RuntimeError(f"Blender 压缩包含有重复路径：{member.name}")
            seen.add(name)
            if member.isfile():
                expanded_size += member.size
                if expanded_size > _MAX_EXTRACTED_BYTES:
                    raise RuntimeError("Blender 压缩包解压大小超过限制。")
            elif member.issym() or member.islnk():
                target = member.linkname
                if not target or target.startswith("/") or "\\" in target:
                    raise RuntimeError(f"Blender 压缩包含有不安全链接：{name}")
                # 符号链接相对于所在目录，tar 硬链接相对于压缩包根目录。
                resolved = posixpath.normpath(
                    posixpath.join(posixpath.dirname(name), target)
                    if member.issym()
                    else target
                )
                if not (resolved == archive_root or resolved.startswith(archive_root + "/")):
                    raise RuntimeError(f"Blender 压缩包链接指向安装目录之外：{name}")
            elif not member.isdir():
                raise RuntimeError(f"Blender 压缩包包含不允许的特殊文件：{name}")
        # Python 3.12 的 data filter 还会检查已解压链接链，拒绝实际越界。
        bundle.extractall(destination, members=members, filter="data")
    return destination / archive_root


def ensure_blender() -> Path:
    """优先使用 BLENDER_BIN、系统 Blender，否则校验后安装官方 Linux x64 版本。"""
    configured = os.environ.get("BLENDER_BIN", "").strip()
    if configured:
        binary = Path(configured).expanduser().resolve()
        if not _valid_executable(binary):
            raise RuntimeError(f"BLENDER_BIN 不存在或不可执行：{binary}")
        return binary
    installed = shutil.which("blender")
    if installed:
        return Path(installed).resolve()

    import platform

    if platform.system() != "Linux" or platform.machine().lower() not in {"x86_64", "amd64"}:
        raise RuntimeError("自动下载只支持 Linux x64；请安装 Blender 并设置 BLENDER_BIN。")
    version = os.environ.get("BLENDER_VERSION", "5.2.2").strip()
    if not _VERSION_RE.fullmatch(version):
        raise RuntimeError("BLENDER_VERSION 必须是类似 5.2.2 的三段数字版本号。")
    cache = Path(
        os.environ.get("BLENDER_CACHE", str(Path.home() / ".cache" / "blender-render-space"))
    ).expanduser().resolve()
    installation = cache / f"blender-{version}"
    binary = installation / "blender"
    if _valid_executable(binary):
        return binary

    with _INSTALL_LOCK:
        cache.mkdir(parents=True, exist_ok=True)
        with _installation_file_lock(cache / f".install-{version}.lock"):
            if _valid_executable(binary):
                return binary
            release = ".".join(version.split(".")[:2])
            base_url = f"https://download.blender.org/release/Blender{release}/"
            archive_root = f"blender-{version}-linux-x64"
            archive_name = archive_root + ".tar.xz"
            try:
                with tempfile.TemporaryDirectory(prefix=".install-", dir=cache) as temporary:
                    staging = Path(temporary)
                    # 部分云平台无法访问 Blender 下载站。优先使用管理员随仓库
                    # 部署的安装包；仍须提供校验清单，并执行同样的安全解压检查。
                    vendor = Path(__file__).resolve().parent / "vendor"
                    bundled = vendor / archive_name
                    if bundled.is_file():
                        archive = bundled
                        if archive.stat().st_size > _MAX_ARCHIVE_BYTES:
                            raise RuntimeError("Blender 安装包超过允许大小。")
                        checksum_file = vendor / f"blender-{version}.sha256"
                        expected = _expected_digest(checksum_file, archive_name)
                    else:
                        checksum_file = staging / "checksums.sha256"
                        archive = staging / archive_name
                        _download(base_url + f"blender-{version}.sha256", checksum_file, 1024**2)
                        expected = _expected_digest(checksum_file, archive_name)
                        _download(base_url + archive_name, archive, _MAX_ARCHIVE_BYTES)
                    with archive.open("rb") as source:
                        actual = hashlib.file_digest(source, "sha256").hexdigest()
                    if actual != expected:
                        raise RuntimeError("Blender 安装包 SHA-256 校验失败，已停止安装。")
                    extracted = _safe_extract(archive, staging, archive_root)
                    if not _valid_executable(extracted / "blender"):
                        raise RuntimeError("官方安装包中未找到可执行的 Blender。")
                    # 与最终目录位于同一文件系统，rename 只会暴露完整安装。
                    previous = staging / "previous-installation"
                    if installation.exists() or installation.is_symlink():
                        installation.rename(previous)
                    try:
                        os.replace(extracted, installation)
                    except OSError:
                        if previous.exists() or previous.is_symlink():
                            previous.rename(installation)
                        raise
            except (OSError, tarfile.TarError, ValueError) as error:
                raise RuntimeError(f"Blender 自动安装失败：{error}。可设置 BLENDER_BIN 使用已有安装。") from error
    return binary


def _worker_environment(config: dict | None = None) -> dict[str, str]:
    """移除应用服务凭据，保留 CUDA、NVIDIA 和动态库搜索路径。"""
    forbidden_prefixes = (
        "HF_", "HUGGINGFACE_", "HUGGING_FACE_", "AWS_", "AZURE_", "GOOGLE_",
        "GCP_", "OPENAI_", "ANTHROPIC_", "SPACES_", "GRADIO_AUTH",
    )
    sensitive = re.compile(r"TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIAL|API_KEY|ACCESS_KEY|PRIVATE_KEY|AUTHORIZATION")
    environment = {
        name: value
        for name, value in os.environ.items()
        if not name.upper().startswith(forbidden_prefixes) and not sensitive.search(name.upper())
    }
    # 部分 OpenMP / TBB 库默认按宿主逻辑核心数开线程，需限制过度并行。
    threads = _thread_count()
    environment["OMP_NUM_THREADS"] = str(threads)
    environment["OPENBLAS_NUM_THREADS"] = str(threads)
    environment["MKL_NUM_THREADS"] = str(threads)
    environment["PYTHONNOUSERSITE"] = "1"
    config = config or {}
    if config.get("engine", "EEVEE") == "EEVEE":
        # Use Mesa's software OpenGL only for explicitly selected CPU jobs.
        for name in ("LIBGL_ALWAYS_SOFTWARE", "GALLIUM_DRIVER", "MESA_LOADER_DRIVER_OVERRIDE",
                     "__EGL_VENDOR_LIBRARY_FILENAMES", "EGL_PLATFORM"):
            environment.pop(name, None)
        if config.get("device", "CPU") == "CPU":
            environment.update(LIBGL_ALWAYS_SOFTWARE="true", GALLIUM_DRIVER="llvmpipe",
                               EGL_PLATFORM="surfaceless", LP_NUM_THREADS=str(threads))
            mesa = Path("/usr/share/glvnd/egl_vendor.d/50_mesa.json")
            if mesa.is_file():
                environment["__EGL_VENDOR_LIBRARY_FILENAMES"] = str(mesa)
    return environment


def _thread_count() -> int:
    try:
        threads = int(os.environ.get("BLENDER_THREADS", "4"))
    except ValueError as error:
        raise RuntimeError("BLENDER_THREADS 必须为 1 至 256 之间的整数。") from error
    if not 1 <= threads <= 256:
        raise RuntimeError("BLENDER_THREADS 必须为 1 至 256 之间的整数。")
    return threads


def _stop_process_group(process: subprocess.Popen[Any]) -> None:
    """先请求终止，再结束整个进程组，包括仍在运行的后代进程。"""
    if os.name == "posix":
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
        try:
            process.wait(timeout=3)
        except subprocess.TimeoutExpired:
            pass
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
    else:
        # 本地 Windows 指定 BLENDER_BIN 时，同样终止进程树。
        subprocess.run(
            ["taskkill", "/PID", str(process.pid), "/T", "/F"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False, timeout=10,
        )
    try:
        process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        process.kill()
        process.wait(timeout=5)


def _log_tail(path: Path) -> str:
    try:
        with path.open("rb") as source:
            source.seek(0, os.SEEK_END)
            source.seek(max(0, source.tell() - 12_000))
            return source.read().decode("utf-8", errors="replace").strip()
    except OSError:
        return "日志不可读取。"


def run_task(
    task: dict,
    timeout: int,
    log_path: Path,
    cancel_path: Path | None = None,
) -> dict:
    """以可信 worker 执行任务；成功时读取 result_path，失败保留完整日志。"""
    if not isinstance(task, dict) or not task.get("result_path"):
        raise ValueError("渲染任务必须包含 result_path。")
    if not isinstance(timeout, (int, float)) or isinstance(timeout, bool) or timeout <= 0:
        raise ValueError("渲染超时必须是正数秒数。")
    if not _WORKER.is_file():
        raise RuntimeError(f"缺少可信 Blender 工作脚本：{_WORKER.name}")
    binary = ensure_blender()
    log_path = Path(log_path).resolve()
    log_path.parent.mkdir(parents=True, exist_ok=True)
    result_path = Path(task["result_path"]).resolve()
    result_path.parent.mkdir(parents=True, exist_ok=True)
    cancel_path = Path(cancel_path) if cancel_path is not None else None
    if cancel_path is not None and cancel_path.exists():
        raise RuntimeError("渲染任务已取消。")
    result_path.unlink(missing_ok=True)
    task_copy = {**task, "result_path": str(result_path)}
    descriptor, task_filename = tempfile.mkstemp(prefix=".render-task-", suffix=".json", dir=log_path.parent)
    process: subprocess.Popen[Any] | None = None
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as output:
            json.dump(task_copy, output, ensure_ascii=False, allow_nan=False)
        command = [
            str(binary), "--background", "--factory-startup", "--disable-autoexec",
            "--threads", str(_thread_count()), "--python-exit-code", "1",
            "--gpu-backend", (task.get("config", {}).get("graphics_backend", "opengl")
                              if task.get("action") != "inspect" and
                              task.get("config", {}).get("engine") == "EEVEE" and
                              task.get("config", {}).get("device") != "CPU" else "opengl"),
            "--python", str(_WORKER), "--", task_filename,
        ]
        with log_path.open("wb") as log:
            process = subprocess.Popen(
                command,
                stdin=subprocess.DEVNULL,
                stdout=log,
                stderr=subprocess.STDOUT,
                cwd=log_path.parent,
                env=_worker_environment(task.get("config", {})),
                start_new_session=os.name == "posix",
                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0,
            )
            started = time.monotonic()
            while process.poll() is None:
                if cancel_path is not None and cancel_path.exists():
                    _stop_process_group(process)
                    raise RuntimeError("渲染任务已取消；已完成的帧保留，可继续渲染。")
                if time.monotonic() - started >= timeout:
                    _stop_process_group(process)
                    raise TimeoutError(f"渲染任务超过 {timeout} 秒，进程组已终止；已完成的帧保留。")
                time.sleep(0.15)
        if process.returncode != 0:
            _stop_process_group(process)
            raise RuntimeError(f"Blender 执行失败（退出码 {process.returncode}）：\n{_log_tail(log_path)}")
        if not result_path.is_file() or result_path.stat().st_size > 16 * 1024**2:
            raise RuntimeError(f"Blender 未生成有效结果文件：\n{_log_tail(log_path)}")
        try:
            result = json.loads(result_path.read_text(encoding="utf-8"))
        except (OSError, ValueError) as error:
            raise RuntimeError(f"Blender 返回的结果 JSON 无效：{error}") from error
        if not isinstance(result, dict):
            raise RuntimeError("Blender 结果必须是 JSON 对象。")
        return result
    finally:
        if process is not None and process.poll() is None:
            _stop_process_group(process)
        Path(task_filename).unlink(missing_ok=True)