Harden Windows Arm64 utility discovery

Author: isVoid

cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py

@@ -91,17 +91,30 @@ def _windows_installed_nsight_root(product: str) -> str | None:
     access = winreg.KEY_READ | winreg.KEY_WOW64_64KEY
     product_key_path = rf"{_NSIGHT_REGISTRY_ROOT}\{product}"
     try:
-        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, product_key_path, 0, access) as product_key:
-            current_version, _ = winreg.QueryValueEx(product_key, "CurrentVersion")
-            if not isinstance(current_version, str) or not current_version:
-                raise RuntimeError(f"Invalid CurrentVersion value in {product_key_path!r}")
-            with winreg.OpenKey(product_key, current_version, 0, access) as version_key:
-                install_root, _ = winreg.QueryValueEx(version_key, None)
+        product_context = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, product_key_path, 0, access)
     except FileNotFoundError:
         return None
 
-    if not isinstance(install_root, str) or not install_root:
-        raise RuntimeError(f"Invalid installation directory for {product_key_path!r} version {current_version!r}")
+    try:
+        with product_context as product_key:
+            current_version, _ = winreg.QueryValueEx(product_key, "CurrentVersion")
+            if not isinstance(current_version, str) or not current_version.strip():
+                raise RuntimeError(
+                    f"Invalid CurrentVersion value {current_version!r} in "
+                    f"Nsight {product!r} registry registration at {product_key_path!r}"
+                )
+            with winreg.OpenKey(product_key, current_version, 0, access) as version_key:
+                install_root, _ = winreg.QueryValueEx(version_key, None)
+    except FileNotFoundError as exc:
+        raise RuntimeError(
+            f"Incomplete Nsight {product!r} registry registration at {product_key_path!r}"
+        ) from exc
+
+    if not isinstance(install_root, str) or not install_root.strip():
+        raise RuntimeError(
+            f"Invalid installation directory {install_root!r} in Nsight {product!r} "
+            f"registry registration at {product_key_path!r} version {current_version!r}"
+        )
     return install_root
 
 
@@ -182,6 +195,9 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None:
     Raises:
         UnsupportedBinaryError: If ``utility_name`` is not in the supported set
             (see ``SUPPORTED_BINARY_UTILITIES``).
+        RuntimeError: If a native Windows architecture needed for an
+            architecture-specific utility layout cannot be determined, or an
+            installed Nsight product has incomplete or invalid registry data.
 
     Windows on ARM (WoA) Note:
         Binary utilities execute in separate processes and do not need to match

cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py

@@ -11,6 +11,8 @@
     "arm64": 0xAA64,
 }
 
+_WINDOWS_ARCH_BY_PE_MACHINE = {machine: arch for arch, machine in WINDOWS_PE_MACHINE_BY_ARCH.items()}
+
 
 class UnsupportedArchError(RuntimeError):
     """Raised when Python reports an unsupported Windows architecture."""
@@ -36,8 +38,8 @@ def windows_python_arch() -> str:
     raise UnsupportedArchError(raw_platform_tag)
 
 
-def windows_machine_arch() -> str:
-    """Return the native Windows machine architecture."""
+def _windows_machine_arch_from_platform() -> str:
+    """Return the Windows architecture reported by Python's platform module."""
     raw_machine = platform.machine()
     machine = raw_machine.lower().replace("_", "-")
 
@@ -50,6 +52,62 @@ def windows_machine_arch() -> str:
     raise RuntimeError(f"Unsupported Windows machine architecture: {raw_machine!r}")
 
 
+def _windows_native_machine() -> int | None:
+    """Return the native Windows PE machine type, or None on older Windows."""
+    import ctypes
+    from ctypes import wintypes
+
+    try:
+        # These ctypes attributes are absent from the type stubs on non-Windows hosts.
+        kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)  # type: ignore[attr-defined]
+    except OSError as exc:
+        raise RuntimeError("Failed to load kernel32 while detecting the native Windows architecture") from exc
+
+    get_current_process = kernel32.GetCurrentProcess
+    try:
+        is_wow64_process2 = kernel32.IsWow64Process2
+    except AttributeError:
+        return None
+
+    get_current_process.argtypes = ()
+    get_current_process.restype = wintypes.HANDLE
+    is_wow64_process2.argtypes = (
+        wintypes.HANDLE,
+        ctypes.POINTER(wintypes.USHORT),
+        ctypes.POINTER(wintypes.USHORT),
+    )
+    is_wow64_process2.restype = wintypes.BOOL
+
+    process_machine = wintypes.USHORT()
+    native_machine = wintypes.USHORT()
+    if not is_wow64_process2(
+        get_current_process(),
+        ctypes.byref(process_machine),
+        ctypes.byref(native_machine),
+    ):
+        error_code = ctypes.get_last_error()  # type: ignore[attr-defined]
+        error = ctypes.WinError(error_code)  # type: ignore[attr-defined]
+        raise RuntimeError(
+            f"IsWow64Process2 failed while detecting the native Windows architecture "
+            f"(Windows error {error_code}): {error}"
+        ) from error
+    return native_machine.value
+
+
+def windows_machine_arch() -> str:
+    """Return the native Windows machine architecture, ignoring process emulation."""
+    native_machine = _windows_native_machine()
+    if native_machine is None:
+        # IsWow64Process2 predates x64-on-Arm emulation, so this fallback is only
+        # needed on older Windows versions where platform.machine() is sufficient.
+        return _windows_machine_arch_from_platform()
+
+    try:
+        return _WINDOWS_ARCH_BY_PE_MACHINE[native_machine]
+    except KeyError:
+        raise RuntimeError(f"Unsupported native Windows PE machine type: 0x{native_machine:04x}") from None
+
+
 def windows_pe_matches_arch(path: str, target_arch: str) -> bool:
     """Return whether a Windows Portable Executable (PE) targets the requested architecture.
 

cuda_pathfinder/docs/source/install.rst

@@ -9,7 +9,7 @@ Runtime Requirements
 
 ``cuda.pathfinder`` is a pure-Python package with no runtime dependencies:
 
-* Linux (x86-64, arm64) and Windows (x86-64)
+* Linux (x86-64, arm64) and Windows (x86-64, arm64)
 * Python 3.10 - 3.14
 
 Installing from PyPI

cuda_pathfinder/tests/test_find_nvidia_binaries.py

@@ -58,6 +58,15 @@ def fake_is_executable_candidate(path):
     return checked
 
 
+def _patch_winreg(mocker):
+    winreg = mocker.MagicMock()
+    winreg.HKEY_LOCAL_MACHINE = object()
+    winreg.KEY_READ = 0x20019
+    winreg.KEY_WOW64_64KEY = 0x0100
+    mocker.patch.object(binary_finder_module.importlib, "import_module", return_value=winreg)
+    return winreg
+
+
 @pytest.mark.usefixtures("clear_find_binary_cache")
 def test_find_binary_search_path_includes_site_packages_conda_cuda(monkeypatch, mocker):
     conda_prefix = os.path.join(os.sep, "conda")
@@ -425,13 +434,9 @@ def test_windows_installed_nsight_root_reads_64_bit_registry(mocker):
     product_context.__enter__.return_value = product_key
     version_context = mocker.MagicMock()
     version_context.__enter__.return_value = version_key
-    winreg = mocker.MagicMock()
-    winreg.HKEY_LOCAL_MACHINE = object()
-    winreg.KEY_READ = 0x20019
-    winreg.KEY_WOW64_64KEY = 0x0100
+    winreg = _patch_winreg(mocker)
     winreg.OpenKey.side_effect = (product_context, version_context)
     winreg.QueryValueEx.side_effect = (("2026.1.3", 1), (install_root, 1))
-    mocker.patch.object(binary_finder_module.importlib, "import_module", return_value=winreg)
 
     assert binary_finder_module._windows_installed_nsight_root("Systems") == install_root
     access = winreg.KEY_READ | winreg.KEY_WOW64_64KEY
@@ -448,6 +453,99 @@ def test_windows_installed_nsight_root_reads_64_bit_registry(mocker):
     )
 
 
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_windows_installed_nsight_root_returns_none_when_product_key_is_absent(mocker):
+    winreg = _patch_winreg(mocker)
+    winreg.OpenKey.side_effect = FileNotFoundError("Nsight Systems is not installed")
+
+    assert binary_finder_module._windows_installed_nsight_root("Systems") is None
+
+
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_windows_installed_nsight_root_rejects_missing_current_version(mocker):
+    product_context = mocker.MagicMock()
+    product_context.__enter__.return_value = mocker.MagicMock()
+    winreg = _patch_winreg(mocker)
+    winreg.OpenKey.return_value = product_context
+    winreg.QueryValueEx.side_effect = FileNotFoundError("CurrentVersion is missing")
+
+    with pytest.raises(RuntimeError, match=r"Incomplete Nsight 'Systems' registry registration") as exc_info:
+        binary_finder_module._windows_installed_nsight_root("Systems")
+
+    assert isinstance(exc_info.value.__cause__, FileNotFoundError)
+
+
+@pytest.mark.parametrize("current_version", (None, "", "   ", 2026))
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_windows_installed_nsight_root_rejects_invalid_current_version(mocker, current_version):
+    product_context = mocker.MagicMock()
+    product_context.__enter__.return_value = mocker.MagicMock()
+    winreg = _patch_winreg(mocker)
+    winreg.OpenKey.return_value = product_context
+    winreg.QueryValueEx.return_value = (current_version, 1)
+
+    with pytest.raises(RuntimeError, match=r"Invalid CurrentVersion value .*Nsight 'Systems' registry registration"):
+        binary_finder_module._windows_installed_nsight_root("Systems")
+
+
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_windows_installed_nsight_root_rejects_missing_version_key(mocker):
+    product_key = mocker.MagicMock()
+    product_context = mocker.MagicMock()
+    product_context.__enter__.return_value = product_key
+    winreg = _patch_winreg(mocker)
+    winreg.OpenKey.side_effect = (product_context, FileNotFoundError("Version key is missing"))
+    winreg.QueryValueEx.return_value = ("2026.1.3", 1)
+
+    with pytest.raises(RuntimeError, match=r"Incomplete Nsight 'Systems' registry registration") as exc_info:
+        binary_finder_module._windows_installed_nsight_root("Systems")
+
+    assert isinstance(exc_info.value.__cause__, FileNotFoundError)
+
+
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_windows_installed_nsight_root_rejects_missing_installation_directory(mocker):
+    product_context = mocker.MagicMock()
+    product_context.__enter__.return_value = mocker.MagicMock()
+    version_context = mocker.MagicMock()
+    version_context.__enter__.return_value = mocker.MagicMock()
+    winreg = _patch_winreg(mocker)
+    winreg.OpenKey.side_effect = (product_context, version_context)
+    winreg.QueryValueEx.side_effect = (("2026.1.3", 1), FileNotFoundError("Installation directory is missing"))
+
+    with pytest.raises(RuntimeError, match=r"Incomplete Nsight 'Systems' registry registration") as exc_info:
+        binary_finder_module._windows_installed_nsight_root("Systems")
+
+    assert isinstance(exc_info.value.__cause__, FileNotFoundError)
+
+
+@pytest.mark.parametrize("install_root", (None, "", "   ", 2026))
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_windows_installed_nsight_root_rejects_invalid_installation_directory(mocker, install_root):
+    product_context = mocker.MagicMock()
+    product_context.__enter__.return_value = mocker.MagicMock()
+    version_context = mocker.MagicMock()
+    version_context.__enter__.return_value = mocker.MagicMock()
+    winreg = _patch_winreg(mocker)
+    winreg.OpenKey.side_effect = (product_context, version_context)
+    winreg.QueryValueEx.side_effect = (("2026.1.3", 1), (install_root, 1))
+
+    with pytest.raises(
+        RuntimeError,
+        match=r"Invalid installation directory .*Nsight 'Systems' registry registration.*version '2026.1.3'",
+    ):
+        binary_finder_module._windows_installed_nsight_root("Systems")
+
+
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_windows_installed_nsight_root_propagates_access_errors(mocker):
+    winreg = _patch_winreg(mocker)
+    winreg.OpenKey.side_effect = PermissionError("Registry access denied")
+
+    with pytest.raises(PermissionError, match="Registry access denied"):
+        binary_finder_module._windows_installed_nsight_root("Systems")
+
+
 @pytest.mark.usefixtures("clear_find_binary_cache")
 def test_find_binary_first_matching_dir_wins(monkeypatch, mocker):
     conda_prefix = os.path.join(os.sep, "conda")

cuda_pathfinder/tests/test_search_steps.py

@@ -5,7 +5,9 @@
 
 from __future__ import annotations
 
+import ctypes
 import os
+from ctypes import wintypes
 
 import pytest
 
@@ -148,22 +150,82 @@ def test_rejects_unknown_sysconfig_tag(self, mocker):
 
 
 class TestWindowsMachineArch:
+    @pytest.mark.parametrize(
+        ("native_machine", "expected"),
+        ((0x8664, "x64"), (0xAA64, "arm64")),
+    )
+    @pytest.mark.agent_authored(model="gpt-5.6")
+    def test_uses_native_pe_machine(self, mocker, native_machine, expected):
+        mocker.patch.object(windows_arch_mod, "_windows_native_machine", return_value=native_machine)
+        platform_machine = mocker.patch.object(windows_arch_mod.platform, "machine", return_value="AMD64")
+
+        assert windows_arch_mod.windows_machine_arch() == expected
+        platform_machine.assert_not_called()
+
+    @pytest.mark.agent_authored(model="gpt-5.6")
+    def test_rejects_unknown_native_pe_machine(self, mocker):
+        mocker.patch.object(windows_arch_mod, "_windows_native_machine", return_value=0x014C)
+
+        with pytest.raises(RuntimeError, match=r"Unsupported native Windows PE machine type: 0x014c"):
+            windows_arch_mod.windows_machine_arch()
+
     @pytest.mark.parametrize(
         ("reported_machine", "expected"),
         (("AMD64", "x64"), ("x86_64", "x64"), ("ARM64", "arm64"), ("aarch64", "arm64")),
     )
     @pytest.mark.agent_authored(model="gpt-5.6")
-    def test_normalizes_platform_machine(self, mocker, reported_machine, expected):
+    def test_old_windows_fallback_normalizes_platform_machine(self, mocker, reported_machine, expected):
+        mocker.patch.object(windows_arch_mod, "_windows_native_machine", return_value=None)
         mocker.patch.object(windows_arch_mod.platform, "machine", return_value=reported_machine)
 
         assert windows_arch_mod.windows_machine_arch() == expected
 
     @pytest.mark.agent_authored(model="gpt-5.6")
-    def test_rejects_unknown_machine(self, mocker):
-        mocker.patch.object(windows_arch_mod.platform, "machine", return_value="mips64")
+    def test_native_machine_returns_none_when_is_wow64_process2_is_unavailable(self, mocker):
+        kernel32 = mocker.Mock(spec=["GetCurrentProcess"])
+        mocker.patch.object(ctypes, "WinDLL", create=True, return_value=kernel32)
 
-        with pytest.raises(RuntimeError, match=r"Unsupported Windows machine architecture: 'mips64'"):
-            windows_arch_mod.windows_machine_arch()
+        assert windows_arch_mod._windows_native_machine() is None
+
+    @pytest.mark.agent_authored(model="gpt-5.6")
+    def test_native_machine_configures_api_and_returns_native_machine(self, mocker):
+        kernel32 = mocker.Mock()
+        kernel32.GetCurrentProcess.return_value = wintypes.HANDLE(1)
+
+        def report_native_machine(_process, _process_machine, native_machine):
+            native_machine._obj.value = 0xAA64
+            return True
+
+        kernel32.IsWow64Process2.side_effect = report_native_machine
+        mocker.patch.object(ctypes, "WinDLL", create=True, return_value=kernel32)
+
+        assert windows_arch_mod._windows_native_machine() == 0xAA64
+        assert kernel32.GetCurrentProcess.argtypes == ()
+        assert kernel32.GetCurrentProcess.restype is wintypes.HANDLE
+        assert kernel32.IsWow64Process2.argtypes == (
+            wintypes.HANDLE,
+            ctypes.POINTER(wintypes.USHORT),
+            ctypes.POINTER(wintypes.USHORT),
+        )
+        assert kernel32.IsWow64Process2.restype is wintypes.BOOL
+
+    @pytest.mark.agent_authored(model="gpt-5.6")
+    def test_native_machine_raises_contextual_error_when_api_call_fails(self, mocker):
+        kernel32 = mocker.Mock()
+        kernel32.GetCurrentProcess.return_value = wintypes.HANDLE(1)
+        kernel32.IsWow64Process2.return_value = False
+        mocker.patch.object(ctypes, "WinDLL", create=True, return_value=kernel32)
+        mocker.patch.object(ctypes, "get_last_error", create=True, return_value=87)
+        windows_error = OSError(87, "The parameter is incorrect")
+        mocker.patch.object(ctypes, "WinError", create=True, return_value=windows_error)
+
+        with pytest.raises(
+            RuntimeError,
+            match=r"IsWow64Process2 failed while detecting the native Windows architecture \(Windows error 87\)",
+        ) as exc_info:
+            windows_arch_mod._windows_native_machine()
+
+        assert exc_info.value.__cause__ is windows_error
 
 
 @pytest.mark.parametrize(