From bd2438ea6580c7b12ea396be2f4975ca89af928e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 24 Jun 2026 06:37:41 -0700 Subject: [PATCH] Verify DiffusionGemma visual-server binary against approved checksums (#6635) ensure_diffusion_visual_server() downloaded the visual-server release asset with the unverified download_file() and marked it executable, bypassing the approved-checksum manifest that gates every other prebuilt llama.cpp artifact. The backend later auto-discovers that binary and launches it through DG_VISUAL_BIN, so a compromised or substituted release asset could place attacker-controlled native code in the install tree and have it executed under the Studio user. Require the matched asset to be present in the approved checksum manifest and download it through download_file_verified() with the published sha256. A name-matching asset that is absent from the manifest is refused rather than executed. Add regression tests covering the verified-download path and the refusal of an unapproved asset. --- studio/install_llama_prebuilt.py | 44 +++++-- .../test_install_llama_prebuilt_logic.py | 119 ++++++++++++++++++ 2 files changed, 154 insertions(+), 9 deletions(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 33cb709ba..c7f34e39f 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 a5a1131ec..9ee8759bb 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 == []