fix(cuda): resolve nvidia-smi via absolute path to block Windows CWD planting (BE-5357)

comfy_cli/_safe_exec.py

@@ -0,0 +1,149 @@
+"""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 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
+    ``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. 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()))
+        # ``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):
+        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:
+    """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``.
+
+    ``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 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. 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 _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)
+            return None
+        return path
+    except Exception:
+        logger.debug("resolving binary %r failed", name, exc_info=True)
+        return None

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,15 +79,34 @@ 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.
+
+    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:
+        return None
+
     try:
         output = subprocess.check_output(
-            ["nvidia-smi"],
+            [nvidia_smi],
             text=True,
             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)

comfy_cli/hardware.py

@@ -20,7 +20,7 @@
 
 import psutil
 
-from comfy_cli import cuda_detect
+from comfy_cli import _safe_exec, cuda_detect
 
 logger = logging.getLogger(__name__)
 
@@ -30,11 +30,21 @@
 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:`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 = _safe_exec.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,

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"
@@ -97,6 +111,82 @@ 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
+    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):