secure_exec_kernel/

process_table.rs

use crate::user::ProcessIdentity;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::error::Error;
use std::fmt;
use std::ops::{BitOr, BitOrAssign};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex, MutexGuard, WaitTimeoutResult, Weak};
use std::thread;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

const ZOMBIE_TTL: Duration = Duration::from_secs(60);
const INIT_PID: u32 = 1;
const MAX_ALLOCATED_PID: u32 = i32::MAX as u32;
pub const DEFAULT_PROCESS_UMASK: u32 = 0o022;
pub const SIGHUP: i32 = 1;
pub const SIGCHLD: i32 = 17;
pub const SIGCONT: i32 = 18;
pub const SIGSTOP: i32 = 19;
pub const SIGTSTP: i32 = 20;
pub const SIGTERM: i32 = 15;
pub const SIGKILL: i32 = 9;
pub const SIGPIPE: i32 = 13;
pub const SIGWINCH: i32 = 28;
const MAX_SIGNAL: i32 = 64;

pub type ProcessResult<T> = Result<T, ProcessTableError>;
pub type ProcessExitCallback = Arc<dyn Fn(i32) + Send + Sync + 'static>;

pub trait DriverProcess: Send + Sync {
    fn kill(&self, signal: i32);
    fn wait(&self, timeout: Duration) -> Option<i32>;
    fn set_on_exit(&self, callback: ProcessExitCallback);
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProcessTableError {
    code: &'static str,
    message: String,
}

impl ProcessTableError {
    pub fn code(&self) -> &'static str {
        self.code
    }

    fn invalid_signal(signal: i32) -> Self {
        Self {
            code: "EINVAL",
            message: format!("invalid signal {signal}"),
        }
    }

    fn no_such_process(pid: u32) -> Self {
        Self {
            code: "ESRCH",
            message: format!("no such process {pid}"),
        }
    }

    fn no_such_process_group(pgid: u32) -> Self {
        Self {
            code: "ESRCH",
            message: format!("no such process group {pgid}"),
        }
    }

    fn no_matching_child(waiter_pid: u32, pid: i32) -> Self {
        Self {
            code: "ECHILD",
            message: format!("process {waiter_pid} has no matching child for waitpid({pid})"),
        }
    }

    fn pid_space_exhausted() -> Self {
        Self {
            code: "EAGAIN",
            message: String::from("process id space exhausted"),
        }
    }

    fn permission_denied(message: impl Into<String>) -> Self {
        Self {
            code: "EPERM",
            message: message.into(),
        }
    }
}

impl fmt::Display for ProcessTableError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}: {}", self.code, self.message)
    }
}

impl Error for ProcessTableError {}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProcessStatus {
    Running,
    Stopped,
    Exited,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct SignalSet {
    bits: u64,
}

impl SignalSet {
    pub const fn empty() -> Self {
        Self { bits: 0 }
    }

    pub const fn is_empty(self) -> bool {
        self.bits == 0
    }

    pub fn from_signal(signal: i32) -> ProcessResult<Self> {
        Ok(Self {
            bits: signal_bit(signal)?,
        })
    }

    pub fn from_signals(signals: impl IntoIterator<Item = i32>) -> ProcessResult<Self> {
        let mut set = Self::empty();
        for signal in signals {
            set.insert(signal)?;
        }
        Ok(set)
    }

    pub fn contains(self, signal: i32) -> bool {
        signal_bit(signal)
            .map(|bit| self.bits & bit != 0)
            .unwrap_or(false)
    }

    pub fn insert(&mut self, signal: i32) -> ProcessResult<()> {
        self.bits |= signal_bit(signal)?;
        Ok(())
    }

    pub fn remove(&mut self, signal: i32) -> ProcessResult<()> {
        self.bits &= !signal_bit(signal)?;
        Ok(())
    }

    pub fn union(self, other: Self) -> Self {
        Self {
            bits: self.bits | other.bits,
        }
    }

    pub fn difference(self, other: Self) -> Self {
        Self {
            bits: self.bits & !other.bits,
        }
    }

    pub fn signals(self) -> Vec<i32> {
        let mut signals = Vec::new();
        for signal in 1..=MAX_SIGNAL {
            if self.contains(signal) {
                signals.push(signal);
            }
        }
        signals
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SigmaskHow {
    Block,
    Unblock,
    SetMask,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WaitPidFlags {
    bits: u32,
}

impl WaitPidFlags {
    pub const WNOHANG: Self = Self { bits: 1 << 0 };
    pub const WUNTRACED: Self = Self { bits: 1 << 1 };
    pub const WCONTINUED: Self = Self { bits: 1 << 2 };

    pub const fn empty() -> Self {
        Self { bits: 0 }
    }

    pub const fn contains(self, other: Self) -> bool {
        (self.bits & other.bits) == other.bits
    }
}

impl Default for WaitPidFlags {
    fn default() -> Self {
        Self::empty()
    }
}

impl BitOr for WaitPidFlags {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        Self {
            bits: self.bits | rhs.bits,
        }
    }
}

impl BitOrAssign for WaitPidFlags {
    fn bitor_assign(&mut self, rhs: Self) {
        self.bits |= rhs.bits;
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProcessWaitEvent {
    Exited,
    Stopped,
    Continued,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProcessWaitResult {
    pub pid: u32,
    pub status: i32,
    pub event: ProcessWaitEvent,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProcessFileDescriptors {
    pub stdin: u32,
    pub stdout: u32,
    pub stderr: u32,
}

impl Default for ProcessFileDescriptors {
    fn default() -> Self {
        Self {
            stdin: 0,
            stdout: 1,
            stderr: 2,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProcessContext {
    pub pid: u32,
    pub ppid: u32,
    pub env: BTreeMap<String, String>,
    pub cwd: String,
    pub umask: u32,
    pub fds: ProcessFileDescriptors,
    pub identity: ProcessIdentity,
    pub blocked_signals: SignalSet,
    pub pending_signals: SignalSet,
}

impl Default for ProcessContext {
    fn default() -> Self {
        Self {
            pid: 0,
            ppid: 0,
            env: BTreeMap::new(),
            cwd: String::from("/"),
            umask: DEFAULT_PROCESS_UMASK,
            fds: ProcessFileDescriptors::default(),
            identity: ProcessIdentity::default(),
            blocked_signals: SignalSet::empty(),
            pending_signals: SignalSet::empty(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProcessEntry {
    pub pid: u32,
    pub ppid: u32,
    pub pgid: u32,
    pub sid: u32,
    pub driver: String,
    pub command: String,
    pub args: Vec<String>,
    pub status: ProcessStatus,
    pub exit_code: Option<i32>,
    pub exit_time_ms: Option<u64>,
    pub env: BTreeMap<String, String>,
    pub cwd: String,
    pub umask: u32,
    pub identity: ProcessIdentity,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProcessInfo {
    pub pid: u32,
    pub ppid: u32,
    pub pgid: u32,
    pub sid: u32,
    pub driver: String,
    pub command: String,
    pub status: ProcessStatus,
    pub exit_code: Option<i32>,
    pub identity: ProcessIdentity,
}

#[derive(Clone)]
pub struct ProcessTable {
    inner: Arc<ProcessTableInner>,
}

struct ProcessTableInner {
    state: Mutex<ProcessTableState>,
    waiters: Condvar,
    reaper: Arc<ZombieReaper>,
}

struct ProcessRecord {
    entry: ProcessEntry,
    driver_process: Arc<dyn DriverProcess>,
    pending_wait_events: VecDeque<PendingWaitEvent>,
    blocked_signals: SignalSet,
    pending_signals: SignalSet,
}

struct ScheduledSignalDelivery {
    pid: u32,
    signal: i32,
    status: ProcessStatus,
    driver_process: Arc<dyn DriverProcess>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct PendingWaitEvent {
    status: i32,
    event: ProcessWaitEvent,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum WaitSelector {
    AnyChild,
    ChildPid(u32),
    ProcessGroup(u32),
}

struct ZombieReaper {
    state: Mutex<ZombieReaperState>,
    wake: Condvar,
    thread_spawns: AtomicUsize,
}

#[derive(Default)]
struct ZombieReaperState {
    deadlines: BTreeMap<u32, Instant>,
    shutdown: bool,
}

struct ProcessTableState {
    entries: BTreeMap<u32, ProcessRecord>,
    next_pid: u32,
    zombie_ttl: Duration,
    on_process_exit: Option<Arc<dyn Fn(u32) + Send + Sync + 'static>>,
    terminating_all: bool,
}

impl Default for ProcessTableState {
    fn default() -> Self {
        Self {
            entries: BTreeMap::new(),
            next_pid: 1,
            zombie_ttl: ZOMBIE_TTL,
            on_process_exit: None,
            terminating_all: false,
        }
    }
}

impl Default for ProcessTable {
    fn default() -> Self {
        let reaper = Arc::new(ZombieReaper::default());
        Self {
            inner: {
                let inner = Arc::new(ProcessTableInner {
                    state: Mutex::new(ProcessTableState::default()),
                    waiters: Condvar::new(),
                    reaper,
                });
                start_zombie_reaper(Arc::downgrade(&inner), Arc::clone(&inner.reaper));
                inner
            },
        }
    }
}

impl ProcessTable {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_zombie_ttl(zombie_ttl: Duration) -> Self {
        let table = Self::new();
        table.inner.lock_state().zombie_ttl = zombie_ttl;
        table
    }

    pub fn allocate_pid(&self) -> ProcessResult<u32> {
        let mut state = self.inner.lock_state();
        let start = normalize_next_pid(state.next_pid);
        let mut pid = start;

        loop {
            if !state.entries.contains_key(&pid) {
                state.next_pid = next_allocated_pid_after(pid);
                return Ok(pid);
            }

            pid = next_allocated_pid_after(pid);
            if pid == start {
                return Err(ProcessTableError::pid_space_exhausted());
            }
        }
    }

    pub fn set_on_process_exit(&self, callback: Option<Arc<dyn Fn(u32) + Send + Sync + 'static>>) {
        self.inner.lock_state().on_process_exit = callback;
    }

    pub fn register(
        &self,
        pid: u32,
        driver: impl Into<String>,
        command: impl Into<String>,
        args: Vec<String>,
        ctx: ProcessContext,
        driver_process: Arc<dyn DriverProcess>,
    ) -> ProcessEntry {
        let (pgid, sid) = {
            let state = self.inner.lock_state();
            match state.entries.get(&ctx.ppid) {
                Some(parent) => (parent.entry.pgid, parent.entry.sid),
                None => (pid, pid),
            }
        };

        let entry = ProcessEntry {
            pid,
            ppid: ctx.ppid,
            pgid,
            sid,
            driver: driver.into(),
            command: command.into(),
            args,
            status: ProcessStatus::Running,
            exit_code: None,
            exit_time_ms: None,
            env: ctx.env,
            cwd: ctx.cwd,
            umask: ctx.umask & 0o777,
            identity: ctx.identity,
        };

        let weak = Arc::downgrade(&self.inner);
        driver_process.set_on_exit(Arc::new(move |code| {
            if let Some(inner) = weak.upgrade() {
                mark_exited_inner(&inner, pid, code);
            }
        }));

        let mut state = self.inner.lock_state();
        state.next_pid = next_pid_after_registered(state.next_pid, pid);
        state.entries.insert(
            pid,
            ProcessRecord {
                entry: entry.clone(),
                driver_process,
                pending_wait_events: VecDeque::new(),
                blocked_signals: ctx.blocked_signals,
                pending_signals: ctx.pending_signals,
            },
        );

        entry
    }

    pub fn get(&self, pid: u32) -> Option<ProcessEntry> {
        self.inner
            .lock_state()
            .entries
            .get(&pid)
            .map(|record| record.entry.clone())
    }

    pub fn zombie_timer_count(&self) -> usize {
        self.inner.reaper.scheduled_count()
    }

    pub fn zombie_reaper_thread_spawn_count(&self) -> usize {
        self.inner.reaper.thread_spawn_count()
    }

    pub fn running_count(&self) -> usize {
        self.inner
            .lock_state()
            .entries
            .values()
            .filter(|record| record.entry.status == ProcessStatus::Running)
            .count()
    }

    pub fn mark_exited(&self, pid: u32, exit_code: i32) {
        mark_exited_inner(&self.inner, pid, exit_code);
    }

    pub fn mark_stopped(&self, pid: u32, signal: i32) {
        mark_wait_event_inner(
            &self.inner,
            pid,
            ProcessStatus::Stopped,
            PendingWaitEvent {
                status: signal,
                event: ProcessWaitEvent::Stopped,
            },
        );
    }

    pub fn mark_continued(&self, pid: u32) {
        mark_wait_event_inner(
            &self.inner,
            pid,
            ProcessStatus::Running,
            PendingWaitEvent {
                status: SIGCONT,
                event: ProcessWaitEvent::Continued,
            },
        );
    }

    pub fn waitpid(&self, pid: u32) -> ProcessResult<(u32, i32)> {
        let mut state = self.inner.lock_state();
        loop {
            let Some(record) = state.entries.get(&pid) else {
                return Err(ProcessTableError::no_such_process(pid));
            };

            if record.entry.status == ProcessStatus::Exited {
                let status = record.entry.exit_code.unwrap_or_default();
                state.entries.remove(&pid);
                drop(state);
                self.inner.reaper.cancel(pid);
                self.inner.waiters.notify_all();
                return Ok((pid, status));
            }

            state = self.inner.wait_for_state(state);
        }
    }

    pub fn waitpid_for(
        &self,
        waiter_pid: u32,
        pid: i32,
        flags: WaitPidFlags,
    ) -> ProcessResult<Option<ProcessWaitResult>> {
        let mut state = self.inner.lock_state();
        loop {
            let selector = resolve_wait_selector(&state, waiter_pid, pid)?;
            let matching_children = matching_child_pids(&state, waiter_pid, selector);
            if matching_children.is_empty() {
                return Err(ProcessTableError::no_matching_child(waiter_pid, pid));
            }

            if let Some(result) = take_waitable_event(&mut state, &matching_children, flags) {
                let should_reap = result.event == ProcessWaitEvent::Exited;
                drop(state);
                if should_reap {
                    self.inner.reaper.cancel(result.pid);
                    self.inner.waiters.notify_all();
                }
                return Ok(Some(result));
            }

            if flags.contains(WaitPidFlags::WNOHANG) {
                return Ok(None);
            }

            state = self.inner.wait_for_state(state);
        }
    }

    pub fn kill(&self, pid: i32, signal: i32) -> ProcessResult<()> {
        if !(0..=MAX_SIGNAL).contains(&signal) {
            return Err(ProcessTableError::invalid_signal(signal));
        }

        let deliveries = {
            let mut state = self.inner.lock_state();
            if pid < 0 {
                let pgid = pid.unsigned_abs();
                let grouped = state
                    .entries
                    .values()
                    .filter(|record| record.entry.pgid == pgid)
                    .map(|record| record.entry.pid)
                    .collect::<Vec<_>>();
                if grouped.is_empty() {
                    return Err(ProcessTableError::no_such_process_group(pgid));
                }
                if signal == 0 {
                    return Ok(());
                }
                collect_signal_deliveries(&mut state, &grouped, signal)?
            } else {
                let pid = pid as u32;
                let Some(record) = state.entries.get(&pid) else {
                    return Err(ProcessTableError::no_such_process(pid));
                };
                if record.entry.status == ProcessStatus::Exited || signal == 0 {
                    return Ok(());
                }
                collect_signal_deliveries(&mut state, &[pid], signal)?
            }
        };

        if signal == 0 {
            return Ok(());
        }

        deliver_signals(&self.inner, deliveries);
        Ok(())
    }

    pub fn setpgid(&self, pid: u32, pgid: u32) -> ProcessResult<()> {
        let mut state = self.inner.lock_state();
        let (current_sid, target_pgid) = {
            let Some(record) = state.entries.get(&pid) else {
                return Err(ProcessTableError::no_such_process(pid));
            };
            (record.entry.sid, if pgid == 0 { pid } else { pgid })
        };

        if target_pgid != pid {
            let mut group_exists = false;
            for record in state.entries.values() {
                if record.entry.pgid != target_pgid || record.entry.status == ProcessStatus::Exited
                {
                    continue;
                }
                if record.entry.sid != current_sid {
                    return Err(ProcessTableError::permission_denied(
                        "cannot join process group in different session",
                    ));
                }
                group_exists = true;
                break;
            }
            if !group_exists {
                return Err(ProcessTableError::permission_denied(format!(
                    "no such process group {target_pgid}"
                )));
            }
        }

        if let Some(record) = state.entries.get_mut(&pid) {
            record.entry.pgid = target_pgid;
        }
        Ok(())
    }

    pub fn getpgid(&self, pid: u32) -> ProcessResult<u32> {
        self.get(pid)
            .map(|entry| entry.pgid)
            .ok_or_else(|| ProcessTableError::no_such_process(pid))
    }

    pub fn setsid(&self, pid: u32) -> ProcessResult<u32> {
        let mut state = self.inner.lock_state();
        let Some(record) = state.entries.get_mut(&pid) else {
            return Err(ProcessTableError::no_such_process(pid));
        };

        if record.entry.pgid == pid {
            return Err(ProcessTableError::permission_denied(format!(
                "process {pid} is already a process group leader"
            )));
        }

        record.entry.sid = pid;
        record.entry.pgid = pid;
        Ok(pid)
    }

    pub fn getsid(&self, pid: u32) -> ProcessResult<u32> {
        self.get(pid)
            .map(|entry| entry.sid)
            .ok_or_else(|| ProcessTableError::no_such_process(pid))
    }

    pub fn getppid(&self, pid: u32) -> ProcessResult<u32> {
        self.get(pid)
            .map(|entry| entry.ppid)
            .ok_or_else(|| ProcessTableError::no_such_process(pid))
    }

    pub fn get_umask(&self, pid: u32) -> ProcessResult<u32> {
        self.get(pid)
            .map(|entry| entry.umask)
            .ok_or_else(|| ProcessTableError::no_such_process(pid))
    }

    pub fn set_umask(&self, pid: u32, umask: u32) -> ProcessResult<u32> {
        let mut state = self.inner.lock_state();
        let record = state
            .entries
            .get_mut(&pid)
            .ok_or_else(|| ProcessTableError::no_such_process(pid))?;
        let previous = record.entry.umask;
        record.entry.umask = umask & 0o777;
        Ok(previous)
    }

    pub fn has_process_group(&self, pgid: u32) -> bool {
        self.inner
            .lock_state()
            .entries
            .values()
            .any(|record| record.entry.pgid == pgid && record.entry.status != ProcessStatus::Exited)
    }

    pub fn list_processes(&self) -> BTreeMap<u32, ProcessInfo> {
        self.inner
            .lock_state()
            .entries
            .values()
            .map(|record| (record.entry.pid, to_process_info(&record.entry)))
            .collect()
    }

    pub fn terminate_all(&self) {
        let running = {
            let mut state = self.inner.lock_state();
            state.terminating_all = true;
            self.inner.reaper.clear();
            state
                .entries
                .values()
                .filter(|record| record.entry.status == ProcessStatus::Running)
                .map(|record| (record.entry.pid, Arc::clone(&record.driver_process)))
                .collect::<Vec<_>>()
        };

        for (_, driver) in &running {
            driver.kill(SIGTERM);
        }
        for (pid, driver) in &running {
            if let Some(exit_code) = driver.wait(Duration::from_secs(1)) {
                self.mark_exited(*pid, exit_code);
            }
        }

        let survivors = {
            let state = self.inner.lock_state();
            running
                .iter()
                .filter(|(pid, _)| {
                    state
                        .entries
                        .get(pid)
                        .map(|record| record.entry.status == ProcessStatus::Running)
                        .unwrap_or(false)
                })
                .cloned()
                .collect::<Vec<_>>()
        };

        for (_, driver) in &survivors {
            driver.kill(SIGKILL);
        }
        for (pid, driver) in &survivors {
            if let Some(exit_code) = driver.wait(Duration::from_millis(500)) {
                self.mark_exited(*pid, exit_code);
            }
        }

        self.inner.lock_state().terminating_all = false;
    }

    pub fn sigprocmask(
        &self,
        pid: u32,
        how: SigmaskHow,
        set: SignalSet,
    ) -> ProcessResult<SignalSet> {
        let (previous, deliveries) = {
            let mut state = self.inner.lock_state();
            let record = state
                .entries
                .get_mut(&pid)
                .ok_or_else(|| ProcessTableError::no_such_process(pid))?;
            let previous = record.blocked_signals;
            record.blocked_signals = match how {
                SigmaskHow::Block => previous.union(set),
                SigmaskHow::Unblock => previous.difference(set),
                SigmaskHow::SetMask => set,
            };

            let unblocked_pending = record.pending_signals.difference(record.blocked_signals);
            let deliveries = collect_pending_signal_deliveries(record, unblocked_pending)?;
            (previous, deliveries)
        };

        deliver_signals(&self.inner, deliveries);
        Ok(previous)
    }

    pub fn sigpending(&self, pid: u32) -> ProcessResult<SignalSet> {
        self.inner
            .lock_state()
            .entries
            .get(&pid)
            .map(|record| record.pending_signals)
            .ok_or_else(|| ProcessTableError::no_such_process(pid))
    }
}

fn to_process_info(entry: &ProcessEntry) -> ProcessInfo {
    ProcessInfo {
        pid: entry.pid,
        ppid: entry.ppid,
        pgid: entry.pgid,
        sid: entry.sid,
        driver: entry.driver.clone(),
        command: entry.command.clone(),
        status: entry.status,
        exit_code: entry.exit_code,
        identity: entry.identity.clone(),
    }
}

fn mark_exited_inner(inner: &Arc<ProcessTableInner>, pid: u32, exit_code: i32) {
    let (callback, zombie_ttl, should_schedule, deliveries) = {
        let mut state = inner.lock_state();
        let (ppid, pgid) = {
            let Some(record) = state.entries.get_mut(&pid) else {
                return;
            };

            if record.entry.status == ProcessStatus::Exited {
                return;
            }

            record.entry.status = ProcessStatus::Exited;
            record.entry.exit_code = Some(exit_code);
            record.entry.exit_time_ms = Some(now_ms());
            let ppid = record.entry.ppid;
            let pgid = record.entry.pgid;
            (ppid, pgid)
        };
        let mut affected_pgids = BTreeSet::from([pgid]);
        reparent_children_to_init(&mut state, pid, &mut affected_pgids);

        let orphaned_group_targets = collect_orphaned_group_signal_targets(&state, &affected_pgids);

        let should_schedule = !state.terminating_all;
        let mut deliveries = Vec::new();
        if should_schedule {
            if let Some(parent) = state
                .entries
                .get_mut(&ppid)
                .filter(|parent| parent.entry.status == ProcessStatus::Running)
            {
                if let Some(delivery) =
                    queue_or_schedule_signal(parent, SIGCHLD).expect("SIGCHLD should be valid")
                {
                    deliveries.push(delivery);
                }
            }
        }

        for target_pid in orphaned_group_targets {
            if let Some(record) = state.entries.get_mut(&target_pid) {
                if let Some(delivery) =
                    queue_or_schedule_signal(record, SIGHUP).expect("SIGHUP should be valid")
                {
                    deliveries.push(delivery);
                }
                if let Some(delivery) =
                    queue_or_schedule_signal(record, SIGCONT).expect("SIGCONT should be valid")
                {
                    deliveries.push(delivery);
                }
            }
        }

        (
            state.on_process_exit.clone(),
            state.zombie_ttl,
            should_schedule,
            deliveries,
        )
    };

    if should_schedule {
        inner.reaper.schedule(pid, zombie_ttl);
    } else {
        inner.reaper.cancel(pid);
    }

    deliver_signals(inner, deliveries);

    if let Some(on_process_exit) = callback {
        on_process_exit(pid);
    }

    inner.waiters.notify_all();
}

fn reparent_children_to_init(
    state: &mut ProcessTableState,
    exiting_pid: u32,
    affected_pgids: &mut BTreeSet<u32>,
) {
    let new_parent = reparent_target_pid(state, exiting_pid);
    for record in state.entries.values_mut() {
        if record.entry.ppid != exiting_pid {
            continue;
        }
        record.entry.ppid = new_parent;
        affected_pgids.insert(record.entry.pgid);
    }
}

fn reparent_target_pid(state: &ProcessTableState, exiting_pid: u32) -> u32 {
    if exiting_pid != INIT_PID
        && state
            .entries
            .get(&INIT_PID)
            .map(|record| record.entry.status != ProcessStatus::Exited)
            .unwrap_or(false)
    {
        INIT_PID
    } else {
        0
    }
}

fn collect_orphaned_group_signal_targets(
    state: &ProcessTableState,
    candidate_pgids: &BTreeSet<u32>,
) -> Vec<u32> {
    let mut targets = Vec::new();
    for &pgid in candidate_pgids {
        if !process_group_is_orphaned(state, pgid) || !process_group_has_stopped_member(state, pgid)
        {
            continue;
        }

        for record in state.entries.values() {
            if record.entry.pgid == pgid && record.entry.status != ProcessStatus::Exited {
                targets.push(record.entry.pid);
            }
        }
    }
    targets
}

fn process_group_is_orphaned(state: &ProcessTableState, pgid: u32) -> bool {
    let mut has_member = false;
    for record in state.entries.values() {
        if record.entry.pgid != pgid || record.entry.status == ProcessStatus::Exited {
            continue;
        }
        has_member = true;
        if has_parent_outside_group_in_same_session(state, &record.entry) {
            return false;
        }
    }

    has_member
}

fn has_parent_outside_group_in_same_session(
    state: &ProcessTableState,
    entry: &ProcessEntry,
) -> bool {
    match entry.ppid {
        0 | INIT_PID => false,
        ppid => state
            .entries
            .get(&ppid)
            .map(|parent| {
                parent.entry.status != ProcessStatus::Exited
                    && parent.entry.sid == entry.sid
                    && parent.entry.pgid != entry.pgid
            })
            .unwrap_or(false),
    }
}

fn process_group_has_stopped_member(state: &ProcessTableState, pgid: u32) -> bool {
    state
        .entries
        .values()
        .any(|record| record.entry.pgid == pgid && record.entry.status == ProcessStatus::Stopped)
}

fn mark_wait_event_inner(
    inner: &Arc<ProcessTableInner>,
    pid: u32,
    next_status: ProcessStatus,
    event: PendingWaitEvent,
) {
    let deliveries = {
        let mut state = inner.lock_state();
        let ppid = {
            let Some(record) = state.entries.get_mut(&pid) else {
                return;
            };

            if record.entry.status == ProcessStatus::Exited || record.entry.status == next_status {
                return;
            }

            record.entry.status = next_status;
            record.pending_wait_events.push_back(event);
            record.entry.ppid
        };

        state
            .entries
            .get_mut(&ppid)
            .filter(|parent| parent.entry.status == ProcessStatus::Running)
            .and_then(|parent| {
                queue_or_schedule_signal(parent, SIGCHLD)
                    .expect("SIGCHLD should be valid")
                    .into_iter()
                    .next()
            })
            .into_iter()
            .collect::<Vec<_>>()
    };

    deliver_signals(inner, deliveries);

    inner.waiters.notify_all();
}

fn signal_bit(signal: i32) -> ProcessResult<u64> {
    if !(1..=MAX_SIGNAL).contains(&signal) {
        return Err(ProcessTableError::invalid_signal(signal));
    }
    Ok(1u64 << (signal - 1))
}

fn normalize_next_pid(pid: u32) -> u32 {
    if (INIT_PID..=MAX_ALLOCATED_PID).contains(&pid) {
        pid
    } else {
        INIT_PID
    }
}

fn next_allocated_pid_after(pid: u32) -> u32 {
    if pid >= MAX_ALLOCATED_PID {
        INIT_PID
    } else {
        pid + 1
    }
}

fn next_pid_after_registered(current: u32, registered: u32) -> u32 {
    let current = normalize_next_pid(current);
    if !(INIT_PID..=MAX_ALLOCATED_PID).contains(&registered) {
        return current;
    }

    if current <= registered {
        next_allocated_pid_after(registered)
    } else {
        current
    }
}

fn signal_can_be_blocked(signal: i32) -> bool {
    !matches!(signal, SIGKILL | SIGSTOP | SIGCONT)
}

fn queue_or_schedule_signal(
    record: &mut ProcessRecord,
    signal: i32,
) -> ProcessResult<Option<ScheduledSignalDelivery>> {
    if signal_can_be_blocked(signal) && record.blocked_signals.contains(signal) {
        record.pending_signals.insert(signal)?;
        return Ok(None);
    }

    Ok(Some(ScheduledSignalDelivery {
        pid: record.entry.pid,
        signal,
        status: record.entry.status,
        driver_process: Arc::clone(&record.driver_process),
    }))
}

fn collect_signal_deliveries(
    state: &mut ProcessTableState,
    target_pids: &[u32],
    signal: i32,
) -> ProcessResult<Vec<ScheduledSignalDelivery>> {
    let mut deliveries = Vec::new();
    for pid in target_pids {
        let Some(record) = state.entries.get_mut(pid) else {
            continue;
        };
        if let Some(delivery) = queue_or_schedule_signal(record, signal)? {
            deliveries.push(delivery);
        }
    }
    Ok(deliveries)
}

fn collect_pending_signal_deliveries(
    record: &mut ProcessRecord,
    signals: SignalSet,
) -> ProcessResult<Vec<ScheduledSignalDelivery>> {
    let mut deliveries = Vec::new();
    for signal in signals.signals() {
        record.pending_signals.remove(signal)?;
        deliveries.push(ScheduledSignalDelivery {
            pid: record.entry.pid,
            signal,
            status: record.entry.status,
            driver_process: Arc::clone(&record.driver_process),
        });
    }
    Ok(deliveries)
}

fn deliver_signals(inner: &Arc<ProcessTableInner>, deliveries: Vec<ScheduledSignalDelivery>) {
    let mut stopped = Vec::new();
    let mut continued = Vec::new();

    for delivery in &deliveries {
        match delivery.signal {
            SIGSTOP | SIGTSTP if delivery.status == ProcessStatus::Running => {
                stopped.push((delivery.pid, delivery.signal))
            }
            SIGCONT if delivery.status == ProcessStatus::Stopped => continued.push(delivery.pid),
            _ => {}
        }
        delivery.driver_process.kill(delivery.signal);
    }

    for (pid, signal) in stopped {
        mark_wait_event_inner(
            inner,
            pid,
            ProcessStatus::Stopped,
            PendingWaitEvent {
                status: signal,
                event: ProcessWaitEvent::Stopped,
            },
        );
    }
    for pid in continued {
        mark_wait_event_inner(
            inner,
            pid,
            ProcessStatus::Running,
            PendingWaitEvent {
                status: SIGCONT,
                event: ProcessWaitEvent::Continued,
            },
        );
    }
}

fn resolve_wait_selector(
    state: &ProcessTableState,
    waiter_pid: u32,
    pid: i32,
) -> ProcessResult<WaitSelector> {
    let waiter = state
        .entries
        .get(&waiter_pid)
        .ok_or_else(|| ProcessTableError::no_such_process(waiter_pid))?;

    Ok(match pid {
        -1 => WaitSelector::AnyChild,
        0 => WaitSelector::ProcessGroup(waiter.entry.pgid),
        p if p < -1 => WaitSelector::ProcessGroup(p.unsigned_abs()),
        p => WaitSelector::ChildPid(p as u32),
    })
}

fn matching_child_pids(
    state: &ProcessTableState,
    waiter_pid: u32,
    selector: WaitSelector,
) -> Vec<u32> {
    state
        .entries
        .values()
        .filter(|record| record.entry.ppid == waiter_pid)
        .filter(|record| match selector {
            WaitSelector::AnyChild => true,
            WaitSelector::ChildPid(pid) => record.entry.pid == pid,
            WaitSelector::ProcessGroup(pgid) => record.entry.pgid == pgid,
        })
        .map(|record| record.entry.pid)
        .collect()
}

fn take_waitable_event(
    state: &mut ProcessTableState,
    matching_children: &[u32],
    flags: WaitPidFlags,
) -> Option<ProcessWaitResult> {
    for child_pid in matching_children {
        let mut non_exit_result = None;
        let mut should_reap = false;
        {
            let record = state.entries.get_mut(child_pid)?;
            if let Some(index) = record
                .pending_wait_events
                .iter()
                .position(|event| is_waitable_event(event.event, flags))
            {
                let event = record
                    .pending_wait_events
                    .remove(index)
                    .expect("pending wait event should exist");
                non_exit_result = Some(ProcessWaitResult {
                    pid: *child_pid,
                    status: event.status,
                    event: event.event,
                });
            } else if record.entry.status == ProcessStatus::Exited {
                should_reap = true;
            }
        }

        if let Some(result) = non_exit_result {
            return Some(result);
        }

        if should_reap {
            let record = state
                .entries
                .remove(child_pid)
                .expect("exited child should still exist");
            return Some(ProcessWaitResult {
                pid: *child_pid,
                status: record.entry.exit_code.unwrap_or_default(),
                event: ProcessWaitEvent::Exited,
            });
        }
    }

    None
}

fn is_waitable_event(event: ProcessWaitEvent, flags: WaitPidFlags) -> bool {
    match event {
        ProcessWaitEvent::Exited => true,
        ProcessWaitEvent::Stopped => flags.contains(WaitPidFlags::WUNTRACED),
        ProcessWaitEvent::Continued => flags.contains(WaitPidFlags::WCONTINUED),
    }
}

fn start_zombie_reaper(inner: Weak<ProcessTableInner>, reaper: Arc<ZombieReaper>) {
    reaper.thread_spawns.fetch_add(1, Ordering::SeqCst);
    thread::spawn(move || loop {
        let Some(pid) = reaper.take_next_due_pid() else {
            return;
        };

        let Some(inner) = inner.upgrade() else {
            return;
        };

        let mut state = inner.lock_state();
        let should_reap = state
            .entries
            .get(&pid)
            .map(|record| {
                record.entry.status == ProcessStatus::Exited
                    && !has_living_parent(&state, record.entry.ppid)
            })
            .unwrap_or(false);
        if should_reap {
            state.entries.remove(&pid);
        } else if state
            .entries
            .get(&pid)
            .map(|record| record.entry.status == ProcessStatus::Exited)
            .unwrap_or(false)
        {
            reaper.schedule(pid, state.zombie_ttl);
        }
        drop(state);
        inner.waiters.notify_all();
    });
}

fn has_living_parent(state: &ProcessTableState, ppid: u32) -> bool {
    ppid != 0
        && state
            .entries
            .get(&ppid)
            .map(|record| record.entry.status != ProcessStatus::Exited)
            .unwrap_or(false)
}

impl ProcessTableInner {
    fn lock_state(&self) -> MutexGuard<'_, ProcessTableState> {
        lock_or_recover(&self.state)
    }

    fn wait_for_state<'a>(
        &self,
        guard: MutexGuard<'a, ProcessTableState>,
    ) -> MutexGuard<'a, ProcessTableState> {
        wait_or_recover(&self.waiters, guard)
    }
}

fn now_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

impl Default for ZombieReaper {
    fn default() -> Self {
        Self {
            state: Mutex::new(ZombieReaperState::default()),
            wake: Condvar::new(),
            thread_spawns: AtomicUsize::new(0),
        }
    }
}

impl ZombieReaper {
    fn schedule(&self, pid: u32, ttl: Duration) {
        let mut state = lock_or_recover(&self.state);
        state.deadlines.insert(pid, Instant::now() + ttl);
        drop(state);
        self.wake.notify_all();
    }

    fn cancel(&self, pid: u32) {
        let mut state = lock_or_recover(&self.state);
        let removed = state.deadlines.remove(&pid).is_some();
        drop(state);
        if removed {
            self.wake.notify_all();
        }
    }

    fn clear(&self) {
        let mut state = lock_or_recover(&self.state);
        let changed = !state.deadlines.is_empty();
        state.deadlines.clear();
        drop(state);
        if changed {
            self.wake.notify_all();
        }
    }

    fn shutdown(&self) {
        let mut state = lock_or_recover(&self.state);
        state.shutdown = true;
        drop(state);
        self.wake.notify_all();
    }

    fn scheduled_count(&self) -> usize {
        lock_or_recover(&self.state).deadlines.len()
    }

    fn thread_spawn_count(&self) -> usize {
        self.thread_spawns.load(Ordering::SeqCst)
    }

    fn take_next_due_pid(&self) -> Option<u32> {
        let mut state = lock_or_recover(&self.state);
        loop {
            if state.shutdown {
                return None;
            }

            let Some((pid, deadline)) = state
                .deadlines
                .iter()
                .min_by_key(|(_, deadline)| **deadline)
                .map(|(&pid, &deadline)| (pid, deadline))
            else {
                state = wait_or_recover(&self.wake, state);
                continue;
            };

            let now = Instant::now();
            if deadline <= now {
                state.deadlines.remove(&pid);
                return Some(pid);
            }

            let timeout = deadline.saturating_duration_since(now);
            let (next_state, _) = wait_timeout_or_recover(&self.wake, state, timeout);
            state = next_state;
        }
    }
}

impl Drop for ProcessTableInner {
    fn drop(&mut self) {
        self.reaper.shutdown();
    }
}

fn lock_or_recover<'a, T>(mutex: &'a Mutex<T>) -> MutexGuard<'a, T> {
    match mutex.lock() {
        Ok(guard) => guard,
        Err(poisoned) => poisoned.into_inner(),
    }
}

fn wait_or_recover<'a, T>(condvar: &Condvar, guard: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
    match condvar.wait(guard) {
        Ok(guard) => guard,
        Err(poisoned) => poisoned.into_inner(),
    }
}

fn wait_timeout_or_recover<'a, T>(
    condvar: &Condvar,
    guard: MutexGuard<'a, T>,
    timeout: Duration,
) -> (MutexGuard<'a, T>, WaitTimeoutResult) {
    match condvar.wait_timeout(guard, timeout) {
        Ok(result) => result,
        Err(poisoned) => poisoned.into_inner(),
    }
}

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

    #[derive(Default)]
    struct TestDriverProcess {
        on_exit: Mutex<Option<ProcessExitCallback>>,
    }

    impl TestDriverProcess {
        fn exit(&self, exit_code: i32) {
            let callback = self
                .on_exit
                .lock()
                .expect("test driver lock poisoned")
                .clone();
            if let Some(callback) = callback {
                callback(exit_code);
            }
        }
    }

    impl DriverProcess for TestDriverProcess {
        fn kill(&self, _signal: i32) {}

        fn wait(&self, _timeout: Duration) -> Option<i32> {
            None
        }

        fn set_on_exit(&self, callback: ProcessExitCallback) {
            *self.on_exit.lock().expect("test driver lock poisoned") = Some(callback);
        }
    }

    fn context(ppid: u32) -> ProcessContext {
        ProcessContext {
            ppid,
            ..ProcessContext::default()
        }
    }

    #[test]
    fn allocate_pid_wraps_without_reusing_live_or_zombie_processes() {
        let table = ProcessTable::with_zombie_ttl(Duration::from_secs(3600));
        let live_high = Arc::new(TestDriverProcess::default());
        let zombie_high = Arc::new(TestDriverProcess::default());
        let live_one = Arc::new(TestDriverProcess::default());
        let max_pid = MAX_ALLOCATED_PID;

        table.register(
            max_pid - 1,
            "test",
            "live-high",
            Vec::new(),
            context(0),
            live_high,
        );
        table.register(
            max_pid,
            "test",
            "zombie-high",
            Vec::new(),
            context(0),
            zombie_high.clone(),
        );
        table.register(1, "test", "live-one", Vec::new(), context(0), live_one);
        zombie_high.exit(0);

        table.inner.lock_state().next_pid = max_pid - 1;

        assert_eq!(table.allocate_pid().expect("allocate pid"), 2);
        assert_eq!(table.allocate_pid().expect("allocate pid"), 3);
    }
}