diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 16b3207256..7809661827 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -37,6 +37,7 @@ from core.inference.sd_cpp_args import ( SdCppModelFiles, build_sd_cpp_command, ) +from utils.process_lifetime import child_popen_kwargs logger = logging.getLogger(__name__) @@ -122,8 +123,12 @@ def find_sd_cpp_binary() -> Optional[str]: if hit: return hit - # 3. Default install root (sibling of ~/.unsloth/llama.cpp). - hit = _first_file(_layout_candidates(Path.home() / ".unsloth" / "stable-diffusion.cpp")) + # 3. Default install root. Honors UNSLOTH_STUDIO_HOME / STUDIO_HOME exactly like the + # installer's default_install_dir(), so a binary installed under a custom Studio root + # is found without also having to set UNSLOTH_SD_CPP_PATH. + studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") + default_base = Path(studio_home).parent if studio_home else Path.home() / ".unsloth" + hit = _first_file(_layout_candidates(default_base / "stable-diffusion.cpp")) if hit: return hit @@ -232,6 +237,9 @@ class SdCppEngine: text = True, errors = "replace", env = run_env, + # Bind the child to the backend's lifetime (PR_SET_PDEATHSIG on Linux), so a + # long sd-cli denoise is SIGKILLed if the backend dies instead of orphaning. + **child_popen_kwargs(), ) # Drain stdout on a background thread and wait on the PROCESS, not the stream: # iterating proc.stdout directly blocks until the stream closes, so a sd-cli that @@ -261,6 +269,12 @@ class SdCppEngine: finally: if proc.poll() is None: proc.kill() + # Reap the SIGKILLed child, or it lingers as a zombie until this process + # exits (the cancel/timeout branches kill without waiting otherwise). + try: + proc.wait(timeout = 5.0) + except Exception: # noqa: BLE001 + pass # The process has exited; let the reader finish draining the buffered output. reader.join(timeout = 5.0) diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index 171ea3e810..b009cd37a6 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -17,7 +17,15 @@ _STUDIO = Path(__file__).resolve().parents[2] if str(_STUDIO) not in sys.path: sys.path.insert(0, str(_STUDIO)) -from install_sd_cpp_prebuilt import default_install_dir, resolve_release_asset # noqa: E402 +import zipfile # noqa: E402 + +import pytest # noqa: E402 + +from install_sd_cpp_prebuilt import ( # noqa: E402 + _safe_extractall, + default_install_dir, + resolve_release_asset, +) # A real stable-diffusion.cpp latest-release asset list. _ASSETS = [ @@ -114,3 +122,45 @@ def test_default_install_dir_is_sibling_of_llama(monkeypatch): d = default_install_dir() assert d.name == "stable-diffusion.cpp" assert d.parent.name == ".unsloth" + + +# ── safe extraction (Zip-Slip guard) ───────────────────────────────────────── + + +def test_safe_extractall_rejects_path_traversal(tmp_path): + target = tmp_path / "install" + target.mkdir() + archive = tmp_path / "evil.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("sd-cli", b"ok") + zf.writestr("../escape.txt", b"pwned") # escapes the install dir + with zipfile.ZipFile(archive) as zf: + with pytest.raises(RuntimeError, match = "unsafe path"): + _safe_extractall(zf, target) + assert not (tmp_path / "escape.txt").exists() + + +def test_safe_extractall_extracts_normal_members(tmp_path): + target = tmp_path / "install" + target.mkdir() + archive = tmp_path / "ok.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("build/bin/sd-cli", b"ok") + with zipfile.ZipFile(archive) as zf: + _safe_extractall(zf, target) + assert (target / "build" / "bin" / "sd-cli").read_bytes() == b"ok" + + +def test_find_sd_cpp_binary_honors_studio_home(tmp_path, monkeypatch): + # A binary installed under a custom Studio root must be discovered without also + # setting UNSLOTH_SD_CPP_PATH (matches default_install_dir's env handling). + from core.inference import sd_cpp_engine as eng + + monkeypatch.delenv("SD_CLI_PATH", raising = False) + monkeypatch.delenv("UNSLOTH_SD_CPP_PATH", raising = False) + studio_home = tmp_path / "studio_root" + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + binary = tmp_path / "stable-diffusion.cpp" / "build" / "bin" / "sd-cli" + binary.parent.mkdir(parents = True) + binary.write_bytes(b"x") + assert eng.find_sd_cpp_binary() == str(binary) diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 3850412369..b3ca9f38a1 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -140,6 +140,54 @@ def _locate_sd_cli(root: Path) -> Optional[Path]: return None +def _download(url: str, dest: Path, *, timeout: float = 300.0) -> None: + """Stream ``url`` to ``dest`` with an explicit timeout. ``urlretrieve`` takes no + timeout and can hang forever on a stalled socket.""" + import shutil + req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-sd-cpp-installer"}) + with urllib.request.urlopen(req, timeout = timeout) as resp, open(dest, "wb") as f: # noqa: S310 + shutil.copyfileobj(resp, f) + + +def _safe_extractall(zf: zipfile.ZipFile, target: Path) -> None: + """``extractall`` with a per-member containment check, so an archive carrying an + absolute path or a ``..`` entry can't write outside ``target`` (Zip-Slip).""" + base = target.resolve() + for member in zf.infolist(): + dest = (base / member.filename).resolve() + if dest != base and base not in dest.parents: + raise RuntimeError(f"unsafe path in archive: {member.filename!r}") + zf.extractall(target) + + +def _maybe_fetch_windows_cudart(release: dict, chosen: str, target: Path) -> None: + """On Windows + a CUDA build, also fetch the separate CUDA-runtime DLL archive. + + Upstream ships the runtime as ``cudart-sd-...-win-cu12-...zip`` (which + ``resolve_release_asset`` filters out); without those DLLs ``sd-cli.exe`` cannot start + on a machine that does not already have the CUDA runtime installed.""" + if platform.system().lower() != "windows" or "cuda" not in chosen.lower(): + return + cudart = next( + ( + a + for a in release.get("assets", []) + if a["name"].lower().startswith("cudart") and "win" in a["name"].lower() + ), + None, + ) + if cudart is None: + return + dest = target / cudart["name"] + print(f"downloading CUDA runtime {cudart['name']} ...", flush = True) + try: + _download(cudart["browser_download_url"], dest) + with zipfile.ZipFile(dest) as zf: + _safe_extractall(zf, target) + finally: + dest.unlink(missing_ok = True) + + def install( *, install_dir: Optional[Path] = None, @@ -170,11 +218,13 @@ def install( target.mkdir(parents = True, exist_ok = True) archive = target / chosen print(f"downloading {chosen} -> {archive}", flush = True) - urllib.request.urlretrieve(url, archive) # noqa: S310 (github release URL) + _download(url, archive) print("extracting ...", flush = True) with zipfile.ZipFile(archive) as zf: - zf.extractall(target) + _safe_extractall(zf, target) archive.unlink(missing_ok = True) + # Windows CUDA builds need the separately-published cudart runtime DLLs. + _maybe_fetch_windows_cudart(release, chosen, target) sd_cli = _locate_sd_cli(target) if not sd_cli: raise RuntimeError(f"archive {chosen} contained no sd-cli binary") @@ -209,7 +259,7 @@ def main(argv: Optional[list[str]] = None) -> int: try: install( - install_dir = Path(args.install_dir) if args.install_dir else None, + install_dir = Path(args.install_dir).expanduser() if args.install_dir else None, accelerator = args.accelerator, ) except RuntimeError as exc: