Mujoco_WASM/training_server/mobile_manipulator/packages.py

"""Authenticated browser scene snapshots; no caller-supplied filesystem paths."""

import hashlib
import json
import re
import shutil
import stat
import tempfile
import threading
import zipfile
from pathlib import Path, PurePosixPath
from xml.etree import ElementTree as ET

from .config import ROBOTS, TASK

MAX_UPLOAD = 128 * 1024**2
MAX_EXPANDED = 512 * 1024**2


class MobilePackages:
    def __init__(self, root):
        self.root = Path(root)
        self.lock = threading.Lock()

    def receive(self, stream, length):
        if not 0 < length <= MAX_UPLOAD:
            raise ValueError("场景上传大小必须在 1–128 MiB 内")
        self.root.mkdir(parents=True, exist_ok=True)
        with self.lock, tempfile.TemporaryDirectory(dir=self.root) as temporary:
            temp = Path(temporary)
            archive = temp / "upload.zip"
            digest = hashlib.sha256()
            with archive.open("wb") as out:
                remaining = length
                while remaining:
                    chunk = stream.read(min(1024 * 1024, remaining))
                    if not chunk:
                        raise ValueError("场景上传不完整")
                    digest.update(chunk)
                    out.write(chunk)
                    remaining -= len(chunk)
            package_id = digest.hexdigest()
            destination = self.root / package_id
            if destination.is_dir():
                return {"id": package_id, **self.describe(package_id)}
            if sum(p.is_dir() and len(p.name) == 64 for p in self.root.iterdir()) >= 20:
                raise ValueError("场景快照已达 20 份，请在停止服务后清理 mobile_packages")
            directory = temp / "package"
            directory.mkdir()
            try:
                with zipfile.ZipFile(archive) as z:
                    infos = z.infolist()
                    if len(infos) > 10000 or sum(i.file_size for i in infos) > MAX_EXPANDED:
                        raise ValueError("场景展开超出 512 MiB / 10000 文件上限")
                    names = set()
                    for info in infos:
                        name = info.filename
                        path = PurePosixPath(name)
                        if (
                            not name
                            or "\\" in name
                            or ":" in name
                            or path.is_absolute()
                            or ".." in path.parts
                            or str(path) in names
                            or stat.S_ISLNK(info.external_attr >> 16)
                        ):
                            raise ValueError("场景包含不安全或重复路径")
                        names.add(str(path))
                        if not info.is_dir():
                            target = directory.joinpath(*path.parts)
                            target.parent.mkdir(parents=True, exist_ok=True)
                            with z.open(info) as src, target.open("wb") as dst:
                                shutil.copyfileobj(src, dst)
                self._validate(directory)
            except (
                zipfile.BadZipFile,
                KeyError,
                TypeError,
                ET.ParseError,
                OSError,
                NotImplementedError,
                RuntimeError,
            ) as error:
                raise ValueError(f"场景快照无效：{error}") from error
            directory.rename(destination)
            return {"id": package_id, **self.describe(package_id)}

    def path(self, package_id):
        if not isinstance(package_id, str) or not re.fullmatch(r"[0-9a-f]{64}", package_id):
            raise ValueError("mobilePackageId 无效")
        directory = self.root / package_id
        if not directory.is_dir():
            raise ValueError("移动操作场景不存在，请重新开始训练以同步场景")
        return directory

    def describe(self, package_id):
        return self._validate(self.path(package_id))

    @staticmethod
    def _validate(directory):
        for name in ("environment.json", "robot.json", "task.json"):
            if (directory / name).stat().st_size > 128 * 1024:
                raise ValueError("场景契约 JSON 超过 128 KiB")
        metadata = json.loads((directory / "environment.json").read_text())
        robot_bytes = (directory / "robot.json").read_bytes()
        robot = json.loads(robot_bytes)
        if (
            not isinstance(metadata, dict)
            or not isinstance(robot, dict)
            or not isinstance(robot.get("id"), str)
            or robot.get("id") not in ROBOTS
            or robot != ROBOTS[robot["id"]]
            or json.loads((directory / "task.json").read_text()) != TASK
            or metadata.get("robotId") != robot["id"]
            or metadata.get("taskId") != TASK["id"]
        ):
            raise ValueError("场景机器人/任务契约与已注册变体不匹配")
        if metadata.get("mujoco") != "3.11.0":
            raise ValueError("场景需要与浏览器一致的 MuJoCo 3.11.0")
        scene_name = metadata.get("scene")
        if not isinstance(scene_name, str) or "\\" in scene_name or ":" in scene_name:
            raise ValueError("scene 路径无效")
        scene = (directory / scene_name).resolve()
        if not scene.is_relative_to(directory.resolve()) or not scene.is_file():
            raise ValueError("scene 越界或不存在")
        if scene.stat().st_size > 16 * 1024**2:
            raise ValueError("场景 MJCF 超过 16 MiB")
        data = scene.read_bytes()
        if b"<!DOCTYPE" in data.upper() or b"<!ENTITY" in data.upper():
            raise ValueError("场景不接受 DTD/entity")
        root = ET.fromstring(data)
        if (
            root.tag != "mujoco"
            or root.find(".//include") is not None
            or root.find(".//plugin") is not None
        ):
            raise ValueError("必须上传展开后的 MJCF，不允许 include/plugin")
        compiler = root.find("compiler")
        for element in root.iter():
            for key, value in element.attrib.items():
                # Includes texture cubemap fileleft/fileright/fileup/... references.
                is_file = key.startswith("file")
                if not is_file and key not in ("meshdir", "texturedir", "assetdir"):
                    continue
                if "\\" in value or ":" in value or PurePosixPath(value).is_absolute():
                    raise ValueError("场景资产必须是包内相对路径")
                prefix = ""
                if is_file and compiler is not None:
                    prefix = compiler.get("meshdir" if element.tag == "mesh" else "texturedir", "")
                    prefix = prefix or compiler.get("assetdir", "")
                target = (scene.parent / prefix / value).resolve()
                if not target.is_relative_to(directory.resolve()):
                    raise ValueError("场景资产引用越界")
                if is_file and not target.is_file():
                    raise ValueError(f"场景缺少资产：{value}")
        return {
            "robotId": robot["id"],
            "sceneSha256": hashlib.sha256(data).hexdigest(),
            "robotConfigSha256": hashlib.sha256(robot_bytes).hexdigest(),
        }