[Windows] use NtQuerySystemInformation to determine process exe() (#1677)

10 files changed, 104 insertions, 121 deletions

diff --git a/HISTORY.rst b/HISTORY.rst
index 848d3559..7bb4fffa 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -14,6 +14,8 @@ XXXX-XX-XX
   Minimum supported Windows version now is Windows Vista.
 - 1667_: added process_iter(new_only=True) parameter.
 - 1671_: [FreeBSD] add CI testing/service for FreeBSD (Cirrus CI).
+- 1677_: [Windows] process exe() will succeed for all process PIDs (instead of
+  raising AccessDenied).
 
 **Bug fixes**
 
diff --git a/docs/index.rst b/docs/index.rst
index b1c4ab74..74dddd23 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1098,9 +1098,9 @@ Process class
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+--------------------------+
     | :meth:`gids`                 |                               | :meth:`name`                 | :meth:`num_ctx_switches`     | :meth:`terminal`         | :meth:`terminal`         |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+--------------------------+
-    | :meth:`num_ctx_switches`     |                               | :meth:`ppid`                 | :meth:`ppid`                 |                          |                          |
+    | :meth:`num_ctx_switches`     | :meth:`exe`                   | :meth:`ppid`                 | :meth:`ppid`                 |                          |                          |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+--------------------------+
-    | :meth:`num_threads`          |                               | :meth:`status`               | :meth:`status`               | :meth:`gids`             | :meth:`gids`             |
+    | :meth:`num_threads`          | :meth:`name`                  | :meth:`status`               | :meth:`status`               | :meth:`gids`             | :meth:`gids`             |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+--------------------------+
     | :meth:`uids`                 |                               | :meth:`terminal`             | :meth:`terminal`             | :meth:`uids`             | :meth:`uids`             |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+--------------------------+
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index cabfe934..e63cf97f 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -446,74 +446,74 @@ psutil_proc_environ(PyObject *self, PyObject *args) {
 
 
 /*
- * Return process executable path.
+ * Return process executable path. Works for all processes regardless of
+ * privilege. NtQuerySystemInformation has some sort of internal cache,
+ * since it succeeds even when a process is gone (but not if a PID never
+ * existed).
  */
 static PyObject *
 psutil_proc_exe(PyObject *self, PyObject *args) {
     DWORD pid;
-    HANDLE hProcess;
-    wchar_t exe[MAX_PATH];
-    unsigned int size = sizeof(exe);
+    NTSTATUS status;
+    PVOID buffer;
+    ULONG bufferSize = 0x100;
+    SYSTEM_PROCESS_ID_INFORMATION processIdInfo;
+    PyObject *py_exe;
 
     if (! PyArg_ParseTuple(args, _Py_PARSE_PID, &pid))
         return NULL;
 
-    hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
-    if (NULL == hProcess)
-        return NULL;
+    if (pid == 0)
+        return AccessDenied("forced for PID 0");
 
-    memset(exe, 0, MAX_PATH);
-    if (QueryFullProcessImageNameW(hProcess, 0, exe, &size) == 0) {
-        // https://github.com/giampaolo/psutil/issues/1662
-        if (GetLastError() == 0)
-            AccessDenied("QueryFullProcessImageNameW (forced EPERM)");
+    buffer = MALLOC_ZERO(bufferSize);
+    if (! buffer)
+        return PyErr_NoMemory();
+    processIdInfo.ProcessId = (HANDLE)(ULONG_PTR)pid;
+    processIdInfo.ImageName.Length = 0;
+    processIdInfo.ImageName.MaximumLength = (USHORT)bufferSize;
+    processIdInfo.ImageName.Buffer = buffer;
+
+    status = NtQuerySystemInformation(
+        SystemProcessIdInformation,
+        &processIdInfo,
+        sizeof(SYSTEM_PROCESS_ID_INFORMATION),
+        NULL);
+
+    if (status == STATUS_INFO_LENGTH_MISMATCH) {
+        // Required length is stored in MaximumLength.
+        FREE(buffer);
+        buffer = MALLOC_ZERO(processIdInfo.ImageName.MaximumLength);
+        if (! buffer)
+            return PyErr_NoMemory();
+        processIdInfo.ImageName.Buffer = buffer;
+
+        status = NtQuerySystemInformation(
+            SystemProcessIdInformation,
+            &processIdInfo,
+            sizeof(SYSTEM_PROCESS_ID_INFORMATION),
+            NULL);
+    }
+
+    if (! NT_SUCCESS(status)) {
+        FREE(buffer);
+        if (psutil_pid_is_running(pid) == 0)
+            NoSuchProcess("NtQuerySystemInformation");
         else
-            PyErr_SetFromOSErrnoWithSyscall("QueryFullProcessImageNameW");
-        CloseHandle(hProcess);
+            psutil_SetFromNTStatusErr(status, "NtQuerySystemInformation");
         return NULL;
     }
-    CloseHandle(hProcess);
-    return PyUnicode_FromWideChar(exe, wcslen(exe));
-}
-
 
-/*
- * Return process base name.
- * Note: psutil_proc_exe() is attempted first because it's faster
- * but it raise AccessDenied for processes owned by other users
- * in which case we fall back on using this.
- */
-static PyObject *
-psutil_proc_name(PyObject *self, PyObject *args) {
-    DWORD pid;
-    int ok;
-    PROCESSENTRY32W pentry;
-    HANDLE hSnapShot;
-
-    if (! PyArg_ParseTuple(args, _Py_PARSE_PID, &pid))
-        return NULL;
-    hSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, pid);
-    if (hSnapShot == INVALID_HANDLE_VALUE)
-        return PyErr_SetFromOSErrnoWithSyscall("CreateToolhelp32Snapshot");
-    pentry.dwSize = sizeof(PROCESSENTRY32W);
-    ok = Process32FirstW(hSnapShot, &pentry);
-    if (! ok) {
-        PyErr_SetFromOSErrnoWithSyscall("Process32FirstW");
-        CloseHandle(hSnapShot);
-        return NULL;
+    if (processIdInfo.ImageName.Buffer == NULL) {
+        // Happens for PID 4.
+        py_exe = Py_BuildValue("s", "");
     }
-    while (ok) {
-        if (pentry.th32ProcessID == pid) {
-            CloseHandle(hSnapShot);
-            return PyUnicode_FromWideChar(
-                pentry.szExeFile, wcslen(pentry.szExeFile));
-        }
-        ok = Process32NextW(hSnapShot, &pentry);
+    else {
+        py_exe = PyUnicode_FromWideChar(processIdInfo.ImageName.Buffer,
+                                        processIdInfo.ImageName.Length / 2);
     }
-
-    CloseHandle(hSnapShot);
-    NoSuchProcess("CreateToolhelp32Snapshot loop (no PID found)");
-    return NULL;
+    FREE(buffer);
+    return py_exe;
 }
 
 
@@ -587,6 +587,10 @@ psutil_GetProcWsetInformation(
 
     bufferSize = 0x8000;
     buffer = MALLOC_ZERO(bufferSize);
+    if (! buffer) {
+        PyErr_NoMemory();
+        return 1;
+    }
 
     while ((status = NtQueryVirtualMemory(
             hProcess,
@@ -605,6 +609,10 @@ psutil_GetProcWsetInformation(
             return 1;
         }
         buffer = MALLOC_ZERO(bufferSize);
+        if (! buffer) {
+            PyErr_NoMemory();
+            return 1;
+        }
     }
 
     if (!NT_SUCCESS(status)) {
@@ -1602,8 +1610,6 @@ PsutilMethods[] = {
      "Return process environment data"},
     {"proc_exe", psutil_proc_exe, METH_VARARGS,
      "Return path of the process executable"},
-    {"proc_name", psutil_proc_name, METH_VARARGS,
-     "Return process name"},
     {"proc_kill", psutil_proc_kill, METH_VARARGS,
      "Kill the process identified by the given PID"},
     {"proc_cpu_times", psutil_proc_cpu_times, METH_VARARGS,
diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py
index 3e14bc3c..83793c5a 100644
--- a/psutil/_pswindows.py
+++ b/psutil/_pswindows.py
@@ -721,9 +721,11 @@ class Process(object):
 
     def oneshot_enter(self):
         self.oneshot_info.cache_activate(self)
+        self.exe.cache_activate(self)
 
     def oneshot_exit(self):
         self.oneshot_info.cache_deactivate(self)
+        self.exe.cache_deactivate(self)
 
     @wrap_exceptions
     @memoize_when_activated
@@ -735,7 +737,6 @@ class Process(object):
         assert len(ret) == len(pinfo_map)
         return ret
 
-    @wrap_exceptions
     def name(self):
         """Return process name, which on Windows is always the final
         part of the executable.
@@ -744,20 +745,19 @@ class Process(object):
         # and process-hacker.
         if self.pid == 0:
             return "System Idle Process"
-        elif self.pid == 4:
+        if self.pid == 4:
             return "System"
-        else:
-            try:
-                # Note: this will fail with AD for most PIDs owned
-                # by another user but it's faster.
-                return py2_strencode(os.path.basename(self.exe()))
-            except AccessDenied:
-                return py2_strencode(cext.proc_name(self.pid))
+        return os.path.basename(self.exe())
 
     @wrap_exceptions
+    @memoize_when_activated
     def exe(self):
         exe = cext.proc_exe(self.pid)
-        return py2_strencode(exe)
+        if not PY3:
+            exe = py2_strencode(exe)
+        if exe.startswith('\\'):
+            return convert_dos_path(exe)
+        return exe  # May be "Registry", "MemCompression", ...
 
     @wrap_exceptions
     @retry_error_partial_copy
@@ -843,7 +843,6 @@ class Process(object):
             for addr, perm, path, rss in raw:
                 path = convert_dos_path(path)
                 if not PY3:
-                    assert isinstance(path, unicode), type(path)
                     path = py2_strencode(path)
                 addr = hex(addr)
                 yield (addr, perm, path, rss)
diff --git a/psutil/arch/windows/ntextapi.h b/psutil/arch/windows/ntextapi.h
index b7b1c976..8cb00430 100644
--- a/psutil/arch/windows/ntextapi.h
+++ b/psutil/arch/windows/ntextapi.h
@@ -33,6 +33,8 @@ typedef LONG NTSTATUS;
 #define ProcessIoPriority 33
 #undef  ProcessWow64Information
 #define ProcessWow64Information 26
+#undef  SystemProcessIdInformation
+#define SystemProcessIdInformation 88
 
 // process suspend() / resume()
 typedef enum _KTHREAD_STATE {
@@ -362,6 +364,12 @@ typedef struct _PSUTIL_PROCESS_WS_COUNTERS {
     SIZE_T NumberOfShareablePages;
 } PSUTIL_PROCESS_WS_COUNTERS, *PPSUTIL_PROCESS_WS_COUNTERS;
 
+// exe()
+typedef struct _SYSTEM_PROCESS_ID_INFORMATION {
+    HANDLE ProcessId;
+    UNICODE_STRING ImageName;
+} SYSTEM_PROCESS_ID_INFORMATION, *PSYSTEM_PROCESS_ID_INFORMATION;
+
 // ====================================================================
 // PEB structs for cmdline(), cwd(), environ()
 // ====================================================================
diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index 0c297b9b..5251983f 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -206,9 +206,6 @@ def _get_py_exe():
         return exe
     else:
         exe = os.path.realpath(sys.executable)
-        if WINDOWS:
-            # avoid subprocess warnings
-            exe = exe.replace('\\', '\\\\')
         assert os.path.exists(exe), exe
         return exe
 
@@ -366,7 +363,7 @@ def create_proc_children_pair():
         s += "f.write(str(os.getpid()));"
         s += "f.close();"
         s += "time.sleep(60);"
-        p = subprocess.Popen(['%s', '-c', s])
+        p = subprocess.Popen([r'%s', '-c', s])
         p.wait()
         """ % (_TESTFN2, PYTHON_EXE))
     # On Windows if we create a subprocess with CREATE_NO_WINDOW flag
diff --git a/psutil/tests/test_contracts.py b/psutil/tests/test_contracts.py
index bdf055d0..5e258b30 100755
--- a/psutil/tests/test_contracts.py
+++ b/psutil/tests/test_contracts.py
@@ -431,6 +431,8 @@ class TestFetchAllProcesses(unittest.TestCase):
         if not ret:
             self.assertEqual(ret, '')
         else:
+            if WINDOWS and not ret.endswith('.exe'):
+                return  # May be "Registry", "MemCompression", ...
             assert os.path.isabs(ret), ret
             # Note: os.stat() may return False even if the file is there
             # hence we skip the test, see:
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index 0b54f5b1..5728f183 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -1325,6 +1325,16 @@ class TestProcess(unittest.TestCase):
             except NotImplementedError:
                 pass
             else:
+                # NtQuerySystemInformation succeeds if process is gone.
+                if WINDOWS and name in ('exe', 'name'):
+                    normcase = os.path.normcase
+                    if name == 'exe':
+                        self.assertEqual(normcase(ret), normcase(PYTHON_EXE))
+                    else:
+                        self.assertEqual(
+                            normcase(ret),
+                            normcase(os.path.basename(PYTHON_EXE)))
+                    continue
                 self.fail(
                     "NoSuchProcess exception not raised for %r, retval=%s" % (
                         name, ret))
diff --git a/psutil/tests/test_unicode.py b/psutil/tests/test_unicode.py
index 6b1d720e..91395ebf 100755..100644
--- a/psutil/tests/test_unicode.py
+++ b/psutil/tests/test_unicode.py
@@ -61,7 +61,6 @@ from psutil import BSD
 from psutil import MACOS
 from psutil import OPENBSD
 from psutil import POSIX
-from psutil import WINDOWS
 from psutil._compat import PY3
 from psutil._compat import u
 from psutil.tests import APPVEYOR
@@ -75,7 +74,6 @@ from psutil.tests import get_test_subprocess
 from psutil.tests import HAS_CONNECTIONS_UNIX
 from psutil.tests import HAS_ENVIRON
 from psutil.tests import HAS_MEMORY_MAPS
-from psutil.tests import mock
 from psutil.tests import PYPY
 from psutil.tests import reap_children
 from psutil.tests import safe_mkdir
@@ -171,16 +169,7 @@ class _BaseFSAPIsTests(object):
 
     def test_proc_name(self):
         subp = get_test_subprocess(cmd=[self.funky_name])
-        if WINDOWS:
-            # On Windows name() is determined from exe() first, because
-            # it's faster; we want to overcome the internal optimization
-            # and test name() instead of exe().
-            with mock.patch("psutil._psplatform.cext.proc_exe",
-                            side_effect=psutil.AccessDenied(os.getpid())) as m:
-                name = psutil.Process(subp.pid).name()
-                assert m.called
-        else:
-            name = psutil.Process(subp.pid).name()
+        name = psutil.Process(subp.pid).name()
         self.assertIsInstance(name, str)
         if self.expect_exact_path_match():
             self.assertEqual(name, os.path.basename(self.funky_name))
@@ -321,19 +310,6 @@ class TestFSAPIsWithInvalidPath(_BaseFSAPIsTests, unittest.TestCase):
         return True
 
 
-@unittest.skipIf(not WINDOWS, "WINDOWS only")
-class TestWinProcessName(unittest.TestCase):
-
-    def test_name_type(self):
-        # On Windows name() is determined from exe() first, because
-        # it's faster; we want to overcome the internal optimization
-        # and test name() instead of exe().
-        with mock.patch("psutil._psplatform.cext.proc_exe",
-                        side_effect=psutil.AccessDenied(os.getpid())) as m:
-            self.assertIsInstance(psutil.Process().name(), str)
-            assert m.called
-
-
 # ===================================================================
 # Non fs APIs
 # ===================================================================
diff --git a/psutil/tests/test_windows.py b/psutil/tests/test_windows.py
index de41aed8..6d4e8599 100755
--- a/psutil/tests/test_windows.py
+++ b/psutil/tests/test_windows.py
@@ -326,20 +326,6 @@ class TestProcess(unittest.TestCase):
         p = psutil.Process(self.pid)
         self.assertRaises(ValueError, p.send_signal, signal.SIGINT)
 
-    def test_exe_and_name(self):
-        for p in psutil.process_iter():
-            # On Windows name() is never supposed to raise AccessDenied,
-            # see https://github.com/giampaolo/psutil/issues/627
-            try:
-                name = p.name()
-            except psutil.NoSuchProcess:
-                pass
-            else:
-                try:
-                    self.assertEqual(os.path.basename(p.exe()), name)
-                except psutil.Error:
-                    continue
-
     def test_num_handles_increment(self):
         p = psutil.Process(os.getpid())
         before = p.num_handles()
@@ -524,6 +510,13 @@ class TestProcess(unittest.TestCase):
                 self.assertRaises(psutil.AccessDenied, p.cwd)
         self.assertGreaterEqual(m.call_count, 5)
 
+    def test_exe(self):
+        # NtQuerySystemInformation succeeds if process is gone. Make sure
+        # it raises NSP for a non existent pid.
+        pid = psutil.pids()[-1] + 99999
+        proc = psutil._psplatform.Process(pid)
+        self.assertRaises(psutil.NoSuchProcess, proc.exe)
+
 
 @unittest.skipIf(not WINDOWS, "WINDOWS only")
 class TestProcessWMI(unittest.TestCase):
@@ -610,16 +603,6 @@ class TestDualProcessImplementation(unittest.TestCase):
     @classmethod
     def tearDownClass(cls):
         reap_children()
-    # ---
-    # same tests as above but mimicks the AccessDenied failure of
-    # the first (fast) method failing with AD.
-
-    def test_name(self):
-        name = psutil.Process(self.pid).name()
-        with mock.patch("psutil._psplatform.cext.proc_exe",
-                        side_effect=psutil.AccessDenied(os.getpid())) as fun:
-            self.assertEqual(psutil.Process(self.pid).name(), name)
-            assert fun.called
 
     def test_memory_info(self):
         mem_1 = psutil.Process(self.pid).memory_info()