diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 33cb709ba7..c7f34e39f2 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -4261,7 +4261,10 @@ def ensure_converter_scripts(install_dir: Path, llama_tag: str) -> None: def ensure_diffusion_visual_server( - install_dir: Path, host: HostInfo, release_tag: str | None + install_dir: Path, + host: HostInfo, + release_tag: str | None, + approved_checksums: ApprovedReleaseChecksums, ) -> None: """Best-effort placement of the DiffusionGemma visual-server binary next to llama-server in the install tree, so Studio can serve DiffusionGemma GGUFs @@ -4293,6 +4296,7 @@ def ensure_diffusion_visual_server( try: assets = github_release_assets(DEFAULT_PUBLISHED_REPO, release_tag) match = None + unapproved_matches: list[str] = [] for asset_name, url in assets.items(): low = asset_name.lower() if "llama-diffusion-gemma-visual-server" not in low: @@ -4301,19 +4305,39 @@ def ensure_diffusion_visual_server( continue if (not host.is_windows) and low.endswith(".exe"): continue - match = (asset_name, url) + # This binary is chmod'd executable and later launched by the + # backend, so it must be covered by the approved checksum manifest + # just like every other prebuilt artifact. An asset that matches the + # name but is missing from the manifest is refused rather than run. + approved = approved_checksums.artifacts.get(asset_name) + if approved is None: + unapproved_matches.append(asset_name) + continue + match = (asset_name, url, approved.sha256) break if match is None: - log( - "diffusion visual server not found in the published release; native " - "DiffusionGemma serving needs DG_VISUAL_BIN or a source build" - ) + if unapproved_matches: + log( + "diffusion visual server asset(s) were present but omitted from the " + "approved checksum manifest; refusing unverified native executable: " + + ", ".join(unapproved_matches) + ) + else: + log( + "diffusion visual server not found in the published release; native " + "DiffusionGemma serving needs DG_VISUAL_BIN or a source build" + ) return bin_dir.mkdir(parents = True, exist_ok = True) - download_file(match[1], target) + download_file_verified( + match[1], + target, + expected_sha256 = match[2], + label = f"diffusion visual server {match[0]}", + ) if not host.is_windows: target.chmod(0o755) - log(f"installed diffusion visual server: {match[0]}") + log(f"installed verified diffusion visual server: {match[0]}") except Exception as exc: log( "diffusion visual server fetch skipped " @@ -6637,7 +6661,9 @@ def install_prebuilt( f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" ) try: - ensure_diffusion_visual_server(install_dir, host, plan.release_tag) + ensure_diffusion_visual_server( + install_dir, host, plan.release_tag, plan.approved_checksums + ) except Exception as exc: log( "diffusion visual server step skipped; install remains valid " diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index a5a1131ecc..9ee8759bb4 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -37,6 +37,41 @@ install_prebuilt = INSTALL_LLAMA_PREBUILT.install_prebuilt write_prebuilt_metadata = INSTALL_LLAMA_PREBUILT.write_prebuilt_metadata existing_install_matches_plan = INSTALL_LLAMA_PREBUILT.existing_install_matches_plan existing_install_matches_choice = INSTALL_LLAMA_PREBUILT.existing_install_matches_choice +ensure_diffusion_visual_server = INSTALL_LLAMA_PREBUILT.ensure_diffusion_visual_server + + +def linux_host() -> HostInfo: + return HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + + +def approved_release_checksums_for_asset(asset_name: str, sha256: str) -> ApprovedReleaseChecksums: + return ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "b9334", + upstream_tag = "b9334", + artifacts = { + asset_name: ApprovedArtifactHash( + asset_name = asset_name, + sha256 = sha256, + repo = "unslothai/llama.cpp", + kind = "diffusion-visual-server", + ) + }, + ) def approved_checksums_for( @@ -2828,3 +2863,87 @@ def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled(tmp monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", True) calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32) assert calls == {"quantize": 1, "server": 1} + + +def test_diffusion_visual_server_uses_approved_checksum_download(monkeypatch, tmp_path: Path): + asset_name = "llama-diffusion-gemma-visual-server-linux-x64" + expected_sha = "a" * 64 + asset_url = "https://github.com/unslothai/llama.cpp/releases/download/b9334/" + asset_name + calls: list[tuple[str, Path, str | None, str | None]] = [] + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: {asset_name: asset_url}, + ) + + def fake_download_file(url, destination): + raise AssertionError("diffusion visual server must not use unverified download_file") + + def fake_download_file_verified(url, destination, *, expected_sha256, label): + calls.append((url, Path(destination), expected_sha256, label)) + Path(destination).write_bytes(b"verified visual server") + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "download_file_verified", fake_download_file_verified + ) + + ensure_diffusion_visual_server( + tmp_path / "install", + linux_host(), + "b9334", + approved_release_checksums_for_asset(asset_name, expected_sha), + ) + + target = tmp_path / "install" / "build" / "bin" / "llama-diffusion-gemma-visual-server" + assert calls == [ + ( + asset_url, + target, + expected_sha, + f"diffusion visual server {asset_name}", + ) + ] + assert target.read_bytes() == b"verified visual server" + assert target.stat().st_mode & 0o777 == 0o755 + + +def test_diffusion_visual_server_refuses_unapproved_release_asset(monkeypatch, tmp_path: Path): + asset_name = "llama-diffusion-gemma-visual-server-attacker-linux" + verified_calls: list[str] = [] + raw_calls: list[str] = [] + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: {asset_name: "https://example.test/" + asset_name}, + ) + + def fake_download_file(url, destination): + raw_calls.append(url) + + def fake_download_file_verified(url, destination, *, expected_sha256, label): + verified_calls.append(url) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "download_file_verified", fake_download_file_verified + ) + + ensure_diffusion_visual_server( + tmp_path / "install", + linux_host(), + "b9334", + ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "b9334", + upstream_tag = "b9334", + artifacts = {}, + ), + ) + + target = tmp_path / "install" / "build" / "bin" / "llama-diffusion-gemma-visual-server" + assert not target.exists() + assert raw_calls == [] + assert verified_calls == []