fix(hardware): resolve probe binaries via absolute path to block Windows CWD planting (BE-3434)

comfy_cli/hardware.py

@@ -16,6 +16,7 @@
 import logging
 import os
 import platform
+import shutil
 import subprocess
 
 import psutil
@@ -27,14 +28,96 @@
 _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.
 
-    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. 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
     try:
         output = subprocess.check_output(
-            cmd,
+            [resolved, *cmd[1:]],
             text=True,
             timeout=_SUBPROCESS_TIMEOUT,
             stderr=subprocess.DEVNULL,

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
 
@@ -351,3 +352,165 @@ 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
+
+    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 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):
+        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()