OpenBSD: rewrite net_connections() from scratch (#2230)

4 files changed, 189 insertions, 224 deletions

diff --git a/psutil/arch/openbsd/proc.c b/psutil/arch/openbsd/proc.c
index 35acab25..344d6010 100644
--- a/psutil/arch/openbsd/proc.c
+++ b/psutil/arch/openbsd/proc.c
@@ -6,25 +6,12 @@
  */
 
 #include <Python.h>
-#include <assert.h>
-#include <errno.h>
-#include <stdlib.h>
-#include <stdio.h>
-#include <string.h>
 #include <fcntl.h>
 #include <sys/types.h>
 #include <sys/param.h>
 #include <sys/sysctl.h>
-#include <sys/user.h>
 #include <sys/proc.h>
-#include <signal.h>
 #include <kvm.h>
-#include <netdb.h>  // for NI_MAXHOST
-#include <sys/socket.h>
-#define _KERNEL  // for DTYPE_*
-#include <sys/file.h>
-#undef _KERNEL
-#include <arpa/inet.h> // for inet_ntoa()
 
 #include "../../_psutil_common.h"
 #include "../../_psutil_posix.h"
@@ -155,6 +142,7 @@ psutil_get_proc_list(struct kinfo_proc **procList, size_t *procCount) {
 }
 
 
+// TODO: refactor this (it's clunky)
 static char **
 _psutil_get_argv(pid_t pid) {
     static char **argv;
@@ -322,212 +310,3 @@ psutil_proc_cwd(PyObject *self, PyObject *args) {
     }
     return PyUnicode_DecodeFSDefault(path);
 }
-
-
-// see sys/kern/kern_sysctl.c lines 1100 and
-// usr.bin/fstat/fstat.c print_inet_details()
-static char *
-psutil_convert_ipv4(int family, uint32_t addr[4]) {
-    struct in_addr a;
-    memcpy(&a, addr, sizeof(a));
-    return inet_ntoa(a);
-}
-
-
-static char *
-psutil_inet6_addrstr(struct in6_addr *p)
-{
-    struct sockaddr_in6 sin6;
-    static char hbuf[NI_MAXHOST];
-    const int niflags = NI_NUMERICHOST;
-
-    memset(&sin6, 0, sizeof(sin6));
-    sin6.sin6_family = AF_INET6;
-    sin6.sin6_len = sizeof(struct sockaddr_in6);
-    sin6.sin6_addr = *p;
-    if (IN6_IS_ADDR_LINKLOCAL(p) &&
-        *(u_int16_t *)&sin6.sin6_addr.s6_addr[2] != 0) {
-        sin6.sin6_scope_id =
-            ntohs(*(u_int16_t *)&sin6.sin6_addr.s6_addr[2]);
-        sin6.sin6_addr.s6_addr[2] = sin6.sin6_addr.s6_addr[3] = 0;
-    }
-
-    if (getnameinfo((struct sockaddr *)&sin6, sin6.sin6_len,
-        hbuf, sizeof(hbuf), NULL, 0, niflags))
-        return "invalid";
-
-    return hbuf;
-}
-
-
-/*
- * List process connections.
- * Note: there is no net_connections() on OpenBSD. The Python
- * implementation will iterate over all processes and use this
- * function.
- * Note: local and remote paths cannot be determined for UNIX sockets.
- */
-PyObject *
-psutil_proc_connections(PyObject *self, PyObject *args) {
-    pid_t pid;
-    int i;
-    int cnt;
-    struct kinfo_file *freep = NULL;
-    struct kinfo_file *kif;
-    char *tcplist = NULL;
-    PyObject *py_retlist = PyList_New(0);
-    PyObject *py_tuple = NULL;
-    PyObject *py_laddr = NULL;
-    PyObject *py_raddr = NULL;
-    PyObject *py_af_filter = NULL;
-    PyObject *py_type_filter = NULL;
-    PyObject *py_family = NULL;
-    PyObject *_type = NULL;
-
-    if (py_retlist == NULL)
-        return NULL;
-    if (! PyArg_ParseTuple(args, _Py_PARSE_PID "OO", &pid, &py_af_filter,
-                           &py_type_filter))
-        goto error;
-    if (!PySequence_Check(py_af_filter) || !PySequence_Check(py_type_filter)) {
-        PyErr_SetString(PyExc_TypeError, "arg 2 or 3 is not a sequence");
-        goto error;
-    }
-
-    freep = kinfo_getfile(pid, &cnt);
-    if (freep == NULL) {
-        goto error;
-    }
-
-    for (i = 0; i < cnt; i++) {
-        int state;
-        int lport;
-        int rport;
-        char addrbuf[NI_MAXHOST + 2];
-        int inseq;
-        struct in6_addr laddr6;
-        py_tuple = NULL;
-        py_laddr = NULL;
-        py_raddr = NULL;
-
-        kif = &freep[i];
-        if (kif->f_type == DTYPE_SOCKET) {
-            // apply filters
-            py_family = PyLong_FromLong((long)kif->so_family);
-            inseq = PySequence_Contains(py_af_filter, py_family);
-            Py_DECREF(py_family);
-            if (inseq == 0)
-                continue;
-            _type = PyLong_FromLong((long)kif->so_type);
-            inseq = PySequence_Contains(py_type_filter, _type);
-            Py_DECREF(_type);
-            if (inseq == 0)
-                continue;
-
-            // IPv4 / IPv6 socket
-            if ((kif->so_family == AF_INET) || (kif->so_family == AF_INET6)) {
-                // fill status
-                if (kif->so_type == SOCK_STREAM)
-                    state = kif->t_state;
-                else
-                    state = PSUTIL_CONN_NONE;
-
-                // ports
-                lport = ntohs(kif->inp_lport);
-                rport = ntohs(kif->inp_fport);
-
-                // local address, IPv4
-                if (kif->so_family == AF_INET) {
-                    py_laddr = Py_BuildValue(
-                        "(si)",
-                        psutil_convert_ipv4(kif->so_family, kif->inp_laddru),
-                        lport);
-                    if (!py_laddr)
-                        goto error;
-                }
-                else {
-                    // local address, IPv6
-                    memcpy(&laddr6, kif->inp_laddru, sizeof(laddr6));
-                    snprintf(addrbuf, sizeof(addrbuf), "%s",
-                             psutil_inet6_addrstr(&laddr6));
-                    py_laddr = Py_BuildValue("(si)", addrbuf, lport);
-                    if (!py_laddr)
-                        goto error;
-                }
-
-                if (rport != 0) {
-                    // remote address, IPv4
-                    if (kif->so_family == AF_INET) {
-                        py_raddr = Py_BuildValue(
-                            "(si)",
-                            psutil_convert_ipv4(
-                                kif->so_family, kif->inp_faddru),
-                            rport);
-                    }
-                    else {
-                        // remote address, IPv6
-                        memcpy(&laddr6, kif->inp_faddru, sizeof(laddr6));
-                        snprintf(addrbuf, sizeof(addrbuf), "%s",
-                                 psutil_inet6_addrstr(&laddr6));
-                        py_raddr = Py_BuildValue("(si)", addrbuf, rport);
-                        if (!py_raddr)
-                            goto error;
-                    }
-                }
-                else {
-                    py_raddr = Py_BuildValue("()");
-                }
-
-                if (!py_raddr)
-                    goto error;
-                py_tuple = Py_BuildValue(
-                    "(iiiNNi)",
-                    kif->fd_fd,
-                    kif->so_family,
-                    kif->so_type,
-                    py_laddr,
-                    py_raddr,
-                    state);
-                if (!py_tuple)
-                    goto error;
-                if (PyList_Append(py_retlist, py_tuple))
-                    goto error;
-                Py_DECREF(py_tuple);
-            }
-            // UNIX socket.
-            // XXX: local addr is supposed to be in "unp_path" but it
-            // always empty; also "fstat" command is not able to show
-            // UNIX socket paths.
-            else if (kif->so_family == AF_UNIX) {
-                py_tuple = Py_BuildValue(
-                    "(iiissi)",
-                    kif->fd_fd,
-                    kif->so_family,
-                    kif->so_type,
-                    "",  // laddr (kif->unp_path is empty)
-                    "",  // raddr
-                    PSUTIL_CONN_NONE);
-                if (!py_tuple)
-                    goto error;
-                if (PyList_Append(py_retlist, py_tuple))
-                    goto error;
-                Py_DECREF(py_tuple);
-                Py_INCREF(Py_None);
-            }
-        }
-    }
-    free(freep);
-    free(tcplist);
-    return py_retlist;
-
-error:
-    Py_XDECREF(py_tuple);
-    Py_XDECREF(py_laddr);
-    Py_XDECREF(py_raddr);
-    Py_DECREF(py_retlist);
-    if (freep != NULL)
-        free(freep);
-    if (tcplist != NULL)
-        free(tcplist);
-    return NULL;
-}
diff --git a/psutil/arch/openbsd/proc.h b/psutil/arch/openbsd/proc.h
index 5ca5890d..747507dd 100644
--- a/psutil/arch/openbsd/proc.h
+++ b/psutil/arch/openbsd/proc.h
@@ -18,5 +18,3 @@ PyObject *psutil_get_cmdline(pid_t pid);
 PyObject *psutil_proc_threads(PyObject *self, PyObject *args);
 PyObject *psutil_proc_num_fds(PyObject *self, PyObject *args);
 PyObject *psutil_proc_cwd(PyObject *self, PyObject *args);
-PyObject *psutil_proc_connections(PyObject *self, PyObject *args);
-
diff --git a/psutil/arch/openbsd/socks.c b/psutil/arch/openbsd/socks.c
new file mode 100644
index 00000000..69daa447
--- /dev/null
+++ b/psutil/arch/openbsd/socks.c
@@ -0,0 +1,180 @@
+/*
+ * Copyright (c) 2009, Giampaolo Rodola'.
+ * All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include <Python.h>
+#include <sys/sysctl.h>
+#include <sys/socket.h>
+#include <kvm.h>
+#define _KERNEL  // silence compiler warning
+#include <sys/file.h>  // DTYPE_SOCKET
+#include <netdb.h>  // INET6_ADDRSTRLEN, in6_addr
+#undef _KERNEL
+
+#include "../../_psutil_common.h"
+#include "../../_psutil_posix.h"
+
+
+PyObject *
+psutil_net_connections(PyObject *self, PyObject *args) {
+    pid_t pid;
+    int i;
+    int cnt;
+    int state;
+    int lport;
+    int rport;
+    char lip[INET6_ADDRSTRLEN];
+    char rip[INET6_ADDRSTRLEN];
+    int inseq;
+
+    char errbuf[_POSIX2_LINE_MAX];
+    kvm_t *kd = NULL;
+
+    struct kinfo_file *kif;
+    struct kinfo_file *ikf;
+    struct in6_addr laddr6;
+
+    PyObject *py_retlist = PyList_New(0);
+    PyObject *py_tuple = NULL;
+    PyObject *py_laddr = NULL;
+    PyObject *py_raddr = NULL;
+    PyObject *py_lpath = NULL;
+    PyObject *py_af_filter = NULL;
+    PyObject *py_type_filter = NULL;
+    PyObject *py_family = NULL;
+    PyObject *_type = NULL;
+
+
+    if (py_retlist == NULL)
+        return NULL;
+    if (! PyArg_ParseTuple(args, _Py_PARSE_PID "OO", &pid, &py_af_filter,
+                           &py_type_filter)) {
+        goto error;
+    }
+    if (!PySequence_Check(py_af_filter) || !PySequence_Check(py_type_filter)) {
+        PyErr_SetString(PyExc_TypeError, "arg 2 or 3 is not a sequence");
+        goto error;
+    }
+
+    kd = kvm_openfiles(NULL, NULL, NULL, KVM_NO_FILES, errbuf);
+    if (! kd) {
+        convert_kvm_err("kvm_openfiles", errbuf);
+        goto error;
+    }
+
+    ikf = kvm_getfiles(kd, KERN_FILE_BYPID, -1, sizeof(*ikf), &cnt);
+    if (! ikf) {
+        PyErr_SetFromOSErrnoWithSyscall("kvm_getfiles");
+        goto error;
+    }
+
+    for (int i = 0; i < cnt; i++) {
+        const struct kinfo_file *kif = ikf + i;
+        py_tuple = NULL;
+        py_laddr = NULL;
+        py_raddr = NULL;
+        py_lpath = NULL;
+
+        // apply filters
+        if (kif->f_type != DTYPE_SOCKET)
+            continue;
+        if (pid != -1 && kif->p_pid != (uint32_t)pid)
+            continue;
+        py_family = PyLong_FromLong((long)kif->so_family);
+        inseq = PySequence_Contains(py_af_filter, py_family);
+        Py_DECREF(py_family);
+        if (inseq == 0)
+            continue;
+        _type = PyLong_FromLong((long)kif->so_type);
+        inseq = PySequence_Contains(py_type_filter, _type);
+        Py_DECREF(_type);
+        if (inseq == 0)
+            continue;
+
+        // IPv4 / IPv6 socket
+        if ((kif->so_family == AF_INET) || (kif->so_family == AF_INET6)) {
+            // status
+            if (kif->so_type == SOCK_STREAM)
+                state = kif->t_state;
+            else
+                state = PSUTIL_CONN_NONE;
+
+            // local & remote port
+            lport = ntohs(kif->inp_lport);
+            rport = ntohs(kif->inp_fport);
+
+            // local addr
+            inet_ntop(kif->so_family, &kif->inp_laddru, lip, sizeof(lip));
+            py_laddr = Py_BuildValue("(si)", lip, lport);
+            if (! py_laddr)
+                goto error;
+
+            // remote addr
+            if (rport != 0) {
+                inet_ntop(kif->so_family, &kif->inp_faddru, rip, sizeof(rip));
+                py_raddr = Py_BuildValue("(si)", rip, rport);
+            }
+            else {
+                py_raddr = Py_BuildValue("()");
+            }
+            if (! py_raddr)
+                goto error;
+
+            // populate tuple and list
+            py_tuple = Py_BuildValue(
+                "(iiiNNil)",
+                kif->fd_fd,
+                kif->so_family,
+                kif->so_type,
+                py_laddr,
+                py_raddr,
+                state,
+                kif->p_pid
+            );
+            if (! py_tuple)
+                goto error;
+            if (PyList_Append(py_retlist, py_tuple))
+                goto error;
+            Py_DECREF(py_tuple);
+        }
+        // UNIX socket
+        else if (kif->so_family == AF_UNIX) {
+            py_lpath = PyUnicode_DecodeFSDefault(kif->unp_path);
+            if (! py_lpath)
+                goto error;
+
+            py_tuple = Py_BuildValue(
+                "(iiiOsil)",
+                kif->fd_fd,
+                kif->so_family,
+                kif->so_type,
+                py_lpath,
+                "",  // raddr
+                PSUTIL_CONN_NONE,
+                kif->p_pid
+            );
+            if (! py_tuple)
+                goto error;
+            if (PyList_Append(py_retlist, py_tuple))
+                goto error;
+            Py_DECREF(py_lpath);
+            Py_DECREF(py_tuple);
+            Py_INCREF(Py_None);
+        }
+    }
+
+    kvm_close(kd);
+    return py_retlist;
+
+error:
+    Py_XDECREF(py_tuple);
+    Py_XDECREF(py_laddr);
+    Py_XDECREF(py_raddr);
+    Py_DECREF(py_retlist);
+    if (kd != NULL)
+        kvm_close(kd);
+    return NULL;
+}
diff --git a/psutil/arch/openbsd/socks.h b/psutil/arch/openbsd/socks.h
new file mode 100644
index 00000000..90b678bb
--- /dev/null
+++ b/psutil/arch/openbsd/socks.h
@@ -0,0 +1,8 @@
+/*
+ * Copyright (c) 2009, Giampaolo Rodola'.
+ * All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+PyObject *psutil_net_connections(PyObject* self, PyObject* args);