amphetamine 0.1.0

//! The process table, plus per-process resident memory, via libproc.
//!
//! A macOS "app" is a tree, not a process: Cursor's real footprint is its main
//! process plus dozens of renderer and extension-host helpers. Anything that
//! reports on an app has to sum the tree, so that lives here.

use anyhow::{Result, anyhow};
use std::collections::{HashMap, HashSet, VecDeque};
use std::ffi::c_void;
use std::mem::{size_of, zeroed};

const PROC_PIDTASKINFO: libc::c_int = 4;
const PROC_PIDTBSDINFO: libc::c_int = 3;

#[derive(Debug, Clone)]
pub struct Proc {
    pub pid: i32,
    pub ppid: i32,
    pub uid: u32,
    pub name: String,
    pub rss: u64,
}

/// Every process we are allowed to inspect, indexed by pid.
///
/// In practice that means our own: macOS refuses `PROC_PIDTBSDINFO` for other
/// users' processes, so root-owned daemons never appear here. That suits us —
/// they are also the ones we could never act on.
#[derive(Debug, Clone, Default)]
pub struct Table {
    pub by_pid: HashMap<i32, Proc>,
    children: HashMap<i32, Vec<i32>>,
}

impl Table {
    pub fn load() -> Result<Self> {
        let procs = list()?;
        let mut children: HashMap<i32, Vec<i32>> = HashMap::new();
        for p in &procs {
            children.entry(p.ppid).or_default().push(p.pid);
        }
        Ok(Self {
            by_pid: procs.into_iter().map(|p| (p.pid, p)).collect(),
            children,
        })
    }

    /// `root` and every descendant, breadth-first.
    ///
    /// The `seen` set is load-bearing, not defensive habit: a process whose
    /// parent exits is reparented mid-walk, which can briefly present a cycle.
    pub fn tree(&self, root: i32) -> Vec<&Proc> {
        let mut seen = HashSet::new();
        let mut queue = VecDeque::from([root]);
        let mut out = Vec::new();
        while let Some(pid) = queue.pop_front() {
            if !seen.insert(pid) {
                continue;
            }
            if let Some(p) = self.by_pid.get(&pid) {
                out.push(p);
            }
            queue.extend(self.children.get(&pid).into_iter().flatten().copied());
        }
        out
    }

    /// Resident memory of an app's entire process tree.
    pub fn tree_rss(&self, root: i32) -> u64 {
        self.tree(root).iter().map(|p| p.rss).sum()
    }

    /// Walks from `pid`'s parent up towards launchd, excluding `pid` itself.
    ///
    /// Bounded by the table size rather than by reaching pid 1, so a reparented
    /// process that forms a loop cannot hang the walk.
    pub fn ancestors(&self, pid: i32) -> impl Iterator<Item = i32> + '_ {
        let mut cur = self.by_pid.get(&pid).map(|p| p.ppid);
        let mut budget = self.by_pid.len();
        std::iter::from_fn(move || {
            let pid = cur.filter(|&p| p > 1)?;
            budget = budget.checked_sub(1)?;
            cur = self.by_pid.get(&pid).map(|p| p.ppid);
            Some(pid)
        })
    }
}

fn list() -> Result<Vec<Proc>> {
    let count = unsafe { libc::proc_listallpids(std::ptr::null_mut(), 0) };
    if count <= 0 {
        return Err(anyhow!(
            "proc_listallpids: {}",
            std::io::Error::last_os_error()
        ));
    }

    // Headroom for processes spawned between sizing and reading.
    let cap = count as usize + 256;
    let mut pids = vec![0i32; cap];
    // The buffer is sized in bytes, but the return value is a count of pids —
    // not, as the name suggests by analogy with the rest of libproc, a byte
    // length. Dividing it by the element size silently loses three quarters of
    // the machine's processes.
    let bytes = (cap * size_of::<i32>()) as libc::c_int;
    let got = unsafe { libc::proc_listallpids(pids.as_mut_ptr() as *mut c_void, bytes) };
    if got <= 0 {
        return Err(anyhow!(
            "proc_listallpids: {}",
            std::io::Error::last_os_error()
        ));
    }
    pids.truncate(got as usize);

    // Processes exit constantly; a pid that vanishes between listing and
    // inspection is normal and is simply dropped.
    Ok(pids
        .into_iter()
        .filter(|&p| p > 0)
        .filter_map(info)
        .collect())
}

fn info(pid: i32) -> Option<Proc> {
    let mut bsd: libc::proc_bsdinfo = unsafe { zeroed() };
    let size = size_of::<libc::proc_bsdinfo>() as libc::c_int;
    let got = unsafe {
        libc::proc_pidinfo(
            pid,
            PROC_PIDTBSDINFO,
            0,
            &mut bsd as *mut _ as *mut c_void,
            size,
        )
    };
    if got != size {
        return None;
    }
    // pbi_name is the fuller name and is empty for some kernel-side processes,
    // where pbi_comm is all that exists.
    let name = match fixed_str(&bsd.pbi_name) {
        n if !n.is_empty() => n,
        _ => fixed_str(&bsd.pbi_comm),
    };
    Some(Proc {
        pid,
        ppid: bsd.pbi_ppid as i32,
        uid: bsd.pbi_uid,
        name,
        rss: rss(pid),
    })
}

/// Reads a fixed-width C char array that is only NUL-terminated when it fits,
/// so a name occupying the full array must stop at the array bound instead.
fn fixed_str(raw: &[libc::c_char]) -> String {
    let bytes: Vec<u8> = raw
        .iter()
        .take_while(|&&c| c != 0)
        .map(|&c| c as u8)
        .collect();
    String::from_utf8_lossy(&bytes).into_owned()
}

/// Resident set size, or 0 for processes we may not inspect.
pub fn rss(pid: i32) -> u64 {
    let mut ti: libc::proc_taskinfo = unsafe { zeroed() };
    let size = size_of::<libc::proc_taskinfo>() as libc::c_int;
    let got = unsafe {
        libc::proc_pidinfo(
            pid,
            PROC_PIDTASKINFO,
            0,
            &mut ti as *mut _ as *mut c_void,
            size,
        )
    };
    if got == size { ti.pti_resident_size } else { 0 }
}

/// Whether the pid still exists. Used to confirm a graceful quit actually landed.
pub fn is_alive(pid: i32) -> bool {
    // Signal 0 runs the permission and existence checks without delivering.
    if unsafe { libc::kill(pid, 0) } == 0 {
        return true;
    }
    // EPERM means it exists but belongs to someone else — still alive.
    std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn table_contains_this_process_with_real_memory() {
        let t = Table::load().expect("process table should be readable");
        let me = std::process::id() as i32;
        let this = t.by_pid.get(&me).expect("our own pid must be in the table");
        assert_eq!(this.uid, unsafe { libc::getuid() });
        assert!(this.rss > 0, "our own RSS should be non-zero");
        assert!(t.by_pid.len() > 20, "suspiciously short process table");
    }

    #[test]
    fn table_matches_the_machine_rather_than_a_fraction_of_it() {
        let t = Table::load().unwrap();
        let mine = String::from_utf8(
            std::process::Command::new("sh")
                .args([
                    "-c",
                    "ps -Ao user=,pid= | awk -v u=\"$(whoami)\" '$1==u' | wc -l",
                ])
                .output()
                .unwrap()
                .stdout,
        )
        .unwrap()
        .trim()
        .parse::<usize>()
        .unwrap();

        // Guards the pid-count-vs-byte-length trap in `list`: getting that
        // wrong still yields a plausible-looking table, just a truncated one.
        // Compared against our own processes because macOS only reports
        // bsdinfo for those — see the note on `Table`.
        let ratio = t.by_pid.len() as f64 / mine as f64;
        assert!(
            (0.9..=1.1).contains(&ratio),
            "table has {} processes but we own {mine}",
            t.by_pid.len()
        );
    }

    #[test]
    fn tree_includes_root_and_terminates() {
        let t = Table::load().unwrap();
        let me = std::process::id() as i32;
        assert!(t.tree(me).iter().any(|p| p.pid == me));
        assert!(t.tree_rss(me) >= t.by_pid[&me].rss);
        // launchd roots everything, so its tree is the whole machine — this
        // exercises the cycle guard and proves traversal terminates.
        assert!(t.tree(1).len() > 10);
    }

    #[test]
    fn fixed_str_stops_at_nul_and_at_bound() {
        let padded: Vec<libc::c_char> = b"Arc\0\0\0".iter().map(|&b| b as libc::c_char).collect();
        assert_eq!(fixed_str(&padded), "Arc");
        // No NUL anywhere: must stop at the array bound rather than run off it.
        let full: Vec<libc::c_char> = b"abcdefgh".iter().map(|&b| b as libc::c_char).collect();
        assert_eq!(fixed_str(&full), "abcdefgh");
    }

    #[test]
    fn liveness_tracks_reality() {
        assert!(is_alive(std::process::id() as i32));
        // Above the kernel's pid ceiling, so it cannot be allocated.
        assert!(!is_alive(i32::MAX - 1));
    }
}