rewrite FreeBSD pids() implementationopenbsd-fixes

1 files changed, 25 insertions, 64 deletions

diff --git a/psutil/arch/freebsd/specific.c b/psutil/arch/freebsd/specific.c
index 90ea81e8..18fa51fa 100644
--- a/psutil/arch/freebsd/specific.c
+++ b/psutil/arch/freebsd/specific.c
@@ -88,81 +88,42 @@ psutil_get_proc_list(struct kinfo_proc **procList, size_t *procCount) {
     // Returns a list of all BSD processes on the system.  This routine
     // allocates the list and puts it in *procList and a count of the
     // number of entries in *procCount.  You are responsible for freeing
-    // this list (use "free" from System framework).
-    // On success, the function returns 0.
-    // On error, the function returns a BSD errno value.
+    // this list. On success returns 0, else 1 with exception set.
     int err;
-    struct kinfo_proc *result;
-    int done;
+    struct kinfo_proc *buf = NULL;
     int name[] = { CTL_KERN, KERN_PROC, KERN_PROC_PROC, 0 };
-    size_t length;
+    size_t length = 0;
 
-    assert( procList != NULL);
+    assert(procList != NULL);
     assert(*procList == NULL);
     assert(procCount != NULL);
 
-    *procCount = 0;
-
-    /*
-     * We start by calling sysctl with result == NULL and length == 0.
-     * That will succeed, and set length to the appropriate length.
-     * We then allocate a buffer of that size and call sysctl again
-     * with that buffer.  If that succeeds, we're done.  If that fails
-     * with ENOMEM, we have to throw away our buffer and loop.  Note
-     * that the loop causes use to call sysctl with NULL again; this
-     * is necessary because the ENOMEM failure case sets length to
-     * the amount of data returned, not the amount of data that
-     * could have been returned.
-     */
-    result = NULL;
-    done = 0;
-    do {
-        assert(result == NULL);
-        // Call sysctl with a NULL buffer.
-        length = 0;
-        err = sysctl((int *)name, (sizeof(name) / sizeof(*name)) - 1,
-                     NULL, &length, NULL, 0);
-        if (err == -1)
-            err = errno;
-
-        // Allocate an appropriately sized buffer based on the results
-        // from the previous call.
-        if (err == 0) {
-            result = malloc(length);
-            if (result == NULL)
-                err = ENOMEM;
-        }
+    // Call sysctl with a NULL buffer in order to get buffer length.
+    err = sysctl(name, 3, NULL, &length, NULL, 0);
+    if (err == -1) {
+        PyErr_SetFromOSErrnoWithSyscall("sysctl (null buffer)");
+        return 1;
+    }
 
-        // Call sysctl again with the new buffer.  If we get an ENOMEM
-        // error, toss away our buffer and start again.
-        if (err == 0) {
-            err = sysctl((int *) name, (sizeof(name) / sizeof(*name)) - 1,
-                         result, &length, NULL, 0);
-            if (err == -1)
-                err = errno;
-            if (err == 0) {
-                done = 1;
-            }
-            else if (err == ENOMEM) {
-                assert(result != NULL);
-                free(result);
-                result = NULL;
-                err = 0;
-            }
-        }
-    } while (err == 0 && ! done);
+    // Allocate an appropriately sized buffer based on the results
+    // from the previous call.
+    buf = malloc(length);
+    if (buf == NULL) {
+        PyErr_NoMemory();
+        return 1;
+    }
 
-    // Clean up and establish post conditions.
-    if (err != 0 && result != NULL) {
-        free(result);
-        result = NULL;
+    // Call sysctl again with the new buffer.
+    err = sysctl(name, 3, buf, &length, NULL, 0);
+    if (err == -1) {
+        PyErr_SetFromOSErrnoWithSyscall("sysctl");
+        free(buf);
+        return 1;
     }
 
-    *procList = result;
+    *procList = buf;
     *procCount = length / sizeof(struct kinfo_proc);
-
-    assert((err == 0) == (*procList != NULL));
-    return err;
+    return 0;
 }