From 6176fa73bb23c665916594f7751eb5943b4b3c53 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 17:40:49 -0700 Subject: [PATCH 1/7] feat: add cross-platform hardware block to comfy env --json (BE-3399) Add comfy_cli/hardware.py with a single detect_hardware() entry point that reports OS, CPU, RAM, and a GPU sub-block (vendor/model/VRAM/unified-memory) so agent surfaces can route weak machines away from local diffusion. - macOS: cpu via sysctl machdep.cpu.brand_string; Apple Silicon -> apple unified-memory GPU; Intel Mac -> unknown non-unified GPU (no system_profiler). - NVIDIA (any OS): nvidia-smi CSV, then ctypes libcuda fallback (reuses cuda_detect._load_libcuda). - AMD (Linux): best-effort rocm-smi --json. - Never raises, never blocks long: every probe wrapped, subprocesses timeout=5. Wire hardware into EnvChecker.fill_data() (JSON) and a Hardware row in fill_print_table() (pretty). Extend schemas/env.json with an OPTIONAL, fully nullable hardware property (not in required) so older/failed-probe payloads stay valid. Envelope schema unchanged (envelope/1). --- comfy_cli/env_checker.py | 34 ++++ comfy_cli/hardware.py | 279 +++++++++++++++++++++++++++++++ comfy_cli/schemas/env.json | 21 +++ tests/comfy_cli/test_hardware.py | 263 +++++++++++++++++++++++++++++ 4 files changed, 597 insertions(+) create mode 100644 comfy_cli/hardware.py create mode 100644 tests/comfy_cli/test_hardware.py diff --git a/comfy_cli/env_checker.py b/comfy_cli/env_checker.py index 2aea4cbd..60cfae45 100644 --- a/comfy_cli/env_checker.py +++ b/comfy_cli/env_checker.py @@ -9,6 +9,7 @@ from rich.console import Console from comfy_cli.config_manager import ConfigManager +from comfy_cli.hardware import detect_hardware from comfy_cli.utils import singleton console = Console() @@ -32,6 +33,36 @@ def format_python_version(version_info): return f"[bold red]{version_info.major}.{version_info.minor}.{version_info.micro}[/bold red]" +def _bytes_to_gb(value) -> str: + """Render a byte count as a whole-number GB string, or ``?`` if unknown.""" + if not isinstance(value, (int | float)): + return "?" + return f"{round(value / (1024**3))} GB" + + +def format_hardware_summary(hw: dict) -> str: + """One-line ``cpu / RAM GB / GPU model (VRAM GB or 'unified')`` summary. + + Used by ``fill_print_table`` (pretty mode). Tolerates missing/None fields so + it never raises on a partial (failed-probe) hardware block. + """ + cpu = hw.get("cpu") or "unknown CPU" + ram = _bytes_to_gb(hw.get("ram_bytes")) + + gpu = hw.get("gpu") + if not gpu: + gpu_part = "no GPU" + else: + model = gpu.get("model") or gpu.get("vendor") or "unknown GPU" + if gpu.get("unified_memory"): + gpu_part = f"{model} (unified)" + elif gpu.get("vram_bytes") is not None: + gpu_part = f"{model} ({_bytes_to_gb(gpu.get('vram_bytes'))} VRAM)" + else: + gpu_part = model + return f"{cpu} / {ram} RAM / {gpu_part}" + + def check_comfy_server_running(port=8188, host="localhost", timeout: float = 5.0): """ Checks if the Comfy server is running by making a GET request to the /history endpoint. @@ -106,6 +137,8 @@ def fill_print_table(self): config_data = ConfigManager().get_env_data() data.extend(config_data) + data.append(("Hardware", format_hardware_summary(detect_hardware()))) + if check_comfy_server_running(): data.append( ( @@ -139,4 +172,5 @@ def fill_data(self) -> dict: "running": server_running, "url": "http://localhost:8188" if server_running else None, }, + "hardware": detect_hardware(), } diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py new file mode 100644 index 00000000..8a4d4b81 --- /dev/null +++ b/comfy_cli/hardware.py @@ -0,0 +1,279 @@ +"""Cross-platform hardware probing for ``comfy env --json``. + +A single entry point, :func:`detect_hardware`, returns a JSON-serializable dict +describing the machine's generation-relevant hardware (OS, CPU, RAM, GPU). Agent +surfaces use this to route weak machines away from local diffusion. + +Contract: **never raise, never block long.** Every probe is wrapped so a failure +yields ``None`` rather than an exception, and every subprocess is bounded by a +``timeout=5``. The RAM figure is the only value expected to always populate +(``psutil`` is a hard dependency); everything else degrades to ``None``. +""" + +from __future__ import annotations + +import ctypes +import logging +import platform +import subprocess + +import psutil + +from comfy_cli import cuda_detect + +logger = logging.getLogger(__name__) + +_SUBPROCESS_TIMEOUT = 5 + + +def _run(cmd: list[str]) -> str | None: + """Run ``cmd`` and return stripped stdout, or ``None`` on any failure. + + Bounded by ``timeout=5`` so a hung binary can never block the probe. + """ + try: + output = subprocess.check_output( + cmd, + text=True, + timeout=_SUBPROCESS_TIMEOUT, + stderr=subprocess.DEVNULL, + ) + except Exception: + logger.debug("hardware probe command failed: %s", cmd, exc_info=True) + return None + output = output.strip() + return output or None + + +def _detect_cpu(system: str) -> str | None: + """Best-effort human-readable CPU/chip name, or ``None`` if unknown.""" + try: + if system == "Darwin": + return _run(["sysctl", "-n", "machdep.cpu.brand_string"]) + + if system == "Linux": + try: + with open("/proc/cpuinfo", encoding="utf-8", errors="replace") as fh: + for line in fh: + if line.startswith("model name"): + _, _, value = line.partition(":") + value = value.strip() + if value: + return value + except OSError: + logger.debug("reading /proc/cpuinfo failed", exc_info=True) + + # Windows / Linux fallback: platform.processor() is often populated on + # Windows and empty on Linux (hence the /proc/cpuinfo pass above first). + return platform.processor() or None + except Exception: + logger.debug("CPU detection failed", exc_info=True) + return None + + +def _detect_gpu_nvidia_smi() -> dict | None: + """Query ``nvidia-smi`` for the first GPU's name and VRAM. + + Returns a gpu dict or ``None`` if the binary is absent / produced no usable + output. ``cuda_detect`` already parses nvidia-smi in production. + """ + output = _run( + [ + "nvidia-smi", + "--query-gpu=name,memory.total", + "--format=csv,noheader,nounits", + ] + ) + if not output: + return None + + # First line is the first GPU: "NVIDIA GeForce RTX 4090, 24564" + first = output.splitlines()[0] + parts = [p.strip() for p in first.split(",")] + if len(parts) < 2: + return None + + model = parts[0] or None + vram_bytes = None + try: + mib = int(parts[1]) + vram_bytes = mib * 1024 * 1024 + except (ValueError, TypeError): + logger.debug("could not parse nvidia-smi VRAM: %r", parts[1]) + + if not model and vram_bytes is None: + return None + + return { + "vendor": "nvidia", + "model": model, + "vram_bytes": vram_bytes, + "unified_memory": False, + } + + +def _detect_gpu_nvidia_ctypes() -> dict | None: + """ctypes fallback for NVIDIA when ``nvidia-smi`` is unavailable. + + Uses the CUDA driver API via ``cuda_detect._load_libcuda``: ``cuInit`` + + ``cuDeviceGet`` + ``cuDeviceGetName`` + ``cuDeviceTotalMem_v2`` — model and + VRAM without the binary. Returns a gpu dict or ``None``. + """ + try: + libcuda = cuda_detect._load_libcuda() + except OSError: + logger.debug("libcuda not loadable for GPU probe") + return None + + try: + if libcuda.cuInit(0) != 0: + return None + + device = ctypes.c_int() + if libcuda.cuDeviceGet(ctypes.byref(device), 0) != 0: + return None + + name_buf = ctypes.create_string_buffer(256) + model = None + if libcuda.cuDeviceGetName(name_buf, len(name_buf), device) == 0: + decoded = name_buf.value.decode("utf-8", errors="replace").strip() + model = decoded or None + + vram_bytes = None + total = ctypes.c_size_t() + total_mem = getattr(libcuda, "cuDeviceTotalMem_v2", None) or getattr(libcuda, "cuDeviceTotalMem", None) + if total_mem is not None and total_mem(ctypes.byref(total), device) == 0: + vram_bytes = int(total.value) + + if not model and vram_bytes is None: + return None + + return { + "vendor": "nvidia", + "model": model, + "vram_bytes": vram_bytes, + "unified_memory": False, + } + except Exception: + logger.debug("ctypes NVIDIA GPU probe failed", exc_info=True) + return None + + +def _detect_gpu_amd() -> dict | None: + """Best-effort AMD GPU probe on Linux via ``rocm-smi``. + + Minimal by design — nulls are acceptable. Returns a gpu dict (vendor "amd") + or ``None`` if ``rocm-smi`` is absent / unparseable. + """ + import json + + output = _run(["rocm-smi", "--showmeminfo", "vram", "--json"]) + if not output: + return None + + try: + payload = json.loads(output) + except (ValueError, TypeError): + logger.debug("could not parse rocm-smi JSON", exc_info=True) + return None + + if not isinstance(payload, dict) or not payload: + return None + + # rocm-smi keys cards as "card0", "card1", ...; take the first. + model = None + vram_bytes = None + for card in payload.values(): + if not isinstance(card, dict): + continue + for key, value in card.items(): + lowered = key.lower() + if model is None and "name" in lowered: + model = str(value).strip() or None + if vram_bytes is None and "total" in lowered and "vram" in lowered: + try: + vram_bytes = int(value) + except (ValueError, TypeError): + pass + break + + return { + "vendor": "amd", + "model": model, + "vram_bytes": vram_bytes, + "unified_memory": False, + } + + +def _detect_gpu(system: str, machine: str, cpu: str | None) -> dict | None: + """Resolve the GPU block, or ``None`` if no GPU is detected. + + Apple Silicon is unified-memory (the RAM figure IS the GPU budget). Every + other path tries NVIDIA (nvidia-smi then ctypes), then AMD. + """ + try: + if system == "Darwin" and machine == "arm64": + return { + "vendor": "apple", + "model": cpu, + "vram_bytes": None, + "unified_memory": True, + } + + gpu = _detect_gpu_nvidia_smi() + if gpu is not None: + return gpu + + gpu = _detect_gpu_nvidia_ctypes() + if gpu is not None: + return gpu + + if system == "Linux": + gpu = _detect_gpu_amd() + if gpu is not None: + return gpu + + return None + except Exception: + logger.debug("GPU detection failed", exc_info=True) + return None + + +def _detect_ram_bytes() -> int | None: + try: + return int(psutil.virtual_memory().total) + except Exception: + logger.debug("RAM detection failed", exc_info=True) + return None + + +def detect_hardware() -> dict: + """Detect generation-relevant hardware. Never raises. + + Returns a dict shaped like ``schemas/env.json``'s ``hardware`` block: + ``os``/``os_version``/``arch``/``cpu``/``ram_bytes`` plus a ``gpu`` sub-dict + (or ``None``). Any probe that fails contributes ``None`` rather than raising. + """ + try: + system = platform.system() + except Exception: + system = "" + try: + os_version = platform.release() + except Exception: + os_version = None + try: + machine = platform.machine() + except Exception: + machine = "" + + cpu = _detect_cpu(system) + + return { + "os": system or None, + "os_version": os_version or None, + "arch": machine or None, + "cpu": cpu, + "ram_bytes": _detect_ram_bytes(), + "gpu": _detect_gpu(system, machine, cpu), + } diff --git a/comfy_cli/schemas/env.json b/comfy_cli/schemas/env.json index 46afcccd..29b0eb72 100644 --- a/comfy_cli/schemas/env.json +++ b/comfy_cli/schemas/env.json @@ -51,6 +51,27 @@ "port": {"type": "integer"}, "pid": {"type": "integer"} } + }, + "hardware": { + "type": ["object", "null"], + "additionalProperties": true, + "properties": { + "os": {"type": ["string", "null"]}, + "os_version": {"type": ["string", "null"]}, + "arch": {"type": ["string", "null"]}, + "cpu": {"type": ["string", "null"]}, + "ram_bytes": {"type": ["integer", "null"]}, + "gpu": { + "type": ["object", "null"], + "additionalProperties": true, + "properties": { + "vendor": {"type": ["string", "null"]}, + "model": {"type": ["string", "null"]}, + "vram_bytes": {"type": ["integer", "null"]}, + "unified_memory": {"type": ["boolean", "null"]} + } + } + } } } } diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py new file mode 100644 index 00000000..8057f13e --- /dev/null +++ b/tests/comfy_cli/test_hardware.py @@ -0,0 +1,263 @@ +"""Tests for :mod:`comfy_cli.hardware`. + +Each branch (macOS/apple, nvidia-smi, ctypes fallback, total failure) is mocked +at the module boundary so the tests run identically on any CI runner. The +overriding contract under test: ``detect_hardware`` never raises, always returns +RAM, and degrades every failed probe to ``None``. +""" + +from __future__ import annotations + +import ctypes +import json +from pathlib import Path +from unittest.mock import patch + +import jsonschema + +from comfy_cli import hardware +from comfy_cli.env_checker import EnvChecker, format_hardware_summary + +_EnvCheckerCls = EnvChecker.__closure__[0].cell_contents + +SCHEMA_PATH = Path(hardware.__file__).parent / "schemas" / "env.json" + + +def _env_schema() -> dict: + return json.loads(SCHEMA_PATH.read_text()) + + +class TestDetectHardwareMacOS: + def test_apple_silicon_unified_block(self): + with ( + patch.object(hardware.platform, "system", return_value="Darwin"), + patch.object(hardware.platform, "machine", return_value="arm64"), + patch.object(hardware.platform, "release", return_value="25.4.0"), + patch.object(hardware, "_run", return_value="Apple M4 Max"), + patch.object(hardware, "_detect_ram_bytes", return_value=68719476736), + ): + hw = hardware.detect_hardware() + + assert hw["os"] == "Darwin" + assert hw["os_version"] == "25.4.0" + assert hw["arch"] == "arm64" + assert hw["cpu"] == "Apple M4 Max" + assert hw["ram_bytes"] == 68719476736 + assert hw["gpu"] == { + "vendor": "apple", + "model": "Apple M4 Max", + "vram_bytes": None, + "unified_memory": True, + } + + def test_intel_mac_has_no_apple_gpu(self): + """Intel Macs are not unified-memory; they fall through to the NVIDIA/AMD + probes (which return nothing here) → gpu null, not an apple block.""" + with ( + patch.object(hardware.platform, "system", return_value="Darwin"), + patch.object(hardware.platform, "machine", return_value="x86_64"), + patch.object(hardware, "_run", return_value=None), + patch.object(hardware.cuda_detect, "_load_libcuda", side_effect=OSError), + patch.object(hardware, "_detect_ram_bytes", return_value=17179869184), + ): + hw = hardware.detect_hardware() + + assert hw["gpu"] is None + assert hw["ram_bytes"] == 17179869184 + + +class TestDetectHardwareNvidiaSmi: + def test_nvidia_smi_happy_path(self): + def fake_run(cmd): + if cmd[0] == "nvidia-smi": + return "NVIDIA GeForce RTX 4090, 24564" + return None + + with ( + patch.object(hardware.platform, "system", return_value="Linux"), + patch.object(hardware.platform, "machine", return_value="x86_64"), + patch.object(hardware, "_run", side_effect=fake_run), + patch.object(hardware, "_detect_ram_bytes", return_value=137438953472), + ): + hw = hardware.detect_hardware() + + assert hw["gpu"] == { + "vendor": "nvidia", + "model": "NVIDIA GeForce RTX 4090", + "vram_bytes": 24564 * 1024 * 1024, + "unified_memory": False, + } + + def test_nvidia_smi_multi_gpu_takes_first(self): + with patch.object( + hardware, + "_run", + return_value="NVIDIA GeForce RTX 4090, 24564\nNVIDIA GeForce RTX 3090, 24576", + ): + gpu = hardware._detect_gpu_nvidia_smi() + assert gpu["model"] == "NVIDIA GeForce RTX 4090" + assert gpu["vram_bytes"] == 24564 * 1024 * 1024 + + +class _FakeLibCuda: + """Minimal stand-in for a loaded ``libcuda`` CDLL for the ctypes path.""" + + def __init__(self, name: str, total_bytes: int): + self._name = name + self._total = total_bytes + + def cuInit(self, flags): + return 0 + + def cuDeviceGet(self, dev_ptr, ordinal): + dev_ptr._obj.value = 0 + return 0 + + def cuDeviceGetName(self, buf, length, dev): + buf.value = self._name.encode("utf-8") + return 0 + + def cuDeviceTotalMem_v2(self, size_ptr, dev): + size_ptr._obj.value = self._total + return 0 + + +class TestDetectHardwareCtypesFallback: + def test_nvidia_smi_absent_libcuda_used(self): + fake = _FakeLibCuda("NVIDIA A100-SXM4-80GB", 85899345920) + + def fake_run(cmd): + # nvidia-smi absent → None so we exercise the ctypes fallback. + return None + + with ( + patch.object(hardware.platform, "system", return_value="Linux"), + patch.object(hardware.platform, "machine", return_value="x86_64"), + patch.object(hardware, "_run", side_effect=fake_run), + patch.object(hardware.cuda_detect, "_load_libcuda", return_value=fake), + patch.object(hardware, "_detect_ram_bytes", return_value=270582939648), + ): + hw = hardware.detect_hardware() + + assert hw["gpu"] == { + "vendor": "nvidia", + "model": "NVIDIA A100-SXM4-80GB", + "vram_bytes": 85899345920, + "unified_memory": False, + } + + def test_byref_obj_contract(self): + """Guard the ctypes mock trick: byref exposes the wrapped object so the + fake driver can write results back through the pointer.""" + val = ctypes.c_int(7) + assert ctypes.byref(val)._obj is val + + +class TestDetectHardwareNeverRaises: + def test_all_probes_fail_yields_nulls_but_ram(self): + with ( + patch.object(hardware.platform, "system", return_value="Windows"), + patch.object(hardware.platform, "machine", return_value="AMD64"), + patch.object(hardware.platform, "processor", return_value=""), + patch.object(hardware, "_run", side_effect=TimeoutError("boom")), + patch.object(hardware.cuda_detect, "_load_libcuda", side_effect=OSError), + ): + # Must not raise even though every subprocess/ctypes probe blows up. + hw = hardware.detect_hardware() + + assert hw["os"] == "Windows" + assert hw["cpu"] is None + assert hw["gpu"] is None + # RAM comes from psutil (a hard dep) and is unaffected by probe failures. + assert isinstance(hw["ram_bytes"], int) + assert hw["ram_bytes"] > 0 + + def test_platform_calls_raising_do_not_escape(self): + with ( + patch.object(hardware.platform, "system", side_effect=RuntimeError), + patch.object(hardware.platform, "machine", side_effect=RuntimeError), + patch.object(hardware.platform, "release", side_effect=RuntimeError), + ): + hw = hardware.detect_hardware() + assert hw["os"] is None + assert hw["arch"] is None + assert hw["gpu"] is None + + +class TestFormatHardwareSummary: + def test_unified_memory_summary(self): + hw = { + "cpu": "Apple M4 Max", + "ram_bytes": 68719476736, + "gpu": {"vendor": "apple", "model": "Apple M4 Max", "vram_bytes": None, "unified_memory": True}, + } + summary = format_hardware_summary(hw) + assert "Apple M4 Max" in summary + assert "64 GB RAM" in summary + assert "unified" in summary + + def test_discrete_vram_summary(self): + hw = { + "cpu": "AMD Ryzen 9", + "ram_bytes": 137438953472, + "gpu": {"vendor": "nvidia", "model": "RTX 4090", "vram_bytes": 24 * 1024**3, "unified_memory": False}, + } + summary = format_hardware_summary(hw) + assert "RTX 4090" in summary + assert "24 GB VRAM" in summary + + def test_no_gpu_and_missing_fields(self): + summary = format_hardware_summary({"cpu": None, "ram_bytes": None, "gpu": None}) + assert "no GPU" in summary + assert "unknown CPU" in summary + + +class TestFillDataHardware: + @patch("comfy_cli.env_checker.check_comfy_server_running", return_value=False) + @patch("comfy_cli.env_checker.ConfigManager") + def test_fill_data_includes_hardware_and_validates(self, mock_cm, _mock_server): + mock_cm.return_value.get_data.return_value = {"path": "/tmp/config"} + + checker = _EnvCheckerCls.__new__(_EnvCheckerCls) + checker.virtualenv_path = None + checker.conda_env = None + + data = checker.fill_data() + assert "hardware" in data + hw = data["hardware"] + assert set(hw) >= {"os", "os_version", "arch", "cpu", "ram_bytes", "gpu"} + + # Complete the payload the way cmdline.py does (workspace is added there) + # and validate the whole thing against env.json. + data["workspace"] = {"path": None, "type": None} + jsonschema.Draft202012Validator(_env_schema()).validate(data) + + def test_fully_nulled_hardware_still_validates(self): + """A failed-probe payload (all nulls, gpu null) must stay schema-valid — + this is why `hardware` is optional and every leaf is nullable.""" + payload = { + "python": {"version": "3.12.1", "executable": "/usr/bin/python"}, + "config": {}, + "workspace": {"path": None, "type": None}, + "server": {"running": False}, + "hardware": { + "os": None, + "os_version": None, + "arch": None, + "cpu": None, + "ram_bytes": None, + "gpu": None, + }, + } + jsonschema.Draft202012Validator(_env_schema()).validate(payload) + + def test_payload_without_hardware_still_validates(self): + """Older consumers / pre-hardware payloads must remain valid — `hardware` + is not in `required`.""" + payload = { + "python": {"version": "3.12.1", "executable": "/usr/bin/python"}, + "config": {}, + "workspace": {"path": None, "type": None}, + "server": {"running": False}, + } + jsonschema.Draft202012Validator(_env_schema()).validate(payload) From 5e5b4f4ff7c918331bd1952369fe2a5329fd74b4 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 18:01:05 -0700 Subject: [PATCH 2/7] fix(hardware): address Cursor review findings for env hardware block (BE-3399) - env_checker: escape untrusted cpu/gpu-model strings in the Rich-rendered hardware summary so markup-like content (e.g. "[/]") can't raise MarkupError and crash `comfy env` in pretty mode. - hardware(_detect_gpu_amd): return None (no phantom all-None AMD block) when nothing parses, matching the NVIDIA probes; gate the card loop on the "card" key prefix so a non-card metadata block isn't parsed as the GPU; exclude the "VRAM Total Used Memory" usage key so total capacity isn't understated. - hardware(_detect_gpu_nvidia_ctypes): pop/restore CUDA_VISIBLE_DEVICES around the ctypes probe (mirrors cuda_detect) so an exported ""/"-1" doesn't hide a present GPU. - tests: cover AMD total-vs-used key, metadata-block skip, all-None->None, and markup escaping. Co-Authored-By: Claude Opus 4.8 --- comfy_cli/env_checker.py | 9 +++- comfy_cli/hardware.py | 34 ++++++++++++--- tests/comfy_cli/test_hardware.py | 72 ++++++++++++++++++++++++++++++++ 3 files changed, 107 insertions(+), 8 deletions(-) diff --git a/comfy_cli/env_checker.py b/comfy_cli/env_checker.py index 60cfae45..c1ad8185 100644 --- a/comfy_cli/env_checker.py +++ b/comfy_cli/env_checker.py @@ -7,6 +7,7 @@ import requests from rich.console import Console +from rich.markup import escape from comfy_cli.config_manager import ConfigManager from comfy_cli.hardware import detect_hardware @@ -46,14 +47,18 @@ def format_hardware_summary(hw: dict) -> str: Used by ``fill_print_table`` (pretty mode). Tolerates missing/None fields so it never raises on a partial (failed-probe) hardware block. """ - cpu = hw.get("cpu") or "unknown CPU" + # cpu/model come from untrusted probe output (/proc/cpuinfo, nvidia-smi, + # rocm-smi, libcuda). This summary is rendered by fill_print_table via Rich + # (markup enabled), so escape those strings — a stray "[" (plausible on + # VMs/hypervisors) would otherwise raise MarkupError and crash `comfy env`. + cpu = escape(hw.get("cpu") or "unknown CPU") ram = _bytes_to_gb(hw.get("ram_bytes")) gpu = hw.get("gpu") if not gpu: gpu_part = "no GPU" else: - model = gpu.get("model") or gpu.get("vendor") or "unknown GPU" + model = escape(str(gpu.get("model") or gpu.get("vendor") or "unknown GPU")) if gpu.get("unified_memory"): gpu_part = f"{model} (unified)" elif gpu.get("vram_bytes") is not None: diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py index 8a4d4b81..61aba2b3 100644 --- a/comfy_cli/hardware.py +++ b/comfy_cli/hardware.py @@ -14,6 +14,7 @@ import ctypes import logging +import os import platform import subprocess @@ -125,7 +126,15 @@ def _detect_gpu_nvidia_ctypes() -> dict | None: logger.debug("libcuda not loadable for GPU probe") return None + # Mirror cuda_detect.detect_cuda_driver_version: an exported + # CUDA_VISIBLE_DEVICES="" or "-1" (common on shared/CI hosts) hides all + # devices from cuDeviceGet, so a physically-present GPU would report as + # gpu:null. Drop it for the duration of the probe and restore it after. + saved_visible = os.environ.get("CUDA_VISIBLE_DEVICES") try: + if saved_visible is not None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + if libcuda.cuInit(0) != 0: return None @@ -157,6 +166,9 @@ def _detect_gpu_nvidia_ctypes() -> dict | None: except Exception: logger.debug("ctypes NVIDIA GPU probe failed", exc_info=True) return None + finally: + if saved_visible is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = saved_visible def _detect_gpu_amd() -> dict | None: @@ -180,23 +192,33 @@ def _detect_gpu_amd() -> dict | None: if not isinstance(payload, dict) or not payload: return None - # rocm-smi keys cards as "card0", "card1", ...; take the first. + # rocm-smi keys cards as "card0", "card1", ...; take the first card entry. + # Gate on the "card" prefix so a non-card metadata block (e.g. a "system" + # dict) that iterates first isn't mistaken for the GPU. model = None vram_bytes = None - for card in payload.values(): - if not isinstance(card, dict): + for key, card in payload.items(): + if not str(key).lower().startswith("card") or not isinstance(card, dict): continue - for key, value in card.items(): - lowered = key.lower() + for field, value in card.items(): + lowered = field.lower() if model is None and "name" in lowered: model = str(value).strip() or None - if vram_bytes is None and "total" in lowered and "vram" in lowered: + # Match the total-capacity key ("VRAM Total Memory (B)"), excluding + # the usage key ("VRAM Total Used Memory (B)") which also contains + # both "vram" and "total" and would otherwise understate capacity. + if vram_bytes is None and "vram" in lowered and "total" in lowered and "used" not in lowered: try: vram_bytes = int(value) except (ValueError, TypeError): pass break + # Report no GPU rather than a phantom all-None AMD block (which would spoof + # GPU presence for routing decisions), matching the NVIDIA probes. + if not model and vram_bytes is None: + return None + return { "vendor": "amd", "model": model, diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py index 8057f13e..c496f990 100644 --- a/tests/comfy_cli/test_hardware.py +++ b/tests/comfy_cli/test_hardware.py @@ -184,7 +184,79 @@ def test_platform_calls_raising_do_not_escape(self): assert hw["gpu"] is None +class TestDetectGpuAmd: + def test_amd_happy_path(self): + payload = json.dumps( + {"card0": {"Card SKU": "gfx90a", "GPU Name": "AMD Instinct MI210", "VRAM Total Memory (B)": "68702699520"}} + ) + with patch.object(hardware, "_run", return_value=payload): + gpu = hardware._detect_gpu_amd() + assert gpu == { + "vendor": "amd", + "model": "AMD Instinct MI210", + "vram_bytes": 68702699520, + "unified_memory": False, + } + + def test_amd_total_key_beats_used_key(self): + """The 'VRAM Total Used Memory' key also contains 'vram'+'total'; the probe + must report the total-capacity key, not usage.""" + payload = json.dumps( + { + "card0": { + "VRAM Total Used Memory (B)": "1073741824", + "VRAM Total Memory (B)": "68702699520", + } + } + ) + with patch.object(hardware, "_run", return_value=payload): + gpu = hardware._detect_gpu_amd() + assert gpu["vram_bytes"] == 68702699520 + + def test_amd_skips_non_card_metadata_block(self): + """A leading non-card metadata dict must not be mistaken for the GPU.""" + payload = json.dumps( + { + "system": {"Driver version": "6.0.0"}, + "card0": {"GPU Name": "AMD Radeon RX 7900 XTX", "VRAM Total Memory (B)": "25757220864"}, + } + ) + with patch.object(hardware, "_run", return_value=payload): + gpu = hardware._detect_gpu_amd() + assert gpu["model"] == "AMD Radeon RX 7900 XTX" + assert gpu["vram_bytes"] == 25757220864 + + def test_amd_all_none_reports_no_gpu(self): + """An error/metadata-only payload with nothing parseable must yield None, + not a phantom {vendor: amd, model: None, vram_bytes: None} block.""" + payload = json.dumps({"card0": {"Something Unrelated": "x"}}) + with patch.object(hardware, "_run", return_value=payload): + assert hardware._detect_gpu_amd() is None + + class TestFormatHardwareSummary: + def test_markup_in_model_is_escaped_not_crashing(self): + """CPU/GPU strings are untrusted; markup-like content (e.g. a close tag + '[/]') must be escaped so rendering through Rich can't raise MarkupError.""" + import pytest + from rich.errors import MarkupError + from rich.text import Text + + raw_model = "GPU [/] X" + # Sanity: the raw model really would crash Rich markup parsing. + with pytest.raises(MarkupError): + Text.from_markup(raw_model) + + hw = { + "cpu": "Fake CPU", + "ram_bytes": 17179869184, + "gpu": {"vendor": "nvidia", "model": raw_model, "vram_bytes": 8 * 1024**3, "unified_memory": False}, + } + summary = format_hardware_summary(hw) + assert "\\[/]" in summary + # Rendering with markup enabled (as cmdline.py does) must not raise. + Text.from_markup(summary) + def test_unified_memory_summary(self): hw = { "cpu": "Apple M4 Max", From 2e44669435f1dddef864255a893628707d80098d Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 18:54:42 -0700 Subject: [PATCH 3/7] fix(hardware): resolve probe binaries via absolute path to block Windows CWD planting (BE-3434) --- comfy_cli/hardware.py | 55 +++++++++++++++++++++- tests/comfy_cli/test_hardware.py | 80 ++++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 2 deletions(-) diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py index 61aba2b3..41537287 100644 --- a/comfy_cli/hardware.py +++ b/comfy_cli/hardware.py @@ -16,6 +16,7 @@ import logging import os import platform +import shutil import subprocess import psutil @@ -27,14 +28,64 @@ _SUBPROCESS_TIMEOUT = 5 +def _is_within_cwd(path: str) -> bool: + """Return ``True`` only if ``path`` provably resolves inside ``os.getcwd()``. + + Used to reject a probe binary that resolves to the current working directory + (see :func:`_resolve_binary`). Ambiguity — a path on a different drive, or an + unresolvable one — is treated as *not* under the CWD so a legitimate system + binary is never rejected. + """ + try: + cwd = os.path.realpath(os.getcwd()) + resolved = os.path.realpath(path) + return os.path.commonpath([cwd, resolved]) == cwd + except (OSError, ValueError): + # Different drives (Windows) or an unresolvable path → not under the CWD. + return False + + +def _resolve_binary(name: str) -> str | None: + """Resolve a probe binary to a trusted absolute path, or ``None`` to skip it. + + :func:`shutil.which` performs a PATH lookup and returns ``None`` when the + binary is absent (so the probe simply degrades to ``None``). Passing the + resolved absolute path to :func:`subprocess.check_output` — rather than the + bare name — prevents Windows ``CreateProcess`` from searching the current + working directory, so running ``comfy env`` from an attacker-controlled + directory cannot execute a planted ``nvidia-smi.exe``. + + ``shutil.which`` on Windows may itself resolve against the current directory, + so as defense-in-depth a resolved path inside the CWD tree is additionally + rejected. A legitimate system binary (e.g. ``nvidia-smi.exe`` under + ``System32``) is unaffected. + """ + try: + path = shutil.which(name) + if path is None: + return None + if platform.system() == "Windows" and _is_within_cwd(path): + logger.debug("skipping hardware probe %r: resolved under CWD (%s)", name, path) + return None + return path + except Exception: + logger.debug("resolving hardware probe binary %r failed", name, exc_info=True) + return None + + def _run(cmd: list[str]) -> str | None: """Run ``cmd`` and return stripped stdout, or ``None`` on any failure. - Bounded by ``timeout=5`` so a hung binary can never block the probe. + ``cmd[0]`` is resolved to a trusted absolute path via :func:`_resolve_binary` + before execution (skipping the probe when absent or CWD-planted), and the run + is bounded by ``timeout=5`` so a hung binary can never block the probe. """ + resolved = _resolve_binary(cmd[0]) + if resolved is None: + return None try: output = subprocess.check_output( - cmd, + [resolved, *cmd[1:]], text=True, timeout=_SUBPROCESS_TIMEOUT, stderr=subprocess.DEVNULL, diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py index c496f990..0b2ded75 100644 --- a/tests/comfy_cli/test_hardware.py +++ b/tests/comfy_cli/test_hardware.py @@ -333,3 +333,83 @@ def test_payload_without_hardware_still_validates(self): "server": {"running": False}, } jsonschema.Draft202012Validator(_env_schema()).validate(payload) + + +class TestRunResolvesBinaryPath: + """``_run`` must resolve the binary to an absolute path (never invoke by bare + name) so Windows ``CreateProcess`` can't pick up a CWD-planted executable.""" + + def test_run_invokes_resolved_absolute_path(self): + captured = {} + + def fake_check_output(cmd, **kwargs): + captured["cmd"] = cmd + return "ok\n" + + with ( + patch.object(hardware.shutil, "which", return_value="/usr/bin/sysctl"), + patch.object(hardware.subprocess, "check_output", side_effect=fake_check_output), + ): + result = hardware._run(["sysctl", "-n", "machdep.cpu.brand_string"]) + + assert result == "ok" + # First element is the resolved absolute path, not the bare name; the + # remaining arguments are preserved verbatim. + assert captured["cmd"] == ["/usr/bin/sysctl", "-n", "machdep.cpu.brand_string"] + + def test_run_skips_when_binary_absent(self): + """A binary missing from PATH resolves to None → probe is skipped and the + subprocess is never spawned.""" + with ( + patch.object(hardware.shutil, "which", return_value=None), + patch.object(hardware.subprocess, "check_output") as mock_run, + ): + assert hardware._run(["nvidia-smi", "--query-gpu=name"]) is None + mock_run.assert_not_called() + + def test_run_never_raises_on_resolve_failure(self): + """Even a broken PATH lookup degrades to None, honoring the never-raise + contract.""" + with patch.object(hardware.shutil, "which", side_effect=RuntimeError("boom")): + assert hardware._run(["nvidia-smi"]) is None + + +class TestResolveBinaryWindowsCwdGuard: + """On Windows, a binary that ``shutil.which`` resolves inside the current + working directory is rejected — closing the CWD binary-planting hole.""" + + def test_windows_rejects_binary_planted_in_cwd(self, tmp_path): + planted = tmp_path / "nvidia-smi.exe" + planted.write_text("") + with ( + patch.object(hardware.platform, "system", return_value="Windows"), + patch.object(hardware.os, "getcwd", return_value=str(tmp_path)), + patch.object(hardware.shutil, "which", return_value=str(planted)), + ): + assert hardware._resolve_binary("nvidia-smi") is None + + def test_windows_allows_system_binary_outside_cwd(self, tmp_path): + cwd = tmp_path / "attacker" + system_dir = tmp_path / "System32" + cwd.mkdir() + system_dir.mkdir() + legit = system_dir / "nvidia-smi.exe" + legit.write_text("") + with ( + patch.object(hardware.platform, "system", return_value="Windows"), + patch.object(hardware.os, "getcwd", return_value=str(cwd)), + patch.object(hardware.shutil, "which", return_value=str(legit)), + ): + assert hardware._resolve_binary("nvidia-smi") == str(legit) + + def test_non_windows_does_not_apply_cwd_guard(self, tmp_path): + """POSIX ``shutil.which`` never searches the CWD, so the guard is + Windows-only; a resolved path is returned as-is.""" + resolved = tmp_path / "sysctl" + resolved.write_text("") + with ( + patch.object(hardware.platform, "system", return_value="Darwin"), + patch.object(hardware.os, "getcwd", return_value=str(tmp_path)), + patch.object(hardware.shutil, "which", return_value=str(resolved)), + ): + assert hardware._resolve_binary("sysctl") == str(resolved) From f9e7341ebb5e3cb7de0e70dcd15212021ab925ca Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 20:41:17 -0700 Subject: [PATCH 4/7] fix(hardware): tighten CWD anti-planting guard per Cursor review (BE-3434) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address cursor-review findings on the probe binary CWD guard: - Medium (5/8): the old subtree check rejected any binary whose CWD is an ancestor, so running `comfy env` from `C:\Windows` (or a drive root) silently disabled GPU/CPU detection for `C:\Windows\System32\nvidia-smi.exe`. Reject only a binary sitting *directly* in the CWD (the actual planting signature); a legitimate system binary in a subdirectory is untouched. - Low (1/8): apply the guard on every platform, not just Windows — a `.`/empty entry in POSIX `$PATH` lets `shutil.which` return a CWD match too. - Low (1/8): normalize both paths with `os.path.normcase` so Windows' case-insensitivity can't fail the guard open. - Low (4/8): guard `_run` against an empty `cmd` so it degrades to None instead of raising IndexError, honoring the never-raise contract. Renames `_is_within_cwd` -> `_is_planted_in_cwd`; updates and extends tests (subdirectory-of-CWD allowed, POSIX plant rejected, empty-cmd). Co-Authored-By: Claude Opus 4.8 --- comfy_cli/hardware.py | 51 ++++++++++++++++++---------- tests/comfy_cli/test_hardware.py | 57 +++++++++++++++++++++++++++----- 2 files changed, 82 insertions(+), 26 deletions(-) diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py index 41537287..d69a2367 100644 --- a/comfy_cli/hardware.py +++ b/comfy_cli/hardware.py @@ -28,20 +28,30 @@ _SUBPROCESS_TIMEOUT = 5 -def _is_within_cwd(path: str) -> bool: - """Return ``True`` only if ``path`` provably resolves inside ``os.getcwd()``. - - Used to reject a probe binary that resolves to the current working directory - (see :func:`_resolve_binary`). Ambiguity — a path on a different drive, or an - unresolvable one — is treated as *not* under the CWD so a legitimate system - binary is never rejected. +def _is_planted_in_cwd(path: str) -> bool: + """Return ``True`` only if ``path`` resolves to a file sitting *directly* in + ``os.getcwd()`` — the signature of a planted probe binary. + + ``shutil.which`` searches the current directory first on Windows (and on any + platform whose ``$PATH`` contains ``.`` or an empty entry), so an attacker who + controls the directory the user runs ``comfy env`` from can drop a malicious + ``nvidia-smi.exe`` there. Such a plant always lands in the CWD *itself*, so we + reject only a resolved binary whose parent directory **is** the CWD. A + legitimate system binary in a *subdirectory* — e.g. ``System32`` even when the + CWD is ``C:\\Windows``, or a drive root — is left untouched, honouring the + "a legitimate system binary is never rejected" guarantee. Paths are compared + with :func:`os.path.normcase` so Windows' case-insensitivity can't fail the + guard open. Ambiguity — a path on a different drive, or an unresolvable one — + is treated as *not* planted so a legitimate binary is never rejected. """ try: - cwd = os.path.realpath(os.getcwd()) - resolved = os.path.realpath(path) - return os.path.commonpath([cwd, resolved]) == cwd + cwd = os.path.normcase(os.path.realpath(os.getcwd())) + # ``os.path.dirname`` of a bare/relative ``which`` result is "", which + # ``realpath`` correctly resolves against the CWD. + parent = os.path.normcase(os.path.realpath(os.path.dirname(path))) + return parent == cwd except (OSError, ValueError): - # Different drives (Windows) or an unresolvable path → not under the CWD. + # Different drives (Windows) or an unresolvable path → not planted. return False @@ -55,17 +65,18 @@ def _resolve_binary(name: str) -> str | None: working directory, so running ``comfy env`` from an attacker-controlled directory cannot execute a planted ``nvidia-smi.exe``. - ``shutil.which`` on Windows may itself resolve against the current directory, - so as defense-in-depth a resolved path inside the CWD tree is additionally - rejected. A legitimate system binary (e.g. ``nvidia-smi.exe`` under - ``System32``) is unaffected. + ``shutil.which`` may itself resolve against the current directory (always on + Windows; on any platform when ``$PATH`` holds ``.`` or an empty entry), so as + defense-in-depth a resolved path planted directly in the CWD is additionally + rejected on every platform. A legitimate system binary (e.g. ``nvidia-smi.exe`` + under ``System32``) is unaffected. """ try: path = shutil.which(name) if path is None: return None - if platform.system() == "Windows" and _is_within_cwd(path): - logger.debug("skipping hardware probe %r: resolved under CWD (%s)", name, path) + if _is_planted_in_cwd(path): + logger.debug("skipping hardware probe %r: resolved into CWD (%s)", name, path) return None return path except Exception: @@ -78,8 +89,12 @@ def _run(cmd: list[str]) -> str | None: ``cmd[0]`` is resolved to a trusted absolute path via :func:`_resolve_binary` before execution (skipping the probe when absent or CWD-planted), and the run - is bounded by ``timeout=5`` so a hung binary can never block the probe. + is bounded by ``timeout=5`` so a hung binary can never block the probe. An + empty ``cmd`` degrades to ``None`` rather than raising, honouring the + module's never-raise contract. """ + if not cmd: + return None resolved = _resolve_binary(cmd[0]) if resolved is None: return None diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py index 0b2ded75..9cd02635 100644 --- a/tests/comfy_cli/test_hardware.py +++ b/tests/comfy_cli/test_hardware.py @@ -373,10 +373,20 @@ def test_run_never_raises_on_resolve_failure(self): with patch.object(hardware.shutil, "which", side_effect=RuntimeError("boom")): assert hardware._run(["nvidia-smi"]) is None + def test_run_empty_cmd_returns_none(self): + """An empty command degrades to None instead of raising IndexError, + honoring the never-raise contract.""" + with patch.object(hardware.subprocess, "check_output") as mock_run: + assert hardware._run([]) is None + mock_run.assert_not_called() + -class TestResolveBinaryWindowsCwdGuard: - """On Windows, a binary that ``shutil.which`` resolves inside the current - working directory is rejected — closing the CWD binary-planting hole.""" +class TestResolveBinaryCwdGuard: + """A binary that ``shutil.which`` resolves *directly inside* the current + working directory is rejected — closing the CWD binary-planting hole. The + guard fires on every platform (``$PATH`` can search the CWD on POSIX too via a + ``.``/empty entry) and rejects only the immediate directory so a legitimate + system binary in a subdirectory is never lost.""" def test_windows_rejects_binary_planted_in_cwd(self, tmp_path): planted = tmp_path / "nvidia-smi.exe" @@ -402,14 +412,45 @@ def test_windows_allows_system_binary_outside_cwd(self, tmp_path): ): assert hardware._resolve_binary("nvidia-smi") == str(legit) - def test_non_windows_does_not_apply_cwd_guard(self, tmp_path): - """POSIX ``shutil.which`` never searches the CWD, so the guard is - Windows-only; a resolved path is returned as-is.""" - resolved = tmp_path / "sysctl" - resolved.write_text("") + def test_allows_system_binary_in_subdirectory_of_cwd(self, tmp_path): + """Running from an ancestor of the binary (e.g. ``C:\\Windows`` with the + real binary under ``System32``) must NOT reject it — only a binary + directly in the CWD is a plant.""" + system_dir = tmp_path / "System32" + system_dir.mkdir() + legit = system_dir / "nvidia-smi.exe" + legit.write_text("") + with ( + patch.object(hardware.platform, "system", return_value="Windows"), + # CWD is the ANCESTOR (tmp_path), binary lives one level deeper. + patch.object(hardware.os, "getcwd", return_value=str(tmp_path)), + patch.object(hardware.shutil, "which", return_value=str(legit)), + ): + assert hardware._resolve_binary("nvidia-smi") == str(legit) + + def test_posix_also_rejects_binary_planted_in_cwd(self, tmp_path): + """A ``.``/empty entry in ``$PATH`` lets ``shutil.which`` return a CWD + match on POSIX too, so the guard applies there as well.""" + planted = tmp_path / "nvidia-smi" + planted.write_text("") with ( patch.object(hardware.platform, "system", return_value="Darwin"), patch.object(hardware.os, "getcwd", return_value=str(tmp_path)), + patch.object(hardware.shutil, "which", return_value=str(planted)), + ): + assert hardware._resolve_binary("nvidia-smi") is None + + def test_posix_allows_system_binary_outside_cwd(self, tmp_path): + """A legitimate binary outside the CWD is returned as-is on POSIX.""" + cwd = tmp_path / "project" + bin_dir = tmp_path / "usr_bin" + cwd.mkdir() + bin_dir.mkdir() + resolved = bin_dir / "sysctl" + resolved.write_text("") + with ( + patch.object(hardware.platform, "system", return_value="Darwin"), + patch.object(hardware.os, "getcwd", return_value=str(cwd)), patch.object(hardware.shutil, "which", return_value=str(resolved)), ): assert hardware._resolve_binary("sysctl") == str(resolved) From dd3b9fad3026eda70e49f49251d70347bb46029c Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Thu, 30 Jul 2026 08:57:28 -0700 Subject: [PATCH 5/7] fix(hardware): reject CWD-anchored relative probe resolutions (BE-3434) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `shutil.which` returns `os.path.join(entry, name)`, so a relative `$PATH` entry (`.`, an empty entry, `subdir`, or Windows' implicitly prepended `os.curdir`) yields a RELATIVE match. `_is_planted_in_cwd` correctly lets `subdir/nvidia-smi` through — its parent is not the CWD — but `_run` then handed that relative string to `subprocess.check_output`, which re-resolves it against the attacker-controlled CWD: the exact hijack this PR closes. `_resolve_binary` now skips any non-absolute `which` result. A relative result means the matching `$PATH` entry was itself relative, so the binary is anchored under the CWD; a binary found through a normal absolute `$PATH` entry always comes back absolute and is unaffected. This is strictly stronger than normalising with `abspath`, which would still spawn `/subdir/nvidia-smi`. Adds four tests: the relative-subdirectory, `./`-relative and bare-name `which` results are all rejected, and `_run` never spawns a relative path. Addresses CodeRabbit review thread r3684167296. Co-Authored-By: Claude Opus 5 --- comfy_cli/hardware.py | 23 ++++++++++++++--- tests/comfy_cli/test_hardware.py | 42 ++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py index 30fe5f06..8ff0e61b 100644 --- a/comfy_cli/hardware.py +++ b/comfy_cli/hardware.py @@ -67,14 +67,31 @@ def _resolve_binary(name: str) -> str | None: ``shutil.which`` may itself resolve against the current directory (always on Windows; on any platform when ``$PATH`` holds ``.`` or an empty entry), so as - defense-in-depth a resolved path planted directly in the CWD is additionally - rejected on every platform. A legitimate system binary (e.g. ``nvidia-smi.exe`` - under ``System32``) is unaffected. + defense-in-depth two CWD-anchored results are additionally rejected on every + platform: + + * a **relative** result. ``which`` returns ``os.path.join(entry, name)``, so a + relative path means the matching ``$PATH`` entry was itself relative (``.``, + an empty entry, ``subdir``, or Windows' implicitly prepended ``os.curdir``) + and the binary therefore lives under the attacker-controlled CWD. Handing + that string to :func:`subprocess.check_output` would re-resolve it against + the CWD — exactly the hijack this function exists to prevent — so the probe + is skipped instead. A binary found through a normal absolute ``$PATH`` entry + always comes back absolute and is unaffected. + * an absolute result sitting directly **in** the CWD (see + :func:`_is_planted_in_cwd`), which covers the CWD appearing in ``$PATH`` as + an absolute entry. + + A legitimate system binary (e.g. ``nvidia-smi.exe`` under ``System32``) is + unaffected by either check. """ try: path = shutil.which(name) if path is None: return None + if not os.path.isabs(path): + logger.debug("skipping hardware probe %r: relative PATH match anchored in CWD (%s)", name, path) + return None if _is_planted_in_cwd(path): logger.debug("skipping hardware probe %r: resolved into CWD (%s)", name, path) return None diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py index 18cb11d7..4edcab45 100644 --- a/tests/comfy_cli/test_hardware.py +++ b/tests/comfy_cli/test_hardware.py @@ -10,6 +10,7 @@ import ctypes import json +import os from pathlib import Path from unittest.mock import patch @@ -472,3 +473,44 @@ def test_posix_allows_system_binary_outside_cwd(self, tmp_path): patch.object(hardware.shutil, "which", return_value=str(resolved)), ): assert hardware._resolve_binary("sysctl") == str(resolved) + + +class TestResolveBinaryRejectsRelativeMatches: + """``shutil.which`` returns ``os.path.join(entry, name)``, so a relative + ``$PATH`` entry yields a relative match anchored in the CWD. Executing that + string would let ``subprocess`` re-resolve it against the attacker-controlled + CWD, so such a match is skipped rather than run.""" + + def test_rejects_relative_subdirectory_match(self): + """``PATH=subdir`` → ``subdir/nvidia-smi``: not *directly* in the CWD, so + the planted-in-CWD guard lets it through — the absolute-path check is what + stops it.""" + relative = os.path.join("subdir", "nvidia-smi") + # Precondition: this is exactly the case the CWD guard does NOT catch. + assert not hardware._is_planted_in_cwd(relative) + with patch.object(hardware.shutil, "which", return_value=relative): + assert hardware._resolve_binary("nvidia-smi") is None + + def test_rejects_dot_relative_match(self): + """Windows prepends ``os.curdir`` to the search path, so a CWD plant comes + back as ``.\\nvidia-smi.exe``.""" + with ( + patch.object(hardware.platform, "system", return_value="Windows"), + patch.object(hardware.shutil, "which", return_value=os.path.join(os.curdir, "nvidia-smi.exe")), + ): + assert hardware._resolve_binary("nvidia-smi") is None + + def test_rejects_bare_name_match(self): + """An empty ``$PATH`` entry joins to a bare name, which ``subprocess`` + would resolve by its own PATH/CWD search — the bare-name invocation this + PR removes.""" + with patch.object(hardware.shutil, "which", return_value="nvidia-smi"): + assert hardware._resolve_binary("nvidia-smi") is None + + def test_run_never_spawns_a_relative_path(self): + with ( + patch.object(hardware.shutil, "which", return_value=os.path.join("subdir", "nvidia-smi")), + patch.object(hardware.subprocess, "check_output") as mock_run, + ): + assert hardware._run(["nvidia-smi", "--query-gpu=name"]) is None + mock_run.assert_not_called() From fc981ea297c3c7a6796899888fbd6a585e364011 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Thu, 30 Jul 2026 09:18:29 -0700 Subject: [PATCH 6/7] fix(cuda): resolve nvidia-smi via absolute path to block Windows CWD planting (BE-5357) Promote hardware.py's _resolve_binary()/_is_planted_in_cwd() into a shared leaf module, comfy_cli/_safe_exec.py, and route cuda_detect's nvidia-smi fallback probe through it. cuda_detect invoked the probe by bare name, so on Windows CreateProcess searched the current working directory before $PATH and a planted nvidia-smi.exe would run - the same vector #567 closed in hardware.py. The helper lives in its own module rather than being imported from hardware because hardware already imports cuda_detect. Resolution returning None means "skip the probe", which lands on _detect_via_nvidia_smi's existing degrade-to-None contract, so no new error handling is required. hardware._run now delegates to the shared helper; its behaviour is unchanged. --- comfy_cli/_safe_exec.py | 99 +++++++++++++++++++++ comfy_cli/cuda_detect.py | 17 +++- comfy_cli/hardware.py | 89 ++----------------- tests/comfy_cli/test_cuda_detect.py | 74 ++++++++++++++++ tests/comfy_cli/test_hardware.py | 120 ++------------------------ tests/comfy_cli/test_safe_exec.py | 129 ++++++++++++++++++++++++++++ 6 files changed, 332 insertions(+), 196 deletions(-) create mode 100644 comfy_cli/_safe_exec.py create mode 100644 tests/comfy_cli/test_safe_exec.py diff --git a/comfy_cli/_safe_exec.py b/comfy_cli/_safe_exec.py new file mode 100644 index 00000000..63fd87a7 --- /dev/null +++ b/comfy_cli/_safe_exec.py @@ -0,0 +1,99 @@ +"""Shared helpers for spawning trusted system binaries. + +Windows' ``CreateProcess`` searches the *current working directory* before +``$PATH``, so invoking a probe by bare name (``["nvidia-smi"]``) from an +attacker-prepared directory executes whatever ``nvidia-smi.exe`` was planted +there. :func:`resolve_binary` closes that vector by resolving the name to a +trusted absolute path up front, and returning ``None`` — "skip this probe" — +whenever the only match is anchored in the CWD. + +The module is a leaf on purpose: it imports nothing from ``comfy_cli``, so both +:mod:`comfy_cli.hardware` and :mod:`comfy_cli.cuda_detect` can use it without +the import cycle that would come from ``cuda_detect`` importing ``hardware`` +(``hardware`` already imports ``cuda_detect``). + +Contract: **never raises.** Every failure — a missing binary, a broken ``$PATH`` +lookup, an unresolvable path — degrades to ``None`` so callers can keep their +existing degrade-to-``None`` behaviour without new error handling. +""" + +from __future__ import annotations + +import logging +import os +import shutil + +logger = logging.getLogger(__name__) + + +def is_planted_in_cwd(path: str) -> bool: + """Return ``True`` only if ``path`` resolves to a file sitting *directly* in + ``os.getcwd()`` — the signature of a planted probe binary. + + ``shutil.which`` searches the current directory first on Windows (and on any + platform whose ``$PATH`` contains ``.`` or an empty entry), so an attacker who + controls the directory the user runs ``comfy`` from can drop a malicious + ``nvidia-smi.exe`` there. Such a plant always lands in the CWD *itself*, so we + reject only a resolved binary whose parent directory **is** the CWD. A + legitimate system binary in a *subdirectory* — e.g. ``System32`` even when the + CWD is ``C:\\Windows``, or a drive root — is left untouched, honouring the + "a legitimate system binary is never rejected" guarantee. Paths are compared + with :func:`os.path.normcase` so Windows' case-insensitivity can't fail the + guard open. Ambiguity — a path on a different drive, or an unresolvable one — + is treated as *not* planted so a legitimate binary is never rejected. + """ + try: + cwd = os.path.normcase(os.path.realpath(os.getcwd())) + # ``os.path.dirname`` of a bare/relative ``which`` result is "", which + # ``realpath`` correctly resolves against the CWD. + parent = os.path.normcase(os.path.realpath(os.path.dirname(path))) + return parent == cwd + except (OSError, ValueError): + # Different drives (Windows) or an unresolvable path → not planted. + return False + + +def resolve_binary(name: str) -> str | None: + """Resolve a system binary to a trusted absolute path, or ``None`` to skip it. + + :func:`shutil.which` performs a PATH lookup and returns ``None`` when the + binary is absent (so the caller simply degrades to ``None``). Passing the + resolved absolute path to :mod:`subprocess` — rather than the bare name — + prevents Windows ``CreateProcess`` from searching the current working + directory, so running ``comfy`` from an attacker-controlled directory cannot + execute a planted ``nvidia-smi.exe``. + + ``shutil.which`` may itself resolve against the current directory (always on + Windows; on any platform when ``$PATH`` holds ``.`` or an empty entry), so as + defense-in-depth two CWD-anchored results are additionally rejected on every + platform: + + * a **relative** result. ``which`` returns ``os.path.join(entry, name)``, so a + relative path means the matching ``$PATH`` entry was itself relative (``.``, + an empty entry, ``subdir``, or Windows' implicitly prepended ``os.curdir``) + and the binary therefore lives under the attacker-controlled CWD. Handing + that string to :mod:`subprocess` would re-resolve it against the CWD — + exactly the hijack this function exists to prevent — so the probe is skipped + instead. A binary found through a normal absolute ``$PATH`` entry always + comes back absolute and is unaffected. + * an absolute result sitting directly **in** the CWD (see + :func:`is_planted_in_cwd`), which covers the CWD appearing in ``$PATH`` as + an absolute entry. + + A legitimate system binary (e.g. ``nvidia-smi.exe`` under ``System32``) is + unaffected by either check. + """ + try: + path = shutil.which(name) + if path is None: + return None + if not os.path.isabs(path): + logger.debug("skipping %r: relative PATH match anchored in CWD (%s)", name, path) + return None + if is_planted_in_cwd(path): + logger.debug("skipping %r: resolved into CWD (%s)", name, path) + return None + return path + except Exception: + logger.debug("resolving binary %r failed", name, exc_info=True) + return None diff --git a/comfy_cli/cuda_detect.py b/comfy_cli/cuda_detect.py index f6339f36..b04cbdc2 100644 --- a/comfy_cli/cuda_detect.py +++ b/comfy_cli/cuda_detect.py @@ -9,6 +9,8 @@ import re import subprocess +from comfy_cli import _safe_exec + logger = logging.getLogger(__name__) PYTORCH_CUDA_WHEELS: list[str] = [ @@ -77,10 +79,21 @@ def _detect_via_ctypes() -> int | None: def _detect_via_nvidia_smi() -> tuple[int, int] | None: - """Parse CUDA version from nvidia-smi output, or return None.""" + """Parse CUDA version from nvidia-smi output, or return None. + + ``nvidia-smi`` is resolved to a trusted absolute path first: invoking it by + bare name would let Windows' ``CreateProcess`` pick up an ``nvidia-smi.exe`` + planted in the current working directory. A binary that is absent, or whose + only match is anchored in the CWD, resolves to ``None`` and the probe is + skipped — the same degrade-to-``None`` outcome as a failed run. + """ + nvidia_smi = _safe_exec.resolve_binary("nvidia-smi") + if nvidia_smi is None: + return None + try: output = subprocess.check_output( - ["nvidia-smi"], + [nvidia_smi], text=True, timeout=10, stderr=subprocess.DEVNULL, diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py index 8ff0e61b..572eb2d4 100644 --- a/comfy_cli/hardware.py +++ b/comfy_cli/hardware.py @@ -16,103 +16,30 @@ import logging import os import platform -import shutil import subprocess import psutil -from comfy_cli import cuda_detect +from comfy_cli import _safe_exec, cuda_detect logger = logging.getLogger(__name__) _SUBPROCESS_TIMEOUT = 5 -def _is_planted_in_cwd(path: str) -> bool: - """Return ``True`` only if ``path`` resolves to a file sitting *directly* in - ``os.getcwd()`` — the signature of a planted probe binary. - - ``shutil.which`` searches the current directory first on Windows (and on any - platform whose ``$PATH`` contains ``.`` or an empty entry), so an attacker who - controls the directory the user runs ``comfy env`` from can drop a malicious - ``nvidia-smi.exe`` there. Such a plant always lands in the CWD *itself*, so we - reject only a resolved binary whose parent directory **is** the CWD. A - legitimate system binary in a *subdirectory* — e.g. ``System32`` even when the - CWD is ``C:\\Windows``, or a drive root — is left untouched, honouring the - "a legitimate system binary is never rejected" guarantee. Paths are compared - with :func:`os.path.normcase` so Windows' case-insensitivity can't fail the - guard open. Ambiguity — a path on a different drive, or an unresolvable one — - is treated as *not* planted so a legitimate binary is never rejected. - """ - try: - cwd = os.path.normcase(os.path.realpath(os.getcwd())) - # ``os.path.dirname`` of a bare/relative ``which`` result is "", which - # ``realpath`` correctly resolves against the CWD. - parent = os.path.normcase(os.path.realpath(os.path.dirname(path))) - return parent == cwd - except (OSError, ValueError): - # Different drives (Windows) or an unresolvable path → not planted. - return False - - -def _resolve_binary(name: str) -> str | None: - """Resolve a probe binary to a trusted absolute path, or ``None`` to skip it. - - :func:`shutil.which` performs a PATH lookup and returns ``None`` when the - binary is absent (so the probe simply degrades to ``None``). Passing the - resolved absolute path to :func:`subprocess.check_output` — rather than the - bare name — prevents Windows ``CreateProcess`` from searching the current - working directory, so running ``comfy env`` from an attacker-controlled - directory cannot execute a planted ``nvidia-smi.exe``. - - ``shutil.which`` may itself resolve against the current directory (always on - Windows; on any platform when ``$PATH`` holds ``.`` or an empty entry), so as - defense-in-depth two CWD-anchored results are additionally rejected on every - platform: - - * a **relative** result. ``which`` returns ``os.path.join(entry, name)``, so a - relative path means the matching ``$PATH`` entry was itself relative (``.``, - an empty entry, ``subdir``, or Windows' implicitly prepended ``os.curdir``) - and the binary therefore lives under the attacker-controlled CWD. Handing - that string to :func:`subprocess.check_output` would re-resolve it against - the CWD — exactly the hijack this function exists to prevent — so the probe - is skipped instead. A binary found through a normal absolute ``$PATH`` entry - always comes back absolute and is unaffected. - * an absolute result sitting directly **in** the CWD (see - :func:`_is_planted_in_cwd`), which covers the CWD appearing in ``$PATH`` as - an absolute entry. - - A legitimate system binary (e.g. ``nvidia-smi.exe`` under ``System32``) is - unaffected by either check. - """ - try: - path = shutil.which(name) - if path is None: - return None - if not os.path.isabs(path): - logger.debug("skipping hardware probe %r: relative PATH match anchored in CWD (%s)", name, path) - return None - if _is_planted_in_cwd(path): - logger.debug("skipping hardware probe %r: resolved into CWD (%s)", name, path) - return None - return path - except Exception: - logger.debug("resolving hardware probe binary %r failed", name, exc_info=True) - return None - - def _run(cmd: list[str]) -> str | None: """Run ``cmd`` and return stripped stdout, or ``None`` on any failure. - ``cmd[0]`` is resolved to a trusted absolute path via :func:`_resolve_binary` - before execution (skipping the probe when absent or CWD-planted), and the run - is bounded by ``timeout=5`` so a hung binary can never block the probe. An - empty ``cmd`` degrades to ``None`` rather than raising, honouring the - module's never-raise contract. + ``cmd[0]`` is resolved to a trusted absolute path via + :func:`comfy_cli._safe_exec.resolve_binary` before execution (skipping the + probe when the binary is absent or CWD-planted), and the run is bounded by + ``timeout=5`` so a hung binary can never block the probe. An empty ``cmd`` + degrades to ``None`` rather than raising, honouring the module's never-raise + contract. """ if not cmd: return None - resolved = _resolve_binary(cmd[0]) + resolved = _safe_exec.resolve_binary(cmd[0]) if resolved is None: return None try: diff --git a/tests/comfy_cli/test_cuda_detect.py b/tests/comfy_cli/test_cuda_detect.py index 2a1cf310..5e4eefaf 100644 --- a/tests/comfy_cli/test_cuda_detect.py +++ b/tests/comfy_cli/test_cuda_detect.py @@ -1,8 +1,10 @@ +import os import subprocess from unittest.mock import MagicMock, patch import pytest +from comfy_cli import _safe_exec from comfy_cli.cuda_detect import ( DEFAULT_CUDA_TAG, PYTORCH_CUDA_WHEELS, @@ -67,7 +69,19 @@ def test_cuinit_fails(self): assert _detect_via_ctypes() is None +_RESOLVED_SMI = os.path.join(os.sep, "usr", "bin", "nvidia-smi") + + class TestDetectViaNvidiaSmi: + @pytest.fixture(autouse=True) + def _resolved_nvidia_smi(self): + """``_detect_via_nvidia_smi`` resolves ``nvidia-smi`` to an absolute path + before spawning it, so these parse-level tests stub the resolution to a + fixed path — they then run identically on a machine with no NVIDIA driver + installed.""" + with patch("comfy_cli.cuda_detect._safe_exec.resolve_binary", return_value=_RESOLVED_SMI): + yield + def test_happy_path(self): output = ( "Mon Mar 30 12:00:00 2026\n" @@ -98,6 +112,66 @@ def test_timeout(self): assert _detect_via_nvidia_smi() is None +class TestDetectViaNvidiaSmiResolvesBinaryPath: + """``_detect_via_nvidia_smi`` must spawn a resolved absolute path, never the + bare name — otherwise Windows' ``CreateProcess`` searches the current working + directory first and runs an ``nvidia-smi.exe`` planted there.""" + + def test_invokes_resolved_absolute_path(self): + captured = {} + + def fake_check_output(cmd, **kwargs): + captured["cmd"] = cmd + return "| Driver Version: 560.35.03 CUDA Version: 12.6 |\n" + + with ( + patch.object(_safe_exec.shutil, "which", return_value=_RESOLVED_SMI), + patch("comfy_cli.cuda_detect.subprocess.check_output", side_effect=fake_check_output), + ): + assert _detect_via_nvidia_smi() == (12, 6) + + assert captured["cmd"] == [_RESOLVED_SMI] + + def test_skips_probe_when_binary_absent(self): + """A missing binary resolves to None → the probe is skipped entirely and + no subprocess is spawned, degrading to the same None as a failed run.""" + with ( + patch.object(_safe_exec.shutil, "which", return_value=None), + patch("comfy_cli.cuda_detect.subprocess.check_output") as mock_run, + ): + assert _detect_via_nvidia_smi() is None + mock_run.assert_not_called() + + def test_rejects_binary_planted_in_cwd(self, tmp_path): + """The BE-3434 vector: an ``nvidia-smi`` sitting directly in the CWD is + never executed.""" + planted = tmp_path / "nvidia-smi" + planted.write_text("") + with ( + patch.object(_safe_exec.os, "getcwd", return_value=str(tmp_path)), + patch.object(_safe_exec.shutil, "which", return_value=str(planted)), + patch("comfy_cli.cuda_detect.subprocess.check_output") as mock_run, + ): + assert _detect_via_nvidia_smi() is None + mock_run.assert_not_called() + + def test_never_spawns_a_relative_match(self): + """A relative ``which`` result is anchored in the CWD; handing it to + ``subprocess`` would re-resolve it there, so the probe is skipped.""" + with ( + patch.object(_safe_exec.shutil, "which", return_value=os.path.join("subdir", "nvidia-smi")), + patch("comfy_cli.cuda_detect.subprocess.check_output") as mock_run, + ): + assert _detect_via_nvidia_smi() is None + mock_run.assert_not_called() + + def test_resolution_failure_degrades_to_none(self): + """A broken PATH lookup keeps the existing degrade-to-None contract — no + new exception escapes the probe.""" + with patch.object(_safe_exec.shutil, "which", side_effect=RuntimeError("boom")): + assert _detect_via_nvidia_smi() is None + + class TestDetectCudaDriverVersion: def test_ctypes_success_skips_smi(self): lib = MagicMock() diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py index 4edcab45..ddc7a9a2 100644 --- a/tests/comfy_cli/test_hardware.py +++ b/tests/comfy_cli/test_hardware.py @@ -16,7 +16,7 @@ import jsonschema -from comfy_cli import hardware +from comfy_cli import _safe_exec, hardware from comfy_cli.env_checker import EnvChecker, format_hardware_summary _EnvCheckerCls = EnvChecker.__closure__[0].cell_contents @@ -366,7 +366,7 @@ def fake_check_output(cmd, **kwargs): return "ok\n" with ( - patch.object(hardware.shutil, "which", return_value="/usr/bin/sysctl"), + patch.object(_safe_exec.shutil, "which", return_value="/usr/bin/sysctl"), patch.object(hardware.subprocess, "check_output", side_effect=fake_check_output), ): result = hardware._run(["sysctl", "-n", "machdep.cpu.brand_string"]) @@ -380,7 +380,7 @@ def test_run_skips_when_binary_absent(self): """A binary missing from PATH resolves to None → probe is skipped and the subprocess is never spawned.""" with ( - patch.object(hardware.shutil, "which", return_value=None), + patch.object(_safe_exec.shutil, "which", return_value=None), patch.object(hardware.subprocess, "check_output") as mock_run, ): assert hardware._run(["nvidia-smi", "--query-gpu=name"]) is None @@ -389,7 +389,7 @@ def test_run_skips_when_binary_absent(self): def test_run_never_raises_on_resolve_failure(self): """Even a broken PATH lookup degrades to None, honoring the never-raise contract.""" - with patch.object(hardware.shutil, "which", side_effect=RuntimeError("boom")): + with patch.object(_safe_exec.shutil, "which", side_effect=RuntimeError("boom")): assert hardware._run(["nvidia-smi"]) is None def test_run_empty_cmd_returns_none(self): @@ -399,117 +399,11 @@ def test_run_empty_cmd_returns_none(self): assert hardware._run([]) is None mock_run.assert_not_called() - -class TestResolveBinaryCwdGuard: - """A binary that ``shutil.which`` resolves *directly inside* the current - working directory is rejected — closing the CWD binary-planting hole. The - guard fires on every platform (``$PATH`` can search the CWD on POSIX too via a - ``.``/empty entry) and rejects only the immediate directory so a legitimate - system binary in a subdirectory is never lost.""" - - def test_windows_rejects_binary_planted_in_cwd(self, tmp_path): - planted = tmp_path / "nvidia-smi.exe" - planted.write_text("") - with ( - patch.object(hardware.platform, "system", return_value="Windows"), - patch.object(hardware.os, "getcwd", return_value=str(tmp_path)), - patch.object(hardware.shutil, "which", return_value=str(planted)), - ): - assert hardware._resolve_binary("nvidia-smi") is None - - def test_windows_allows_system_binary_outside_cwd(self, tmp_path): - cwd = tmp_path / "attacker" - system_dir = tmp_path / "System32" - cwd.mkdir() - system_dir.mkdir() - legit = system_dir / "nvidia-smi.exe" - legit.write_text("") - with ( - patch.object(hardware.platform, "system", return_value="Windows"), - patch.object(hardware.os, "getcwd", return_value=str(cwd)), - patch.object(hardware.shutil, "which", return_value=str(legit)), - ): - assert hardware._resolve_binary("nvidia-smi") == str(legit) - - def test_allows_system_binary_in_subdirectory_of_cwd(self, tmp_path): - """Running from an ancestor of the binary (e.g. ``C:\\Windows`` with the - real binary under ``System32``) must NOT reject it — only a binary - directly in the CWD is a plant.""" - system_dir = tmp_path / "System32" - system_dir.mkdir() - legit = system_dir / "nvidia-smi.exe" - legit.write_text("") - with ( - patch.object(hardware.platform, "system", return_value="Windows"), - # CWD is the ANCESTOR (tmp_path), binary lives one level deeper. - patch.object(hardware.os, "getcwd", return_value=str(tmp_path)), - patch.object(hardware.shutil, "which", return_value=str(legit)), - ): - assert hardware._resolve_binary("nvidia-smi") == str(legit) - - def test_posix_also_rejects_binary_planted_in_cwd(self, tmp_path): - """A ``.``/empty entry in ``$PATH`` lets ``shutil.which`` return a CWD - match on POSIX too, so the guard applies there as well.""" - planted = tmp_path / "nvidia-smi" - planted.write_text("") - with ( - patch.object(hardware.platform, "system", return_value="Darwin"), - patch.object(hardware.os, "getcwd", return_value=str(tmp_path)), - patch.object(hardware.shutil, "which", return_value=str(planted)), - ): - assert hardware._resolve_binary("nvidia-smi") is None - - def test_posix_allows_system_binary_outside_cwd(self, tmp_path): - """A legitimate binary outside the CWD is returned as-is on POSIX.""" - cwd = tmp_path / "project" - bin_dir = tmp_path / "usr_bin" - cwd.mkdir() - bin_dir.mkdir() - resolved = bin_dir / "sysctl" - resolved.write_text("") - with ( - patch.object(hardware.platform, "system", return_value="Darwin"), - patch.object(hardware.os, "getcwd", return_value=str(cwd)), - patch.object(hardware.shutil, "which", return_value=str(resolved)), - ): - assert hardware._resolve_binary("sysctl") == str(resolved) - - -class TestResolveBinaryRejectsRelativeMatches: - """``shutil.which`` returns ``os.path.join(entry, name)``, so a relative - ``$PATH`` entry yields a relative match anchored in the CWD. Executing that - string would let ``subprocess`` re-resolve it against the attacker-controlled - CWD, so such a match is skipped rather than run.""" - - def test_rejects_relative_subdirectory_match(self): - """``PATH=subdir`` → ``subdir/nvidia-smi``: not *directly* in the CWD, so - the planted-in-CWD guard lets it through — the absolute-path check is what - stops it.""" - relative = os.path.join("subdir", "nvidia-smi") - # Precondition: this is exactly the case the CWD guard does NOT catch. - assert not hardware._is_planted_in_cwd(relative) - with patch.object(hardware.shutil, "which", return_value=relative): - assert hardware._resolve_binary("nvidia-smi") is None - - def test_rejects_dot_relative_match(self): - """Windows prepends ``os.curdir`` to the search path, so a CWD plant comes - back as ``.\\nvidia-smi.exe``.""" - with ( - patch.object(hardware.platform, "system", return_value="Windows"), - patch.object(hardware.shutil, "which", return_value=os.path.join(os.curdir, "nvidia-smi.exe")), - ): - assert hardware._resolve_binary("nvidia-smi") is None - - def test_rejects_bare_name_match(self): - """An empty ``$PATH`` entry joins to a bare name, which ``subprocess`` - would resolve by its own PATH/CWD search — the bare-name invocation this - PR removes.""" - with patch.object(hardware.shutil, "which", return_value="nvidia-smi"): - assert hardware._resolve_binary("nvidia-smi") is None - def test_run_never_spawns_a_relative_path(self): + """A relative ``which`` match is anchored in the CWD, so ``_run`` skips + the probe rather than letting ``subprocess`` re-resolve it there.""" with ( - patch.object(hardware.shutil, "which", return_value=os.path.join("subdir", "nvidia-smi")), + patch.object(_safe_exec.shutil, "which", return_value=os.path.join("subdir", "nvidia-smi")), patch.object(hardware.subprocess, "check_output") as mock_run, ): assert hardware._run(["nvidia-smi", "--query-gpu=name"]) is None diff --git a/tests/comfy_cli/test_safe_exec.py b/tests/comfy_cli/test_safe_exec.py new file mode 100644 index 00000000..83100de0 --- /dev/null +++ b/tests/comfy_cli/test_safe_exec.py @@ -0,0 +1,129 @@ +"""Tests for :mod:`comfy_cli._safe_exec`. + +The contract under test: ``resolve_binary`` hands back only a *trusted absolute +path*, and returns ``None`` — "skip this probe" — for every CWD-anchored match, +while never rejecting a legitimate system binary and never raising. +""" + +from __future__ import annotations + +import os +from unittest.mock import patch + +from comfy_cli import _safe_exec + + +class TestResolveBinaryCwdGuard: + """A binary that ``shutil.which`` resolves *directly inside* the current + working directory is rejected — closing the CWD binary-planting hole. The + guard fires on every platform (``$PATH`` can search the CWD on POSIX too via a + ``.``/empty entry) and rejects only the immediate directory so a legitimate + system binary in a subdirectory is never lost.""" + + def test_rejects_binary_planted_in_cwd(self, tmp_path): + planted = tmp_path / "nvidia-smi.exe" + planted.write_text("") + with ( + patch.object(_safe_exec.os, "getcwd", return_value=str(tmp_path)), + patch.object(_safe_exec.shutil, "which", return_value=str(planted)), + ): + assert _safe_exec.resolve_binary("nvidia-smi") is None + + def test_allows_system_binary_outside_cwd(self, tmp_path): + cwd = tmp_path / "attacker" + system_dir = tmp_path / "System32" + cwd.mkdir() + system_dir.mkdir() + legit = system_dir / "nvidia-smi.exe" + legit.write_text("") + with ( + patch.object(_safe_exec.os, "getcwd", return_value=str(cwd)), + patch.object(_safe_exec.shutil, "which", return_value=str(legit)), + ): + assert _safe_exec.resolve_binary("nvidia-smi") == str(legit) + + def test_allows_system_binary_in_subdirectory_of_cwd(self, tmp_path): + """Running from an ancestor of the binary (e.g. ``C:\\Windows`` with the + real binary under ``System32``) must NOT reject it — only a binary + directly in the CWD is a plant.""" + system_dir = tmp_path / "System32" + system_dir.mkdir() + legit = system_dir / "nvidia-smi.exe" + legit.write_text("") + with ( + # CWD is the ANCESTOR (tmp_path), binary lives one level deeper. + patch.object(_safe_exec.os, "getcwd", return_value=str(tmp_path)), + patch.object(_safe_exec.shutil, "which", return_value=str(legit)), + ): + assert _safe_exec.resolve_binary("nvidia-smi") == str(legit) + + def test_posix_style_plant_in_cwd_is_rejected(self, tmp_path): + """A ``.``/empty entry in ``$PATH`` lets ``shutil.which`` return a CWD + match on POSIX too, so the guard applies there as well.""" + planted = tmp_path / "nvidia-smi" + planted.write_text("") + with ( + patch.object(_safe_exec.os, "getcwd", return_value=str(tmp_path)), + patch.object(_safe_exec.shutil, "which", return_value=str(planted)), + ): + assert _safe_exec.resolve_binary("nvidia-smi") is None + + def test_posix_allows_system_binary_outside_cwd(self, tmp_path): + """A legitimate binary outside the CWD is returned as-is on POSIX.""" + cwd = tmp_path / "project" + bin_dir = tmp_path / "usr_bin" + cwd.mkdir() + bin_dir.mkdir() + resolved = bin_dir / "sysctl" + resolved.write_text("") + with ( + patch.object(_safe_exec.os, "getcwd", return_value=str(cwd)), + patch.object(_safe_exec.shutil, "which", return_value=str(resolved)), + ): + assert _safe_exec.resolve_binary("sysctl") == str(resolved) + + +class TestResolveBinaryRejectsRelativeMatches: + """``shutil.which`` returns ``os.path.join(entry, name)``, so a relative + ``$PATH`` entry yields a relative match anchored in the CWD. Executing that + string would let ``subprocess`` re-resolve it against the attacker-controlled + CWD, so such a match is skipped rather than run.""" + + def test_rejects_relative_subdirectory_match(self): + """``PATH=subdir`` → ``subdir/nvidia-smi``: not *directly* in the CWD, so + the planted-in-CWD guard lets it through — the absolute-path check is what + stops it.""" + relative = os.path.join("subdir", "nvidia-smi") + # Precondition: this is exactly the case the CWD guard does NOT catch. + assert not _safe_exec.is_planted_in_cwd(relative) + with patch.object(_safe_exec.shutil, "which", return_value=relative): + assert _safe_exec.resolve_binary("nvidia-smi") is None + + def test_rejects_dot_relative_match(self): + """Windows prepends ``os.curdir`` to the search path, so a CWD plant comes + back as ``.\\nvidia-smi.exe``.""" + with patch.object(_safe_exec.shutil, "which", return_value=os.path.join(os.curdir, "nvidia-smi.exe")): + assert _safe_exec.resolve_binary("nvidia-smi") is None + + def test_rejects_bare_name_match(self): + """An empty ``$PATH`` entry joins to a bare name, which ``subprocess`` + would resolve by its own PATH/CWD search — the bare-name invocation this + helper exists to remove.""" + with patch.object(_safe_exec.shutil, "which", return_value="nvidia-smi"): + assert _safe_exec.resolve_binary("nvidia-smi") is None + + +class TestResolveBinaryNeverRaises: + def test_missing_binary_returns_none(self): + with patch.object(_safe_exec.shutil, "which", return_value=None): + assert _safe_exec.resolve_binary("nvidia-smi") is None + + def test_broken_lookup_returns_none(self): + with patch.object(_safe_exec.shutil, "which", side_effect=RuntimeError("boom")): + assert _safe_exec.resolve_binary("nvidia-smi") is None + + def test_is_planted_in_cwd_swallows_path_errors(self): + """An unresolvable path (e.g. a different Windows drive) is treated as + *not* planted so a legitimate binary is never rejected.""" + with patch.object(_safe_exec.os, "getcwd", side_effect=OSError("gone")): + assert _safe_exec.is_planted_in_cwd("/usr/bin/nvidia-smi") is False From 980039e3c1942c56aba0363f82fd0964d2e2afee Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Thu, 30 Jul 2026 09:48:47 -0700 Subject: [PATCH 7/7] fix(cuda): harden resolve_binary per the Cursor review panel (BE-5357) Four review findings on the promoted _safe_exec helper: - is_planted_in_cwd failed open: an OSError/ValueError from getcwd()/realpath() returned "not planted", so resolve_binary handed the unvetted path back and it got spawned. It now fails closed - we cannot prove the binary sits outside the CWD, so the probe is skipped, which is the same degradation as the binary being absent. This was the module's only error path that did not degrade safely. - os.path.isabs is not "fully qualified" on Windows: ntpath.isabs accepts a drive-less rooted path (\tools\nvidia-smi.exe) on the 3.10-3.12 interpreters we support, and CreateProcess then re-resolves it against the process's current drive. _is_fully_qualified additionally requires a drive (or UNC share) on Windows so the "trusted absolute path" assumption actually holds. - shutil.which short-circuits its PATH search when the name carries a directory component, returning the caller's string after only an isfile+X_OK check - so resolve_binary("/tmp/attacker/evil") passed both CWD guards. The helper is documented as shared and lost its leading underscore in the move, so a name with a path part (separator or Windows drive prefix) is now refused up front. - _detect_via_nvidia_smi caught only (FileNotFoundError, SubprocessError), but spawning a resolved absolute path reaches other OSError variants: PermissionError on a noexec/SELinux-restricted mount, or Errno 8 Exec format error for a +x file that is not a valid executable. Those escaped and aborted comfy install with a traceback instead of degrading to None as the docstring promises; the except now covers OSError, matching hardware._run. Docstrings also stop overclaiming: the CWD guard's one known false positive (running comfy from a directory that is itself an absolute $PATH entry) is now stated rather than asserted away. --- comfy_cli/_safe_exec.py | 92 +++++++++++++++++++++------- comfy_cli/cuda_detect.py | 10 +++- tests/comfy_cli/test_cuda_detect.py | 16 +++++ tests/comfy_cli/test_safe_exec.py | 93 ++++++++++++++++++++++++++--- 4 files changed, 182 insertions(+), 29 deletions(-) diff --git a/comfy_cli/_safe_exec.py b/comfy_cli/_safe_exec.py index 63fd87a7..5826a918 100644 --- a/comfy_cli/_safe_exec.py +++ b/comfy_cli/_safe_exec.py @@ -20,11 +20,43 @@ from __future__ import annotations import logging +import ntpath import os import shutil logger = logging.getLogger(__name__) +_PATH_SEPARATORS = ("/", "\\") + + +def _is_bare_name(name: str) -> bool: + """Return ``True`` if ``name`` is a plain binary name with no path part. + + :func:`shutil.which` short-circuits its ``$PATH`` search when the name holds + a directory component — it looks that path up directly — so a caller-supplied + path would sail through both CWD guards below. Both separators and a Windows + drive prefix (``C:nvidia-smi`` is *drive-relative*) are checked on every + platform, because the caller's string is not necessarily native to the host. + """ + return bool(name) and not any(sep in name for sep in _PATH_SEPARATORS) and not ntpath.splitdrive(name)[0] + + +def _is_fully_qualified(path: str) -> bool: + """Return ``True`` if ``path`` names one unambiguous location. + + :func:`os.path.isabs` is not the same thing as "fully qualified" on Windows: + ``ntpath.isabs(r"\\tools\\nvidia-smi.exe")`` is ``True`` on Python ≤ 3.12, yet + ``CreateProcess`` re-resolves such a drive-less rooted path against the + process's *current drive*. Requiring a drive (or a UNC share) makes the + "trusted absolute path" assumption actually hold. POSIX has no drives, so the + extra requirement applies only where it means something. + """ + if not os.path.isabs(path): + return False + if os.name == "nt": + return bool(os.path.splitdrive(path)[0]) + return True + def is_planted_in_cwd(path: str) -> bool: """Return ``True`` only if ``path`` resolves to a file sitting *directly* in @@ -33,14 +65,21 @@ def is_planted_in_cwd(path: str) -> bool: ``shutil.which`` searches the current directory first on Windows (and on any platform whose ``$PATH`` contains ``.`` or an empty entry), so an attacker who controls the directory the user runs ``comfy`` from can drop a malicious - ``nvidia-smi.exe`` there. Such a plant always lands in the CWD *itself*, so we - reject only a resolved binary whose parent directory **is** the CWD. A - legitimate system binary in a *subdirectory* — e.g. ``System32`` even when the - CWD is ``C:\\Windows``, or a drive root — is left untouched, honouring the - "a legitimate system binary is never rejected" guarantee. Paths are compared - with :func:`os.path.normcase` so Windows' case-insensitivity can't fail the - guard open. Ambiguity — a path on a different drive, or an unresolvable one — - is treated as *not* planted so a legitimate binary is never rejected. + ``nvidia-smi.exe`` there. Those relative-``$PATH``-entry matches come back as + *relative* paths and are rejected by :func:`resolve_binary` directly; this + guard covers the remaining shape, an absolute ``$PATH`` entry that happens to + be the CWD (``PATH="$(pwd):$PATH"`` build wrappers, and Windows' implicit + current-directory search). Only the binary's immediate parent is compared, so + a legitimate system binary in a *subdirectory* — e.g. ``System32`` even when + the CWD is ``C:\\Windows`` — is left untouched. Paths are compared with + :func:`os.path.normcase` so Windows' case-insensitivity can't fail the guard + open. + + A resolution error (an unreadable/deleted CWD, an unresolvable path) means we + cannot prove the binary is *outside* the CWD, so it is reported as planted: + the caller then skips the probe, which is the same degradation as the binary + being absent. Failing open here would be the module's only error path that + hands an unvetted string to :mod:`subprocess`. """ try: cwd = os.path.normcase(os.path.realpath(os.getcwd())) @@ -49,8 +88,8 @@ def is_planted_in_cwd(path: str) -> bool: parent = os.path.normcase(os.path.realpath(os.path.dirname(path))) return parent == cwd except (OSError, ValueError): - # Different drives (Windows) or an unresolvable path → not planted. - return False + logger.debug("cannot place %r relative to the CWD; treating as planted", path, exc_info=True) + return True def resolve_binary(name: str) -> str | None: @@ -63,32 +102,43 @@ def resolve_binary(name: str) -> str | None: directory, so running ``comfy`` from an attacker-controlled directory cannot execute a planted ``nvidia-smi.exe``. + ``name`` must be a bare binary name (see :func:`_is_bare_name`); anything + carrying a path component is refused rather than looked up, because + :func:`shutil.which` would hand such a string straight back. + ``shutil.which`` may itself resolve against the current directory (always on Windows; on any platform when ``$PATH`` holds ``.`` or an empty entry), so as defense-in-depth two CWD-anchored results are additionally rejected on every platform: - * a **relative** result. ``which`` returns ``os.path.join(entry, name)``, so a - relative path means the matching ``$PATH`` entry was itself relative (``.``, - an empty entry, ``subdir``, or Windows' implicitly prepended ``os.curdir``) - and the binary therefore lives under the attacker-controlled CWD. Handing - that string to :mod:`subprocess` would re-resolve it against the CWD — - exactly the hijack this function exists to prevent — so the probe is skipped - instead. A binary found through a normal absolute ``$PATH`` entry always - comes back absolute and is unaffected. + * a result that is not **fully qualified** (see :func:`_is_fully_qualified`). + ``which`` returns ``os.path.join(entry, name)``, so a relative path means the + matching ``$PATH`` entry was itself relative (``.``, an empty entry, + ``subdir``, or Windows' implicitly prepended ``os.curdir``) and the binary + therefore lives under the attacker-controlled CWD. Handing that string to + :mod:`subprocess` would re-resolve it against the CWD — exactly the hijack + this function exists to prevent — so the probe is skipped instead. A binary + found through a normal absolute ``$PATH`` entry always comes back fully + qualified and is unaffected. * an absolute result sitting directly **in** the CWD (see :func:`is_planted_in_cwd`), which covers the CWD appearing in ``$PATH`` as an absolute entry. A legitimate system binary (e.g. ``nvidia-smi.exe`` under ``System32``) is - unaffected by either check. + unaffected by either check. The one known false positive is running ``comfy`` + from a directory that is *itself* an absolute ``$PATH`` entry (``/usr/bin``, + ``C:\\Windows\\System32``): the probe is then skipped and the caller degrades + exactly as it would if the binary were not installed. """ try: + if not _is_bare_name(name): + logger.debug("refusing to resolve %r: not a bare binary name", name) + return None path = shutil.which(name) if path is None: return None - if not os.path.isabs(path): - logger.debug("skipping %r: relative PATH match anchored in CWD (%s)", name, path) + if not _is_fully_qualified(path): + logger.debug("skipping %r: match is not fully qualified (%s)", name, path) return None if is_planted_in_cwd(path): logger.debug("skipping %r: resolved into CWD (%s)", name, path) diff --git a/comfy_cli/cuda_detect.py b/comfy_cli/cuda_detect.py index b04cbdc2..475f446b 100644 --- a/comfy_cli/cuda_detect.py +++ b/comfy_cli/cuda_detect.py @@ -86,6 +86,13 @@ def _detect_via_nvidia_smi() -> tuple[int, int] | None: planted in the current working directory. A binary that is absent, or whose only match is anchored in the CWD, resolves to ``None`` and the probe is skipped — the same degrade-to-``None`` outcome as a failed run. + + Spawning a resolved absolute path surfaces ``OSError`` variants that a bare + name never reached: ``PermissionError`` on a ``noexec``/SELinux-restricted + mount, or ``OSError: [Errno 8] Exec format error`` for a file that is ``+x`` + but not a valid executable. Those are caught alongside + :class:`subprocess.SubprocessError` so the promised degradation to ``None`` + holds instead of aborting ``comfy install`` with a traceback. """ nvidia_smi = _safe_exec.resolve_binary("nvidia-smi") if nvidia_smi is None: @@ -98,7 +105,8 @@ def _detect_via_nvidia_smi() -> tuple[int, int] | None: timeout=10, stderr=subprocess.DEVNULL, ) - except (FileNotFoundError, subprocess.SubprocessError): + except (OSError, subprocess.SubprocessError): + logger.debug("nvidia-smi probe failed", exc_info=True) return None match = re.search(r"CUDA Version:\s*(\d+)\.(\d+)", output) diff --git a/tests/comfy_cli/test_cuda_detect.py b/tests/comfy_cli/test_cuda_detect.py index 5e4eefaf..9fd14b9d 100644 --- a/tests/comfy_cli/test_cuda_detect.py +++ b/tests/comfy_cli/test_cuda_detect.py @@ -111,6 +111,22 @@ def test_timeout(self): ): assert _detect_via_nvidia_smi() is None + @pytest.mark.parametrize( + "error", + [ + PermissionError(13, "Permission denied"), + OSError(8, "Exec format error"), + ], + ids=["noexec_or_selinux_denial", "exec_format_error"], + ) + def test_spawn_oserrors_degrade_to_none(self, error): + """Spawning a *resolved absolute path* reaches ``OSError`` variants a bare + name never did — a probe on a ``noexec`` mount, or a ``+x`` file that is + not a valid executable. Both must degrade to ``None`` rather than abort + ``comfy install`` with a traceback.""" + with patch("comfy_cli.cuda_detect.subprocess.check_output", side_effect=error): + assert _detect_via_nvidia_smi() is None + class TestDetectViaNvidiaSmiResolvesBinaryPath: """``_detect_via_nvidia_smi`` must spawn a resolved absolute path, never the diff --git a/tests/comfy_cli/test_safe_exec.py b/tests/comfy_cli/test_safe_exec.py index 83100de0..a1eb5772 100644 --- a/tests/comfy_cli/test_safe_exec.py +++ b/tests/comfy_cli/test_safe_exec.py @@ -1,15 +1,18 @@ """Tests for :mod:`comfy_cli._safe_exec`. -The contract under test: ``resolve_binary`` hands back only a *trusted absolute -path*, and returns ``None`` — "skip this probe" — for every CWD-anchored match, -while never rejecting a legitimate system binary and never raising. +The contract under test: ``resolve_binary`` hands back only a *fully qualified, +unambiguous* path, and returns ``None`` — "skip this probe" — for every +CWD-anchored match and every resolution it cannot vet, without ever raising. """ from __future__ import annotations +import ntpath import os from unittest.mock import patch +import pytest + from comfy_cli import _safe_exec @@ -122,8 +125,84 @@ def test_broken_lookup_returns_none(self): with patch.object(_safe_exec.shutil, "which", side_effect=RuntimeError("boom")): assert _safe_exec.resolve_binary("nvidia-smi") is None - def test_is_planted_in_cwd_swallows_path_errors(self): - """An unresolvable path (e.g. a different Windows drive) is treated as - *not* planted so a legitimate binary is never rejected.""" + def test_is_planted_in_cwd_fails_closed_on_path_errors(self): + """A CWD that cannot be read (deleted out from under the process) leaves + us unable to prove the binary sits *outside* it, so it is reported as + planted and the caller skips the probe rather than spawning an unvetted + path.""" with patch.object(_safe_exec.os, "getcwd", side_effect=OSError("gone")): - assert _safe_exec.is_planted_in_cwd("/usr/bin/nvidia-smi") is False + assert _safe_exec.is_planted_in_cwd("/usr/bin/nvidia-smi") is True + + def test_resolve_binary_skips_probe_when_cwd_unreadable(self): + with ( + patch.object(_safe_exec.shutil, "which", return_value=os.path.join(os.sep, "usr", "bin", "nvidia-smi")), + patch.object(_safe_exec.os, "getcwd", side_effect=OSError("gone")), + ): + assert _safe_exec.resolve_binary("nvidia-smi") is None + + +class TestResolveBinaryRejectsNonBareNames: + """``shutil.which`` looks a name with a directory component up *directly* + instead of searching ``$PATH``, handing the caller's own string back after a + bare ``isfile``+``X_OK`` check — which would sail past both CWD guards. Such a + name is refused before the lookup.""" + + @pytest.mark.parametrize( + "name", + [ + "/tmp/attacker/evil", + "attacker/evil", + r"C:\attacker\evil.exe", + r"..\evil.exe", + "C:evil.exe", # drive-relative: no separator, still not a bare name + "", + ], + ) + def test_rejects_name_with_path_component(self, name): + with patch.object(_safe_exec.shutil, "which") as which: + assert _safe_exec.resolve_binary(name) is None + which.assert_not_called() + + def test_accepts_bare_name(self, tmp_path): + legit = tmp_path / "nvidia-smi" + legit.write_text("") + with ( + patch.object(_safe_exec.os, "getcwd", return_value=str(tmp_path / "elsewhere")), + patch.object(_safe_exec.shutil, "which", return_value=str(legit)), + ): + assert _safe_exec.resolve_binary("nvidia-smi") == str(legit) + + +class TestResolveBinaryRequiresFullyQualifiedMatch: + """``ntpath.isabs`` accepts a drive-less rooted path, which ``CreateProcess`` + re-resolves against the process's *current drive* — so "absolute" alone is not + enough to call a Windows match trusted.""" + + def test_drive_less_rooted_windows_path_is_not_fully_qualified(self): + """``ntpath.isabs`` accepts this shape on the 3.10–3.12 interpreters this + package still supports (3.13 tightened it), so the guard must not lean on + ``isabs`` alone — it has to reject the path on every version.""" + with ( + patch.object(_safe_exec.os, "name", "nt"), + patch.object(_safe_exec.os, "path", ntpath), + ): + assert _safe_exec._is_fully_qualified(r"\tools\nvidia-smi.exe") is False + + @pytest.mark.parametrize("path", [r"C:\Windows\System32\nvidia-smi.exe", r"\\host\share\nvidia-smi.exe"]) + def test_drive_and_unc_paths_are_fully_qualified(self, path): + with ( + patch.object(_safe_exec.os, "name", "nt"), + patch.object(_safe_exec.os, "path", ntpath), + ): + assert _safe_exec._is_fully_qualified(path) is True + + def test_posix_absolute_path_is_fully_qualified(self): + assert _safe_exec._is_fully_qualified(os.path.join(os.sep, "usr", "bin", "nvidia-smi")) is True + + def test_resolve_binary_skips_drive_less_windows_match(self): + with ( + patch.object(_safe_exec.os, "name", "nt"), + patch.object(_safe_exec.os, "path", ntpath), + patch.object(_safe_exec.shutil, "which", return_value=r"\tools\nvidia-smi.exe"), + ): + assert _safe_exec.resolve_binary("nvidia-smi") is None