amphetamine 0.1.1

//! Winning scheduler contention — the only honest form of "acceleration"
//! available on Apple Silicon.
//!
//! There is no clock control here, because none exists in user space. What
//! exists is the nice value: deprioritising the processes competing with your
//! editor leaves more CPU for the editor. Note what this is *not* — it does not
//! move anything onto the efficiency cores. `PRIO_DARWIN_PROCESS`, the
//! mechanism behind `taskpolicy -b`, silently does nothing when aimed at
//! another process on current macOS; it reports success and changes no state.
//! Nice is the lever that actually moves.
//!
//! The central rule of this module: **never demote what we cannot un-demote.**
//! Lowering a nice value requires root, so a session verifies its restore path
//! before it touches a single process.

use crate::{apps::App, config::Config, guard, privilege, proc};
use anyhow::Result;
use std::ffi::{CString, c_char, c_void};
use std::io;

type CFStringRef = *const c_void;
type CFAllocatorRef = *const c_void;

const CF_STRING_ENCODING_UTF8: u32 = 0x0800_0100;
const IOPM_ASSERTION_LEVEL_ON: u32 = 255;

#[link(name = "CoreFoundation", kind = "framework")]
unsafe extern "C" {
    fn CFStringCreateWithCString(alloc: CFAllocatorRef, s: *const c_char, enc: u32) -> CFStringRef;
    fn CFRelease(cf: *const c_void);
}

#[link(name = "IOKit", kind = "framework")]
unsafe extern "C" {
    fn IOPMAssertionCreateWithName(
        kind: CFStringRef,
        level: u32,
        name: CFStringRef,
        id: *mut u32,
    ) -> i32;
    fn IOPMAssertionRelease(id: u32) -> i32;
}

/// Current nice value, or `None` if the process is gone or unreadable.
pub fn nice_of(pid: i32) -> Option<i32> {
    // getpriority returns -1 legitimately, so errno is the only way to tell a
    // real value from a failure.
    unsafe { *libc::__error() = 0 };
    let v = unsafe { libc::getpriority(libc::PRIO_PROCESS, pid as u32) };
    (unsafe { *libc::__error() } == 0).then_some(v)
}

/// Raises a process's nice value. Needs no privileges for our own processes.
pub fn demote(pid: i32, level: i32) -> io::Result<()> {
    match unsafe { libc::setpriority(libc::PRIO_PROCESS, pid as u32, level) } {
        0 => Ok(()),
        _ => Err(io::Error::last_os_error()),
    }
}

/// A held power assertion, released when dropped.
struct Assertion(u32);

impl Assertion {
    /// Holds off idle sleep so a long build is not interrupted by the machine
    /// dozing off. Display sleep is deliberately left alone — this should not
    /// keep your screen awake.
    fn prevent_idle_sleep() -> Option<Self> {
        let kind = cfstring("PreventUserIdleSystemSleep")?;
        let name = cfstring("Amphetamine focus session")?;
        let mut id = 0u32;
        let rc =
            unsafe { IOPMAssertionCreateWithName(kind, IOPM_ASSERTION_LEVEL_ON, name, &mut id) };
        unsafe {
            CFRelease(kind);
            CFRelease(name);
        }
        (rc == 0).then_some(Self(id))
    }
}

impl Drop for Assertion {
    fn drop(&mut self) {
        unsafe { IOPMAssertionRelease(self.0) };
    }
}

fn cfstring(s: &str) -> Option<CFStringRef> {
    let c = CString::new(s).ok()?;
    let r =
        unsafe { CFStringCreateWithCString(std::ptr::null(), c.as_ptr(), CF_STRING_ENCODING_UTF8) };
    (!r.is_null()).then_some(r)
}

#[derive(Debug, Clone)]
pub struct Demotion {
    pub name: String,
    pub pids: usize,
    pub rss: u64,
    pub note: Option<String>,
}

/// An active focus session. Dropping it returns every process it demoted to
/// nice 0 and releases the sleep assertion.
pub struct Session {
    demoted: Vec<i32>,
    _assertion: Option<Assertion>,
    pub report: Vec<Demotion>,
    pub holding_sleep: bool,
    /// False when the privilege grant is missing, in which case nothing was
    /// demoted at all.
    pub restore_ready: bool,
}

impl Session {
    pub fn start(cfg: &Config, apps: &[App], table: &proc::Table, dry_run: bool) -> Result<Self> {
        Self::start_with(cfg, apps, table, dry_run, privilege::can_restore())
    }

    /// `restore_ready` is injected so the refusal path is testable without
    /// installing or removing a real sudoers file.
    pub fn start_with(
        cfg: &Config,
        apps: &[App],
        table: &proc::Table,
        dry_run: bool,
        restore_ready: bool,
    ) -> Result<Self> {
        let level = cfg.focus.nice_level.clamp(1, 20);
        let targets: Vec<&App> = apps
            .iter()
            .filter(|a| guard::any_matches(&a.identities(), &cfg.focus.demote))
            .collect();

        // Check before acting, not after: demoting first and discovering the
        // restore path is missing afterwards is the exact failure this whole
        // module is arranged to prevent.
        let act = !dry_run && restore_ready;
        let mut demoted = Vec::new();
        let mut report = Vec::new();

        for app in targets {
            let ids = app.identities();
            if let Err(refusal) = guard::vet_process(app.pid, app.uid, &ids) {
                report.push(Demotion {
                    name: app.name.clone(),
                    pids: 0,
                    rss: 0,
                    note: Some(refusal.to_string()),
                });
                continue;
            }

            // Demoting only the main process would accomplish nothing: a
            // browser or Electron app does its work in child processes.
            //
            // Only processes sitting at nice 0 are eligible. Anything already
            // carrying a nice value was set by something else, and restoring it
            // to 0 would be us guessing at a value that was never ours.
            let eligible: Vec<i32> = table
                .tree(app.pid)
                .iter()
                .map(|p| p.pid)
                .filter(|&pid| nice_of(pid) == Some(0))
                .collect();

            let mut moved = 0;
            if act {
                for pid in &eligible {
                    if demote(*pid, level).is_ok() {
                        demoted.push(*pid);
                        moved += 1;
                    }
                }
            }
            report.push(Demotion {
                name: app.name.clone(),
                pids: if act { moved } else { eligible.len() },
                rss: app.rss,
                note: (!act && !dry_run).then(|| "no restore path — run `amph setup`".into()),
            });
        }

        let assertion = (cfg.focus.prevent_sleep && !dry_run)
            .then(Assertion::prevent_idle_sleep)
            .flatten();

        Ok(Self {
            demoted,
            holding_sleep: assertion.is_some(),
            _assertion: assertion,
            report,
            restore_ready,
        })
    }

    pub fn demoted_count(&self) -> usize {
        self.demoted.len()
    }

    /// Returns every demoted process to nice 0, reporting anything that failed.
    /// Called by `Drop`, but exposed so a caller can surface the errors.
    pub fn restore_now(&mut self) -> Vec<String> {
        self.demoted
            .drain(..)
            .filter_map(|pid| match privilege::renice_to_zero(pid) {
                Ok(()) => None,
                // A process that exited on its own needs no restoring.
                Err(_) if !proc::is_alive(pid) => None,
                Err(e) => Some(e.to_string()),
            })
            .collect()
    }
}

impl Drop for Session {
    fn drop(&mut self) {
        self.restore_now();
    }
}

/// Returns anything matching the configured demote list to nice 0. The safety
/// net for a session that died without unwinding — a `SIGKILL` or a lost
/// terminal skips `Drop`, and this puts things back.
pub fn restore_all(cfg: &Config, apps: &[App], table: &proc::Table) -> Vec<(String, usize)> {
    apps.iter()
        .filter(|a| guard::any_matches(&a.identities(), &cfg.focus.demote))
        .map(|app| {
            let n = table
                .tree(app.pid)
                .iter()
                .filter(|p| nice_of(p.pid).is_some_and(|v| v > 0))
                .filter(|p| privilege::renice_to_zero(p.pid).is_ok())
                .count();
            (app.name.clone(), n)
        })
        .filter(|(_, n)| *n > 0)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::process::{Command, Stdio};

    /// A real child process, reaped on drop.
    struct Child(std::process::Child);
    impl Drop for Child {
        fn drop(&mut self) {
            self.0.kill().ok();
            self.0.wait().ok();
        }
    }

    fn spawn_sleeper() -> Child {
        Child(
            Command::new("/bin/sleep")
                .arg("30")
                .stdout(Stdio::null())
                .spawn()
                .expect("spawn sleep"),
        )
    }

    #[test]
    fn demotion_raises_nice_on_a_real_process() {
        let child = spawn_sleeper();
        let pid = child.0.id() as i32;

        assert_eq!(nice_of(pid), Some(0), "a fresh process starts at nice 0");
        demote(pid, 10).expect("raising nice needs no privileges");
        assert_eq!(nice_of(pid), Some(10));
    }

    #[test]
    fn lowering_nice_back_is_refused_without_privileges() {
        // The asymmetry this module is built around. If this ever starts
        // passing, the sudoers grant is no longer necessary.
        let child = spawn_sleeper();
        let pid = child.0.id() as i32;
        demote(pid, 10).unwrap();

        let err = demote(pid, 0).expect_err("unprivileged lowering must fail");
        assert_eq!(err.raw_os_error(), Some(libc::EACCES));
        assert_eq!(nice_of(pid), Some(10), "and the value must be unchanged");
    }

    #[test]
    fn nothing_is_demoted_without_a_proven_restore_path() {
        let table = proc::Table::load().unwrap();
        let apps = crate::apps::list(&table);
        let Some(victim) = apps.iter().find(|a| {
            guard::protected_match(&a.identities()).is_none() && a.uid == unsafe { libc::getuid() }
        }) else {
            return;
        };
        let cfg = Config {
            focus: crate::config::Focus {
                demote: vec![victim.name.clone()],
                prevent_sleep: false,
                nice_level: 10,
            },
            ..Default::default()
        };

        let before: Vec<Option<i32>> = table
            .tree(victim.pid)
            .iter()
            .map(|p| nice_of(p.pid))
            .collect();
        let s = Session::start_with(&cfg, &apps, &table, false, false).unwrap();

        assert_eq!(s.demoted_count(), 0, "demoted despite having no way back");
        assert!(!s.restore_ready);
        assert!(s.report.iter().all(|d| d.note.is_some()));
        let after: Vec<Option<i32>> = table
            .tree(victim.pid)
            .iter()
            .map(|p| nice_of(p.pid))
            .collect();
        assert_eq!(before, after, "nice values must be untouched");
    }

    #[test]
    fn dry_run_reports_targets_without_changing_them() {
        let table = proc::Table::load().unwrap();
        let apps = crate::apps::list(&table);
        let Some(victim) = apps.iter().find(|a| {
            guard::protected_match(&a.identities()).is_none() && a.uid == unsafe { libc::getuid() }
        }) else {
            return;
        };
        let cfg = Config {
            focus: crate::config::Focus {
                demote: vec![victim.name.clone()],
                prevent_sleep: false,
                nice_level: 10,
            },
            ..Default::default()
        };
        let before: Vec<Option<i32>> = table
            .tree(victim.pid)
            .iter()
            .map(|p| nice_of(p.pid))
            .collect();
        let s = Session::start_with(&cfg, &apps, &table, true, true).unwrap();

        assert_eq!(s.demoted_count(), 0);
        assert!(!s.report.is_empty(), "dry run should still report targets");
        let after: Vec<Option<i32>> = table
            .tree(victim.pid)
            .iter()
            .map(|p| nice_of(p.pid))
            .collect();
        assert_eq!(before, after);
    }

    #[test]
    fn empty_demote_list_touches_nothing() {
        let table = proc::Table::load().unwrap();
        let apps = crate::apps::list(&table);
        let s = Session::start_with(&Config::default(), &apps, &table, false, true).unwrap();
        assert_eq!(s.demoted_count(), 0);
        assert!(s.report.is_empty());
    }

    #[test]
    fn protected_apps_are_refused_not_demoted() {
        let table = proc::Table::load().unwrap();
        let apps = crate::apps::list(&table);
        let cfg = Config {
            focus: crate::config::Focus {
                demote: vec!["Finder".into()],
                prevent_sleep: false,
                nice_level: 10,
            },
            ..Default::default()
        };
        let s = Session::start_with(&cfg, &apps, &table, false, true).unwrap();
        assert_eq!(s.demoted_count(), 0, "Finder must never be demoted");
        assert!(s.report.iter().all(|d| d.note.is_some()));
        let finder = apps
            .iter()
            .find(|a| a.bundle_id.as_deref() == Some("com.apple.finder"))
            .unwrap();
        assert_eq!(
            nice_of(finder.pid),
            Some(0),
            "Finder's nice must be untouched"
        );
    }

    #[test]
    fn sleep_assertion_can_be_taken_and_released() {
        let a = Assertion::prevent_idle_sleep();
        assert!(
            a.is_some(),
            "IOKit assertion should be grantable unprivileged"
        );
        drop(a);
    }
}