from __future__ import annotations import hashlib import json import os import re import ssl import tempfile import urllib.parse import urllib.request from pathlib import Path from typing import Any _ALLOWED_KINDS = { "lora": "loras", "checkpoint": "checkpoints", "unet": "diffusion_models", "vae": "vae", "clip": "text_encoders", "gguf": "diffusion_models", } _SHA256 = re.compile(r"^[0-9a-f]{64}$") _FILENAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._+ -]{0,180}\.(?:safetensors|gguf)$") _MAX_BYTES = 24 * 1024 * 1024 * 1024 class ArtifactInstallError(RuntimeError): pass def validate_request(payload: dict[str, Any]) -> dict[str, str]: kind = str(payload.get("artifact_type") or "").strip().lower() url = str(payload.get("source_url") or "").strip() sha256 = str(payload.get("sha256") or "").strip().lower() filename = str(payload.get("filename") or "").strip() license_name = str(payload.get("license") or "").strip() source_id = str(payload.get("source_id") or "").strip() parsed = urllib.parse.urlsplit(url) if kind not in _ALLOWED_KINDS: raise ArtifactInstallError("artifact_type is not allow-listed") if parsed.scheme != "https" or not parsed.hostname: raise ArtifactInstallError("source_url must be an absolute HTTPS URL") if parsed.username or parsed.password or parsed.fragment: raise ArtifactInstallError("source_url credentials/fragments are forbidden") if not _SHA256.fullmatch(sha256): raise ArtifactInstallError("sha256 must be a pinned lowercase SHA-256 digest") expected_extension = ".gguf" if kind == "gguf" else ".safetensors" if ( not _FILENAME.fullmatch(filename) or "/" in filename or "\\" in filename or not filename.casefold().endswith(expected_extension) ): raise ArtifactInstallError( f"filename must be a bounded {expected_extension} basename for {kind}" ) if not license_name or len(license_name) > 240: raise ArtifactInstallError("license provenance is required") if not source_id or len(source_id) > 240: raise ArtifactInstallError("source_id/version provenance is required") return { "artifact_type": kind, "source_url": url, "sha256": sha256, "filename": filename, "license": license_name, "source_id": source_id, } def preflight(payload: dict[str, Any], *, model_root: Path) -> dict[str, Any]: request = validate_request(payload) root = model_root.resolve(strict=True) if not root.is_dir() or not os.path.ismount(root): raise ArtifactInstallError("model root must be an existing persistent mount") destination_dir = (root / _ALLOWED_KINDS[request["artifact_type"]]).resolve(strict=True) if root not in destination_dir.parents: raise ArtifactInstallError("resolved destination escaped the model root") if not os.access(destination_dir, os.W_OK | os.X_OK): raise ArtifactInstallError("artifact destination is not writable by the AppOps service") destination = destination_dir / request["filename"] existing_digest = None if destination.exists(): if not destination.is_file() or destination.is_symlink(): raise ArtifactInstallError("destination is not a regular artifact") existing_digest = _hash_file(destination) if existing_digest != request["sha256"]: raise ArtifactInstallError("destination already exists with a different digest") evidence = {"request": request, "model_root": str(root), "existing_digest": existing_digest} fingerprint = hashlib.sha256(json.dumps(evidence, sort_keys=True, separators=(",", ":")).encode()).hexdigest() return {"ready": True, "fingerprint": fingerprint, "artifact_type": request["artifact_type"], "filename": request["filename"], "sha256": request["sha256"], "model_root": str(root), "already_installed": existing_digest is not None} def _hash_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _write_receipt(request: dict[str, str], destination: Path, size: int, receipt_root: Path) -> Path: receipt_root.mkdir(parents=True, exist_ok=True) receipt = { "schema": "privatemuse-comfyui-artifact.v1", "artifact_type": request["artifact_type"], "filename": request["filename"], "sha256": request["sha256"], "bytes": size, "source_id": request["source_id"], "source_host": urllib.parse.urlsplit(request["source_url"]).hostname, "license": request["license"], "destination": str(destination), } receipt_path = receipt_root / f'{request["sha256"]}.json' if receipt_path.exists(): return receipt_path with tempfile.NamedTemporaryFile( mode="w", prefix=".receipt-", suffix=".json", dir=receipt_root, delete=False, encoding="utf-8", ) as handle: json.dump(receipt, handle, sort_keys=True, separators=(",", ":")) handle.write("\n") handle.flush() os.fsync(handle.fileno()) temporary = Path(handle.name) try: os.link(temporary, receipt_path) except FileExistsError: pass finally: temporary.unlink(missing_ok=True) return receipt_path def install( payload: dict[str, Any], *, model_root: Path, receipt_root: Path, ) -> dict[str, Any]: request = validate_request(payload) root = model_root.resolve() destination_dir = (root / _ALLOWED_KINDS[request["artifact_type"]]).resolve() if root not in destination_dir.parents: raise ArtifactInstallError("resolved destination escaped the model root") destination_dir.mkdir(parents=True, exist_ok=True) destination = destination_dir / request["filename"] if destination.is_symlink(): raise ArtifactInstallError("destination symlinks are forbidden") if destination.exists(): digest = _hash_file(destination) if digest != request["sha256"]: raise ArtifactInstallError( "destination already exists with a different digest" ) receipt_path = _write_receipt(request, destination, destination.stat().st_size, receipt_root) return { "status": "already_installed", "path": str(destination), "sha256": digest, "bytes": destination.stat().st_size, "receipt": str(receipt_path), } request_obj = urllib.request.Request( request["source_url"], headers={"User-Agent": "ITWorx-AppOps/PrivateMuse-artifact-installer"}, ) hasher = hashlib.sha256() total = 0 tmp_path: Path | None = None try: with urllib.request.urlopen( request_obj, timeout=30, context=ssl.create_default_context(), ) as response: final = urllib.parse.urlsplit(response.geturl()) if final.scheme != "https" or not final.hostname: raise ArtifactInstallError("download redirected outside HTTPS") declared = response.headers.get("Content-Length") if declared and int(declared) > _MAX_BYTES: raise ArtifactInstallError( "artifact exceeds the 24 GiB bounded installer limit" ) with tempfile.NamedTemporaryFile( prefix=".appops-", suffix=".part", dir=destination_dir, delete=False, ) as handle: tmp_path = Path(handle.name) while True: chunk = response.read(1024 * 1024) if not chunk: break total += len(chunk) if total > _MAX_BYTES: raise ArtifactInstallError( "artifact exceeded the 24 GiB bounded installer limit" ) hasher.update(chunk) handle.write(chunk) handle.flush() os.fsync(handle.fileno()) digest = hasher.hexdigest() if digest != request["sha256"]: raise ArtifactInstallError( "downloaded artifact SHA-256 did not match the pinned digest" ) # Link is atomic and fails if another installer created the final name. # os.replace would silently overwrite a different digest in that race. try: os.link(tmp_path, destination) except FileExistsError as exc: if destination.is_symlink() or not destination.is_file() or _hash_file(destination) != digest: raise ArtifactInstallError("destination appeared with a different digest") from exc tmp_path.unlink() tmp_path = None receipt_path = _write_receipt(request, destination, total, receipt_root) return { "status": "installed", "path": str(destination), "sha256": digest, "bytes": total, "receipt": str(receipt_path), } finally: if tmp_path is not None: tmp_path.unlink(missing_ok=True)