[macOS] dynamic set buffer size for process connections/fds (fixes #1901) (#1903)

5 files changed, 102 insertions, 111 deletions

diff --git a/HISTORY.rst b/HISTORY.rst
index 55c6b2b9..a8284ad2 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -24,6 +24,11 @@ XXXX-XX-XX
   where it first finds one
 - 1874_: [Solaris] swap output error due to incorrect range.
 - 1892_: [macOS] psutil.cpu_freq() broken on Apple M1.
+- 1901_: [macOS] different functions, especially process' open_files() and
+  connections() methods, could randomly raise AccessDenied because the internal
+  buffer of `proc_pidinfo(PROC_PIDLISTFDS)` syscall was not big enough. We now
+  dynamically increase the buffer size until it's big enough instead of giving
+  up and raising AccessDenied, which was a fallback to avoid crashing.
 - 1904_: [Windows] OpenProcess fails with ERROR_SUCCESS due to GetLastError()
   called after sprintf().  (patch by alxchk)
 - 1913_: [Linux] wait_procs seemingly ignoring timeout, TimeoutExpired thrown
diff --git a/psutil/_psosx.py b/psutil/_psosx.py
index d948cc15..eda70d21 100644
--- a/psutil/_psosx.py
+++ b/psutil/_psosx.py
@@ -4,7 +4,6 @@
 
 """macOS platform implementation."""
 
-import contextlib
 import errno
 import functools
 import os
@@ -354,32 +353,6 @@ def wrap_exceptions(fun):
     return wrapper
 
 
-@contextlib.contextmanager
-def catch_zombie(proc):
-    """There are some poor C APIs which incorrectly raise ESRCH when
-    the process is still alive or it's a zombie, or even RuntimeError
-    (those who don't set errno). This is here in order to solve:
-    https://github.com/giampaolo/psutil/issues/1044
-    """
-    try:
-        yield
-    except (OSError, RuntimeError) as err:
-        if isinstance(err, RuntimeError) or err.errno == errno.ESRCH:
-            try:
-                # status() is not supposed to lie and correctly detect
-                # zombies so if it raises ESRCH it's true.
-                status = proc.status()
-            except NoSuchProcess:
-                raise err
-            else:
-                if status == _common.STATUS_ZOMBIE:
-                    raise ZombieProcess(proc.pid, proc._name, proc._ppid)
-                else:
-                    raise AccessDenied(proc.pid, proc._name)
-        else:
-            raise
-
-
 class Process(object):
     """Wrapper class around underlying C implementation."""
 
@@ -402,8 +375,7 @@ class Process(object):
     @memoize_when_activated
     def _get_pidtaskinfo(self):
         # Note: should work for PIDs owned by user only.
-        with catch_zombie(self):
-            ret = cext.proc_pidtaskinfo_oneshot(self.pid)
+        ret = cext.proc_pidtaskinfo_oneshot(self.pid)
         assert len(ret) == len(pidtaskinfo_map)
         return ret
 
@@ -422,18 +394,15 @@ class Process(object):
 
     @wrap_exceptions
     def exe(self):
-        with catch_zombie(self):
-            return cext.proc_exe(self.pid)
+        return cext.proc_exe(self.pid)
 
     @wrap_exceptions
     def cmdline(self):
-        with catch_zombie(self):
-            return cext.proc_cmdline(self.pid)
+        return cext.proc_cmdline(self.pid)
 
     @wrap_exceptions
     def environ(self):
-        with catch_zombie(self):
-            return parse_environ_block(cext.proc_environ(self.pid))
+        return parse_environ_block(cext.proc_environ(self.pid))
 
     @wrap_exceptions
     def ppid(self):
@@ -442,8 +411,7 @@ class Process(object):
 
     @wrap_exceptions
     def cwd(self):
-        with catch_zombie(self):
-            return cext.proc_cwd(self.pid)
+        return cext.proc_cwd(self.pid)
 
     @wrap_exceptions
     def uids(self):
@@ -516,8 +484,7 @@ class Process(object):
         if self.pid == 0:
             return []
         files = []
-        with catch_zombie(self):
-            rawlist = cext.proc_open_files(self.pid)
+        rawlist = cext.proc_open_files(self.pid)
         for path, fd in rawlist:
             if isfile_strict(path):
                 ntuple = _common.popenfile(path, fd)
@@ -530,8 +497,7 @@ class Process(object):
             raise ValueError("invalid %r kind argument; choose between %s"
                              % (kind, ', '.join([repr(x) for x in conn_tmap])))
         families, types = conn_tmap[kind]
-        with catch_zombie(self):
-            rawlist = cext.proc_connections(self.pid, families, types)
+        rawlist = cext.proc_connections(self.pid, families, types)
         ret = []
         for item in rawlist:
             fd, fam, type, laddr, raddr, status = item
@@ -544,8 +510,7 @@ class Process(object):
     def num_fds(self):
         if self.pid == 0:
             return 0
-        with catch_zombie(self):
-            return cext.proc_num_fds(self.pid)
+        return cext.proc_num_fds(self.pid)
 
     @wrap_exceptions
     def wait(self, timeout=None):
@@ -553,13 +518,11 @@ class Process(object):
 
     @wrap_exceptions
     def nice_get(self):
-        with catch_zombie(self):
-            return cext_posix.getpriority(self.pid)
+        return cext_posix.getpriority(self.pid)
 
     @wrap_exceptions
     def nice_set(self, value):
-        with catch_zombie(self):
-            return cext_posix.setpriority(self.pid, value)
+        return cext_posix.setpriority(self.pid, value)
 
     @wrap_exceptions
     def status(self):
diff --git a/psutil/_psutil_osx.c b/psutil/_psutil_osx.c
index 5a77de14..6f824eb2 100644
--- a/psutil/_psutil_osx.c
+++ b/psutil/_psutil_osx.c
@@ -109,6 +109,72 @@ psutil_task_for_pid(pid_t pid, mach_port_t *task)
 
 
 /*
+ * A wrapper around proc_pidinfo(PROC_PIDLISTFDS), which dynamically sets
+ * the buffer size.
+ */
+static struct proc_fdinfo*
+psutil_proc_list_fds(pid_t pid, int *num_fds) {
+    int ret;
+    int fds_size = 0;
+    int max_size = 24 * 1024 * 1024;  // 24M
+    struct proc_fdinfo *fds_pointer = NULL;
+
+    errno = 0;
+    ret = proc_pidinfo(pid, PROC_PIDLISTFDS, 0, NULL, 0);
+    if (ret <= 0) {
+        psutil_raise_for_pid(pid, "proc_pidinfo(PROC_PIDLISTFDS) 1/2");
+        goto error;
+    }
+
+    while (1) {
+        if (ret > fds_size) {
+            while (ret > fds_size) {
+                fds_size += PROC_PIDLISTFD_SIZE * 32;
+                if (fds_size > max_size) {
+                    PyErr_Format(PyExc_RuntimeError,
+                                 "prevent malloc() to allocate > 24M");
+                    goto error;
+                }
+            }
+
+            if (fds_pointer != NULL) {
+                free(fds_pointer);
+            }
+            fds_pointer = malloc(fds_size);
+
+            if (fds_pointer == NULL) {
+                PyErr_NoMemory();
+                goto error;
+            }
+        }
+
+        errno = 0;
+        ret = proc_pidinfo(pid, PROC_PIDLISTFDS, 0, fds_pointer, fds_size);
+        if (ret <= 0) {
+            psutil_raise_for_pid(pid, "proc_pidinfo(PROC_PIDLISTFDS) 2/2");
+            goto error;
+        }
+
+        if (ret + (int)PROC_PIDLISTFD_SIZE >= fds_size) {
+            psutil_debug("PROC_PIDLISTFDS: make room for 1 extra fd");
+            ret = fds_size + (int)PROC_PIDLISTFD_SIZE;
+            continue;
+        }
+
+        break;
+    }
+
+    *num_fds = (ret / (int)PROC_PIDLISTFD_SIZE);
+    return fds_pointer;
+
+error:
+    if (fds_pointer != NULL)
+        free(fds_pointer);
+    return NULL;
+}
+
+
+/*
  * Return a Python list of all the PIDs running on the system.
  */
 static PyObject *
@@ -791,7 +857,7 @@ psutil_proc_threads(PyObject *self, PyObject *args) {
     if (err != KERN_SUCCESS) {
         // errcode 4 is "invalid argument" (access denied)
         if (err == 4) {
-            AccessDenied("task_info");
+            AccessDenied("task_info(TASK_BASIC_INFO)");
         }
         else {
             // otherwise throw a runtime error with appropriate error code
@@ -866,15 +932,12 @@ error:
 static PyObject *
 psutil_proc_open_files(PyObject *self, PyObject *args) {
     pid_t pid;
-    int pidinfo_result;
-    int iterations;
+    int num_fds;
     int i;
     unsigned long nb;
-
     struct proc_fdinfo *fds_pointer = NULL;
     struct proc_fdinfo *fdp_pointer;
     struct vnode_fdinfowithpath vi;
-
     PyObject *py_retlist = PyList_New(0);
     PyObject *py_tuple = NULL;
     PyObject *py_path = NULL;
@@ -885,23 +948,11 @@ psutil_proc_open_files(PyObject *self, PyObject *args) {
     if (! PyArg_ParseTuple(args, _Py_PARSE_PID, &pid))
         goto error;
 
-    pidinfo_result = psutil_proc_pidinfo(pid, PROC_PIDLISTFDS, 0, NULL, 0);
-    if (pidinfo_result <= 0)
-        goto error;
-
-    fds_pointer = malloc(pidinfo_result);
-    if (fds_pointer == NULL) {
-        PyErr_NoMemory();
-        goto error;
-    }
-    pidinfo_result = psutil_proc_pidinfo(
-        pid, PROC_PIDLISTFDS, 0, fds_pointer, pidinfo_result);
-    if (pidinfo_result <= 0)
+    fds_pointer = psutil_proc_list_fds(pid, &num_fds);
+    if (fds_pointer == NULL)
         goto error;
 
-    iterations = (pidinfo_result / PROC_PIDLISTFD_SIZE);
-
-    for (i = 0; i < iterations; i++) {
+    for (i = 0; i < num_fds; i++) {
         fdp_pointer = &fds_pointer[i];
 
         if (fdp_pointer->proc_fdtype == PROX_FDTYPE_VNODE) {
@@ -968,15 +1019,12 @@ error:
 static PyObject *
 psutil_proc_connections(PyObject *self, PyObject *args) {
     pid_t pid;
-    int pidinfo_result;
-    int iterations;
+    int num_fds;
     int i;
     unsigned long nb;
-
     struct proc_fdinfo *fds_pointer = NULL;
     struct proc_fdinfo *fdp_pointer;
     struct socket_fdinfo si;
-
     PyObject *py_retlist = PyList_New(0);
     PyObject *py_tuple = NULL;
     PyObject *py_laddr = NULL;
@@ -997,25 +1045,11 @@ psutil_proc_connections(PyObject *self, PyObject *args) {
         goto error;
     }
 
-    if (pid == 0)
-        return py_retlist;
-    pidinfo_result = psutil_proc_pidinfo(pid, PROC_PIDLISTFDS, 0, NULL, 0);
-    if (pidinfo_result <= 0)
-        goto error;
-
-    fds_pointer = malloc(pidinfo_result);
-    if (fds_pointer == NULL) {
-        PyErr_NoMemory();
-        goto error;
-    }
-
-    pidinfo_result = psutil_proc_pidinfo(
-        pid, PROC_PIDLISTFDS, 0, fds_pointer, pidinfo_result);
-    if (pidinfo_result <= 0)
+    fds_pointer = psutil_proc_list_fds(pid, &num_fds);
+    if (fds_pointer == NULL)
         goto error;
 
-    iterations = (pidinfo_result / PROC_PIDLISTFD_SIZE);
-    for (i = 0; i < iterations; i++) {
+    for (i = 0; i < num_fds; i++) {
         py_tuple = NULL;
         py_laddr = NULL;
         py_raddr = NULL;
@@ -1176,30 +1210,18 @@ error:
 static PyObject *
 psutil_proc_num_fds(PyObject *self, PyObject *args) {
     pid_t pid;
-    int pidinfo_result;
-    int num;
+    int num_fds;
     struct proc_fdinfo *fds_pointer;
 
     if (! PyArg_ParseTuple(args, _Py_PARSE_PID, &pid))
         return NULL;
 
-    pidinfo_result = proc_pidinfo(pid, PROC_PIDLISTFDS, 0, NULL, 0);
-    if (pidinfo_result <= 0)
-        return PyErr_SetFromErrno(PyExc_OSError);
-
-    fds_pointer = malloc(pidinfo_result);
+    fds_pointer = psutil_proc_list_fds(pid, &num_fds);
     if (fds_pointer == NULL)
-        return PyErr_NoMemory();
-    pidinfo_result = proc_pidinfo(pid, PROC_PIDLISTFDS, 0, fds_pointer,
-                                  pidinfo_result);
-    if (pidinfo_result <= 0) {
-        free(fds_pointer);
-        return PyErr_SetFromErrno(PyExc_OSError);
-    }
+        return NULL;
 
-    num = (pidinfo_result / PROC_PIDLISTFD_SIZE);
     free(fds_pointer);
-    return Py_BuildValue("i", num);
+    return Py_BuildValue("i", num_fds);
 }
 
 
diff --git a/psutil/_psutil_posix.c b/psutil/_psutil_posix.c
index 305cec76..7c0d548c 100644
--- a/psutil/_psutil_posix.c
+++ b/psutil/_psutil_posix.c
@@ -147,7 +147,7 @@ psutil_pid_exists(pid_t pid) {
  */
 void
 psutil_raise_for_pid(long pid, char *syscall) {
-    if (errno != 0)  // unlikely
+    if (errno != 0)
         PyErr_SetFromOSErrnoWithSyscall(syscall);
     else if (psutil_pid_exists(pid) == 0)
         NoSuchProcess(syscall);
diff --git a/psutil/arch/osx/process_info.c b/psutil/arch/osx/process_info.c
index 15ad3b89..6ab42750 100644
--- a/psutil/arch/osx/process_info.c
+++ b/psutil/arch/osx/process_info.c
@@ -361,7 +361,7 @@ psutil_get_kinfo_proc(pid_t pid, struct kinfo_proc *kp) {
 
     // sysctl succeeds but len is zero, happens when process has gone away
     if (len == 0) {
-        NoSuchProcess("sysctl (len == 0)");
+        NoSuchProcess("sysctl(kinfo_proc), len == 0");
         return -1;
     }
     return 0;
@@ -370,8 +370,8 @@ psutil_get_kinfo_proc(pid_t pid, struct kinfo_proc *kp) {
 
 /*
  * A wrapper around proc_pidinfo().
- * https://opensource.apple.com/source/xnu/xnu-2050.7.9/bsd/kern/proc_info.c.
- * Returns 0 on failure (and Python exception gets already set).
+ * https://opensource.apple.com/source/xnu/xnu-2050.7.9/bsd/kern/proc_info.c
+ * Returns 0 on failure.
  */
 int
 psutil_proc_pidinfo(pid_t pid, int flavor, uint64_t arg, void *pti, int size) {
@@ -380,11 +380,12 @@ psutil_proc_pidinfo(pid_t pid, int flavor, uint64_t arg, void *pti, int size) {
 
     ret = proc_pidinfo(pid, flavor, arg, pti, size);
     if (ret <= 0) {
-        psutil_raise_for_pid(pid, "proc_pidinfo() failed");
+        psutil_raise_for_pid(pid, "proc_pidinfo()");
         return 0;
     }
-    else if ((unsigned long )ret < sizeof(pti)) {
-        psutil_raise_for_pid(pid, "proc_pidinfo() len mismatch");
+    if ((unsigned long)ret < sizeof(pti)) {
+        psutil_raise_for_pid(
+            pid, "proc_pidinfo() return size < sizeof(struct_pointer)");
         return 0;
     }
     return ret;