diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py index 9d1e1a56..8ff0e61b 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,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, diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py index c6bcbb55..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 @@ -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()