234 lines
9.2 KiB
Python
234 lines
9.2 KiB
Python
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)
|