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

@@ -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 "

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 == []