diff --git a/comfy_cli/_safe_exec.py b/comfy_cli/_safe_exec.py new file mode 100644 index 00000000..5826a918 --- /dev/null +++ b/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 diff --git a/comfy_cli/cuda_detect.py b/comfy_cli/cuda_detect.py index f6339f36..475f446b 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,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) diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py index 9d1e1a56..572eb2d4 100644 --- a/comfy_cli/hardware.py +++ b/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, diff --git a/tests/comfy_cli/test_cuda_detect.py b/tests/comfy_cli/test_cuda_detect.py index 2a1cf310..9fd14b9d 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" @@ -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): diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py index c6bcbb55..ddc7a9a2 100644 --- a/tests/comfy_cli/test_hardware.py +++ b/tests/comfy_cli/test_hardware.py @@ -10,12 +10,13 @@ import ctypes import json +import os from pathlib import Path from unittest.mock import patch 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 @@ -351,3 +352,59 @@ 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(_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"]) + + 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(_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 + 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(_safe_exec.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() + + 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(_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 + mock_run.assert_not_called() diff --git a/tests/comfy_cli/test_safe_exec.py b/tests/comfy_cli/test_safe_exec.py new file mode 100644 index 00000000..a1eb5772 --- /dev/null +++ b/tests/comfy_cli/test_safe_exec.py @@ -0,0 +1,208 @@ +"""Tests for :mod:`comfy_cli._safe_exec`. + +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 + + +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_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 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