bug: add bindings to grab ppid in some cases on macos (#825)

6 files changed, 393 insertions, 26 deletions

diff --git a/Cargo.lock b/Cargo.lock
index df094d40..d757044a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -227,6 +227,7 @@ dependencies = [
  "itertools",
  "libc",
  "log",
+ "mach2",
  "nvml-wrapper",
  "once_cell",
  "predicates",
@@ -968,6 +969,15 @@ dependencies = [
 ]
 
 [[package]]
+name = "mach2"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d0d1830bcd151a6fc4aea1369af235b36c1528fe976b8ff678683c9995eade8"
+dependencies = [
+ "libc",
+]
+
+[[package]]
 name = "memchr"
 version = "2.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 8af54df2..d2c8386e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -61,7 +61,6 @@ concat-string = "1.0.1"
 crossterm = "0.18.2"
 ctrlc = { version = "3.1.9", features = ["termination"] }
 dirs = "4.0.0"
-
 fern = { version = "0.6.1", optional = true }
 futures = "0.3.21"
 futures-timer = "3.0.2"
@@ -93,6 +92,7 @@ smol = "1.2.5"
 
 [target.'cfg(target_os = "macos")'.dependencies]
 heim = { version = "0.1.0-rc.1", features = ["cpu", "disk", "memory", "net"] }
+mach2 = "0.4.1"
 
 [target.'cfg(target_os = "windows")'.dependencies]
 heim = { version = "0.1.0-rc.1", features = ["cpu", "disk", "memory"] }
@@ -114,9 +114,21 @@ clap_mangen = "0.1.6"
 [package.metadata.deb]
 section = "utility"
 assets = [
-    ["target/release/btm", "usr/bin/", "755"],
-    ["LICENSE", "usr/share/doc/btm/", "644"],
-    ["manpage/btm.1.gz", "usr/share/man/man1/btm.1.gz", "644"],
+    [
+        "target/release/btm",
+        "usr/bin/",
+        "755",
+    ],
+    [
+        "LICENSE",
+        "usr/share/doc/btm/",
+        "644",
+    ],
+    [
+        "manpage/btm.1.gz",
+        "usr/share/man/man1/btm.1.gz",
+        "644",
+    ],
     [
         "completion/btm.bash",
         "usr/share/bash-completion/completions/btm",
@@ -127,7 +139,11 @@ assets = [
         "usr/share/fish/vendor_completions.d/btm.fish",
         "644",
     ],
-    ["completion/_btm", "usr/share/zsh/vendor-completions/", "644"],
+    [
+        "completion/_btm",
+        "usr/share/zsh/vendor-completions/",
+        "644",
+    ],
 ]
 extended-description = """\
 A customizable cross-platform graphical process/system monitor for the terminal. Supports Linux, macOS, and Windows.
diff --git a/src/app/data_farmer.rs b/src/app/data_farmer.rs
index 3b6f6dc1..a6670712 100644
--- a/src/app/data_farmer.rs
+++ b/src/app/data_farmer.rs
@@ -110,21 +110,17 @@ impl ProcessData {
             .collect();
         self.process_harvest = process_pid_map;
 
-        // This also needs a quick sort + reverse to be in the correct order.
+        // We collect all processes that either:
+        // - Do not have a parent PID (that is, they are orphan processes)
+        // - Have a parent PID but we don't have the parent (we promote them as orphans)
+        // Note this also needs a quick sort + reverse to be in the correct order.
         self.orphan_pids = {
             let mut res: Vec<Pid> = self
                 .process_harvest
                 .iter()
-                .filter_map(|(pid, process_harvest)| {
-                    if let Some(parent_pid) = process_harvest.parent_pid {
-                        if self.process_harvest.contains_key(&parent_pid) {
-                            None
-                        } else {
-                            Some(*pid)
-                        }
-                    } else {
-                        Some(*pid)
-                    }
+                .filter_map(|(pid, process_harvest)| match process_harvest.parent_pid {
+                    Some(parent_pid) if self.process_harvest.contains_key(&parent_pid) => None,
+                    _ => Some(*pid),
                 })
                 .sorted()
                 .collect();
diff --git a/src/app/data_harvester/processes/macos.rs b/src/app/data_harvester/processes/macos.rs
index c9a3f601..2431bf9b 100644
--- a/src/app/data_harvester/processes/macos.rs
+++ b/src/app/data_harvester/processes/macos.rs
@@ -1,9 +1,10 @@
-//! Process data collection for macOS.  Uses sysinfo.
+//! Process data collection for macOS.  Uses sysinfo and custom bindings.
 
 use super::ProcessHarvest;
 use sysinfo::System;
 
-use crate::data_harvester::processes::UserTable;
+use crate::{data_harvester::processes::UserTable, Pid};
+mod sysctl_bindings;
 
 pub fn get_process_data(
     sys: &System, use_current_cpu_total: bool, mem_total_kb: u64, user_table: &mut UserTable,
@@ -17,8 +18,14 @@ pub fn get_process_data(
     )
 }
 
+pub(crate) fn fallback_macos_ppid(pid: Pid) -> Option<Pid> {
+    sysctl_bindings::kinfo_process(pid)
+        .map(|kinfo| kinfo.kp_eproc.e_ppid)
+        .ok()
+}
+
 fn get_macos_process_cpu_usage(
-    pids: &[i32],
+    pids: &[Pid],
 ) -> std::io::Result<std::collections::HashMap<i32, f64>> {
     use itertools::Itertools;
     let output = std::process::Command::new("ps")
diff --git a/src/app/data_harvester/processes/macos/sysctl_bindings.rs b/src/app/data_harvester/processes/macos/sysctl_bindings.rs
new file mode 100644
index 00000000..77b8901b
--- /dev/null
+++ b/src/app/data_harvester/processes/macos/sysctl_bindings.rs
@@ -0,0 +1,322 @@
+//! Partial bindings from Apple's open source code for getting process information.
+//! Some of this is based on [heim's binding implementation](https://github.com/heim-rs/heim/blob/master/heim-process/src/sys/macos/bindings/process.rs).
+
+use std::mem;
+
+use anyhow::{bail, Result};
+use libc::{
+    boolean_t, c_char, c_long, c_short, c_uchar, c_ushort, c_void, dev_t, gid_t, itimerval, pid_t,
+    rusage, sigset_t, timeval, uid_t, xucred, CTL_KERN, KERN_PROC, KERN_PROC_PID, MAXCOMLEN,
+};
+use mach2::vm_types::user_addr_t;
+
+use crate::Pid;
+
+#[allow(non_camel_case_types)]
+#[repr(C)]
+pub(crate) struct kinfo_proc {
+    pub kp_proc: extern_proc,
+    pub kp_eproc: eproc,
+}
+
+#[allow(non_camel_case_types)]
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct p_st1 {
+    /// Doubly-linked run/sleep queue.
+    p_forw: user_addr_t,
+    p_back: user_addr_t,
+}
+
+#[allow(non_camel_case_types)]
+#[repr(C)]
+pub union p_un {
+    pub p_st1: p_st1,
+
+    /// process start time
+    pub p_starttime: timeval,
+}
+
+/// Exported fields for kern sysctl. See
+/// [`proc.h`](https://opensource.apple.com/source/xnu/xnu-201/bsd/sys/proc.h)
+#[allow(non_camel_case_types)]
+#[repr(C)]
+pub(crate) struct extern_proc {
+    pub p_un: p_un,
+
+    /// Address space.
+    pub p_vmspace: *mut vmspace,
+
+    /// Signal actions, state (PROC ONLY). Should point to
+    /// a `sigacts` but we don't really seem to need this.
+    pub p_sigacts: user_addr_t,
+
+    /// P_* flags.
+    pub p_flag: i32,
+
+    /// S* process status.
+    pub p_stat: c_char,
+
+    /// Process identifier.
+    pub p_pid: pid_t,
+
+    /// Save parent pid during ptrace. XXX
+    pub p_oppid: pid_t,
+
+    /// Sideways return value from fdopen. XXX
+    pub p_dupfd: i32,
+
+    /// where user stack was allocated
+    pub user_stack: caddr_t,
+
+    /// XXX Which thread is exiting?
+    pub exit_thread: *mut c_void,
+
+    /// allow to debug
+    pub p_debugger: i32,
+
+    /// indication to suspend
+    pub sigwait: boolean_t,
+
+    /// Time averaged value of p_cpticks.
+    pub p_estcpu: u32,
+
+    /// Ticks of cpu time.
+    pub p_cpticks: i32,
+
+    /// %cpu for this process during p_swtime
+    pub p_pctcpu: fixpt_t,
+
+    /// Sleep address.
+    pub p_wchan: *mut c_void,
+
+    /// Reason for sleep.
+    pub p_wmesg: *mut c_char,
+
+    /// Time swapped in or out.
+    pub p_swtime: u32,
+
+    /// Time since last blocked.
+    pub p_slptime: u32,
+
+    /// Alarm timer.
+    pub p_realtimer: itimerval,
+
+    /// Real time.
+    pub p_rtime: timeval,
+
+    /// Statclock hit in user mode.
+    pub p_uticks: u64,
+
+    /// Statclock hits in system mode.
+    pub p_sticks: u64,
+
+    /// Statclock hits processing intr.
+    pub p_iticks: u64,
+
+    /// Kernel trace points.
+    pub p_traceflag: i32,
+
+    /// Trace to vnode. Originally a pointer to a struct of vnode.
+    pub p_tracep: *mut c_void,
+
+    /// DEPRECATED.
+    pub p_siglist: i32,
+
+    /// Vnode of executable. Originally a pointer to a struct of vnode.
+    pub p_textvp: *mut c_void,
+
+    /// If non-zero, don't swap.
+    pub p_holdcnt: i32,
+
+    /// DEPRECATED.
+    pub p_sigmask: sigset_t,
+
+    /// Signals being ignored.
+    pub p_sigignore: sigset_t,
+
+    /// Signals being caught by user.
+    pub p_sigcatch: sigset_t,
+
+    /// Process priority.
+    pub p_priority: c_uchar,
+
+    /// User-priority based on p_cpu and p_nice.
+    pub p_usrpri: c_uchar,
+
+    /// Process "nice" value.
+    pub p_nice: c_char,
+
+    pub p_comm: [c_char; MAXCOMLEN + 1],
+
+    /// Pointer to process group. Originally a pointer to a `pgrp`.
+    pub p_pgrp: *mut c_void,
+
+    /// Kernel virtual addr of u-area (PROC ONLY). Originally a pointer to a `user`.
+    pub p_addr: *mut c_void,
+
+    /// Exit status for wait; also stop signal.
+    pub p_xstat: c_ushort,
+
+    /// Accounting flags.
+    pub p_acflag: c_ushort,
+
+    /// Exit information. XXX
+    pub p_ru: *mut rusage,
+}
+
+const WMESGLEN: usize = 7;
+const COMAPT_MAXLOGNAME: usize = 12;
+
+/// See `_caddr_t.h`.
+#[allow(non_camel_case_types)]
+type caddr_t = *const libc::c_char;
+
+/// See `types.h`.
+#[allow(non_camel_case_types)]
+type segsz_t = i32;
+
+/// See `types.h`.
+#[allow(non_camel_case_types)]
+type fixpt_t = u32;
+
+/// See [`proc.h`](https://opensource.apple.com/source/xnu/xnu-201/bsd/sys/proc.h)
+#[allow(non_camel_case_types)]
+#[repr(C)]
+pub(crate) struct pcred {
+    pub pc_lock: [c_char; 72],
+    pub pc_ucred: *mut xucred,
+    pub p_ruid: uid_t,
+    pub p_svuid: uid_t,
+    pub p_rgid: gid_t,
+    pub p_svgid: gid_t,
+    pub p_refcnt: i32,
+}
+
+/// See `vm.h`.
+#[allow(non_camel_case_types)]
+#[repr(C)]
+pub(crate) struct vmspace {
+    pub dummy: i32,
+    pub dummy2: caddr_t,
+    pub dummy3: [i32; 5],
+    pub dummy4: [caddr_t; 3],
+}
+
+/// See [`sysctl.h`](https://opensource.apple.com/source/xnu/xnu-344/bsd/sys/sysctl.h).
+#[allow(non_camel_case_types)]
+#[repr(C)]
+pub(crate) struct eproc {
+    /// Address of proc. We just cheat and use a c_void pointer since we aren't using this.
+    pub e_paddr: *mut c_void,
+
+    /// Session pointer.  We just cheat and use a c_void pointer since we aren't using this.
+    pub e_sess: *mut c_void,
+
+    /// Process credentials
+    pub e_pcred: pcred,
+
+    /// Current credentials
+    pub e_ucred: xucred,
+
+    /// Address space
+    pub e_vm: vmspace,
+
+    /// Parent process ID
+    pub e_ppid: pid_t,
+
+    /// Process group ID
+    pub e_pgid: pid_t,
+
+    /// Job control counter
+    pub e_jobc: c_short,
+
+    /// Controlling tty dev
+    pub e_tdev: dev_t,
+
+    /// tty process group id
+    pub e_tpgid: pid_t,
+
+    /// tty session pointer.  We just cheat and use a c_void pointer since we aren't using this.
+    pub e_tsess: *mut c_void,
+
+    /// wchan message
+    pub e_wmesg: [c_char; WMESGLEN + 1],
+
+    /// text size
+    pub e_xsize: segsz_t,
+
+    /// text rss
+    pub e_xrssize: c_short,
+
+    /// text references
+    pub e_xccount: c_short,
+
+    pub e_xswrss: c_short,
+
+    pub e_flag: c_long,
+
+    /// short setlogin() name
+    pub e_login: [c_char; COMAPT_MAXLOGNAME],
+
+    pub e_spare: [c_long; 4],
+}
+
+/// Obtains the [`kinfo_proc`] given a process PID.
+///
+/// From [heim](https://github.com/heim-rs/heim/blob/master/heim-process/src/sys/macos/bindings/process.rs#L235).
+pub(crate) fn kinfo_process(pid: Pid) -> Result<kinfo_proc> {
+    let mut name: [i32; 4] = [CTL_KERN, KERN_PROC, KERN_PROC_PID, pid];
+    let mut size = mem::size_of::<kinfo_proc>();
+    let mut info = mem::MaybeUninit::<kinfo_proc>::uninit();
+
+    let result = unsafe {
+        libc::sysctl(
+            name.as_mut_ptr(),
+            4,
+            info.as_mut_ptr() as *mut libc::c_void,
+            &mut size,
+            std::ptr::null_mut(),
+            0,
+        )
+    };
+
+    if result < 0 {
+        bail!("failed to get process for pid {pid}");
+    }
+
+    // sysctl succeeds but size is zero, happens when process has gone away
+    if size == 0 {
+        bail!("failed to get process for pid {pid}");
+    }
+
+    unsafe { Ok(info.assume_init()) }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::mem;
+
+    /// A quick test to ensure that things are sized correctly.
+    #[test]
+    fn test_struct_sizes() {
+        assert_eq!(mem::size_of::<p_st1>(), 16);
+        assert_eq!(mem::align_of::<p_st1>(), 8);
+
+        assert_eq!(mem::size_of::<pcred>(), 104);
+        assert_eq!(mem::align_of::<pcred>(), 8);
+
+        assert_eq!(mem::size_of::<vmspace>(), 64);
+        assert_eq!(mem::align_of::<vmspace>(), 8);
+
+        assert_eq!(mem::size_of::<extern_proc>(), 296);
+        assert_eq!(mem::align_of::<extern_proc>(), 8);
+
+        assert_eq!(mem::size_of::<eproc>(), 376);
+        assert_eq!(mem::align_of::<eproc>(), 8);
+
+        assert_eq!(mem::size_of::<kinfo_proc>(), 672);
+        assert_eq!(mem::align_of::<kinfo_proc>(), 8);
+    }
+}
diff --git a/src/app/data_harvester/processes/macos_freebsd.rs b/src/app/data_harvester/processes/macos_freebsd.rs
index 39b57e05..65725732 100644
--- a/src/app/data_harvester/processes/macos_freebsd.rs
+++ b/src/app/data_harvester/processes/macos_freebsd.rs
@@ -6,12 +6,15 @@ use std::io;
 use super::ProcessHarvest;
 use sysinfo::{CpuExt, PidExt, ProcessExt, ProcessStatus, System, SystemExt};
 
-use crate::data_harvester::processes::UserTable;
+use crate::{data_harvester::processes::UserTable, utils::error::Result, Pid};
 
-pub fn get_process_data(
+pub fn get_process_data<F>(
     sys: &System, use_current_cpu_total: bool, mem_total_kb: u64, user_table: &mut UserTable,
-    get_process_cpu_usage: impl Fn(&[i32]) -> io::Result<HashMap<i32, f64>>,
-) -> crate::utils::error::Result<Vec<ProcessHarvest>> {
+    get_process_cpu_usage: F,
+) -> Result<Vec<ProcessHarvest>>
+where
+    F: Fn(&[Pid]) -> io::Result<HashMap<Pid, f64>>,
+{
     let mut process_vector: Vec<ProcessHarvest> = Vec::new();
     let process_hashmap = sys.processes();
     let cpu_usage = sys.global_cpu_info().cpu_usage() as f64 / 100.0;
@@ -66,9 +69,22 @@ pub fn get_process_data(
             (ps.to_string(), convert_process_status_to_char(ps))
         };
         let uid = process_val.user_id().map(|u| **u);
+        let pid = process_val.pid().as_u32() as Pid;
         process_vector.push(ProcessHarvest {
-            pid: process_val.pid().as_u32() as _,
-            parent_pid: process_val.parent().map(|p| p.as_u32() as _),
+            pid,
+            parent_pid: {
+                #[cfg(target_os = "macos")]
+                {
+                    process_val
+                        .parent()
+                        .map(|p| p.as_u32() as _)
+                        .or_else(|| super::fallback_macos_ppid(pid))
+                }
+                #[cfg(not(target_os = "macos"))]
+                {
+                    process_val.parent().map(|p| p.as_u32() as _)
+                }
+            },
             name,
             command,
             mem_usage_percent: if mem_total_kb > 0 {
@@ -96,7 +112,7 @@ pub fn get_process_data(
     }
 
     let unknown_state = ProcessStatus::Unknown(0).to_string();
-    let cpu_usage_unknown_pids: Vec<i32> = process_vector
+    let cpu_usage_unknown_pids: Vec<Pid> = process_vector
         .iter()
         .filter(|process| process.process_state.0 == unknown_state)
         .map(|process| process.pid)