agent_tui_infra/infra/ipc/
process.rs

1//! Process control utilities.
2
3use std::process::Command;
4
5use sysinfo::Pid;
6use sysinfo::ProcessRefreshKind;
7use sysinfo::ProcessesToUpdate;
8use sysinfo::System;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum Signal {
12    Term,
13    Kill,
14}
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum ProcessStatus {
18    Running,
19    NotFound,
20    NoPermission,
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub struct ProcessIdentity {
25    pub pid: u32,
26    pub started_at: Option<u64>,
27}
28
29pub fn current_process_identity() -> ProcessIdentity {
30    let pid = std::process::id();
31    let controller = UnixProcessController;
32    match controller.process_identity(pid) {
33        Ok(Some(identity)) => identity,
34        Ok(None) | Err(_) => ProcessIdentity {
35            pid,
36            started_at: None,
37        },
38    }
39}
40
41pub fn check_expected_process<C: ProcessController>(
42    controller: &C,
43    expected: ProcessIdentity,
44) -> std::io::Result<ProcessStatus> {
45    let status = controller.check_process(expected.pid)?;
46    if !matches!(status, ProcessStatus::Running | ProcessStatus::NoPermission) {
47        return Ok(status);
48    }
49
50    let Some(expected_started_at) = expected.started_at else {
51        return Ok(status);
52    };
53
54    let actual = controller.process_identity(expected.pid)?;
55    match actual {
56        Some(actual) if actual.started_at == Some(expected_started_at) => Ok(status),
57        Some(_) | None => Ok(ProcessStatus::NotFound),
58    }
59}
60
61pub trait ProcessController: Send + Sync {
62    fn check_process(&self, pid: u32) -> std::io::Result<ProcessStatus>;
63
64    fn send_signal(&self, pid: u32, signal: Signal) -> std::io::Result<()>;
65
66    fn process_identity(&self, pid: u32) -> std::io::Result<Option<ProcessIdentity>>;
67}
68
69pub struct UnixProcessController;
70
71impl ProcessController for UnixProcessController {
72    fn check_process(&self, pid: u32) -> std::io::Result<ProcessStatus> {
73        let pid_t: libc::pid_t = pid.try_into().map_err(|_| {
74            std::io::Error::new(std::io::ErrorKind::InvalidInput, "PID out of range")
75        })?;
76
77        // SAFETY: `kill` with signal 0 performs a permission check without sending a signal.
78        // `pid_t` is derived from a validated u32 and is safe for libc calls.
79        let result = unsafe { libc::kill(pid_t, 0) };
80        if result == 0 {
81            if is_defunct_process(pid) {
82                return Ok(ProcessStatus::NotFound);
83            }
84            return Ok(ProcessStatus::Running);
85        }
86
87        let err = std::io::Error::last_os_error();
88        match err.raw_os_error() {
89            Some(libc::ESRCH) => Ok(ProcessStatus::NotFound),
90            Some(libc::EPERM) => Ok(ProcessStatus::NoPermission),
91            _ => Err(err),
92        }
93    }
94
95    fn send_signal(&self, pid: u32, signal: Signal) -> std::io::Result<()> {
96        let pid_t: libc::pid_t = pid.try_into().map_err(|_| {
97            std::io::Error::new(std::io::ErrorKind::InvalidInput, "PID out of range")
98        })?;
99
100        let sig = match signal {
101            Signal::Term => libc::SIGTERM,
102            Signal::Kill => libc::SIGKILL,
103        };
104
105        // SAFETY: `pid_t` is validated and `sig` is a valid libc signal constant.
106        let result = unsafe { libc::kill(pid_t, sig) };
107        if result == 0 {
108            Ok(())
109        } else {
110            Err(std::io::Error::last_os_error())
111        }
112    }
113
114    fn process_identity(&self, pid: u32) -> std::io::Result<Option<ProcessIdentity>> {
115        match self.check_process(pid)? {
116            ProcessStatus::NotFound => Ok(None),
117            ProcessStatus::Running | ProcessStatus::NoPermission => Ok(Some(ProcessIdentity {
118                pid,
119                started_at: process_started_at_secs(pid),
120            })),
121        }
122    }
123}
124
125fn is_defunct_process(pid: u32) -> bool {
126    matches!(process_status_code(pid), Some('Z' | 'X'))
127}
128
129fn process_status_code(pid: u32) -> Option<char> {
130    let output = Command::new("ps")
131        .args(["-o", "stat=", "-p", &pid.to_string()])
132        .output()
133        .ok()?;
134    if !output.status.success() {
135        return None;
136    }
137
138    String::from_utf8(output.stdout).ok()?.trim().chars().next()
139}
140
141fn process_started_at_secs(pid: u32) -> Option<u64> {
142    let pid = Pid::from_u32(pid);
143    let mut system = System::new();
144    system.refresh_processes_specifics(
145        ProcessesToUpdate::Some(&[pid]),
146        true,
147        ProcessRefreshKind::nothing(),
148    );
149    let process = system.process(pid)?;
150    let started_at = process.start_time();
151    (started_at > 0).then(|| System::boot_time().saturating_add(started_at))
152}
153
154#[allow(clippy::expect_used)]
155pub mod mock {
156    use super::*;
157    use crate::common::mutex_lock_or_recover;
158    use std::collections::HashMap;
159    use std::sync::Mutex;
160
161    #[derive(Debug, Clone, Copy)]
162    struct MockProcessRecord {
163        status: ProcessStatus,
164        started_at: Option<u64>,
165    }
166
167    pub struct MockProcessController {
168        process_states: Mutex<HashMap<u32, MockProcessRecord>>,
169        signals_sent: Mutex<Vec<(u32, Signal)>>,
170        check_error: Mutex<Option<std::io::Error>>,
171        signal_error: Mutex<Option<std::io::Error>>,
172        signal_kills_process_on: Option<Signal>,
173    }
174
175    impl Default for MockProcessController {
176        fn default() -> Self {
177            Self::new()
178        }
179    }
180
181    impl MockProcessController {
182        pub fn new() -> Self {
183            Self {
184                process_states: Mutex::new(HashMap::new()),
185                signals_sent: Mutex::new(Vec::new()),
186                check_error: Mutex::new(None),
187                signal_error: Mutex::new(None),
188                signal_kills_process_on: None,
189            }
190        }
191
192        pub fn with_signal_kills_process(mut self) -> Self {
193            self.signal_kills_process_on = Some(Signal::Term);
194            self
195        }
196
197        pub fn with_signal_kills_process_on(mut self, signal: Signal) -> Self {
198            self.signal_kills_process_on = Some(signal);
199            self
200        }
201
202        pub fn with_process(self, pid: u32, status: ProcessStatus) -> Self {
203            mutex_lock_or_recover(&self.process_states).insert(
204                pid,
205                MockProcessRecord {
206                    status,
207                    started_at: None,
208                },
209            );
210            self
211        }
212
213        pub fn with_process_identity(
214            self,
215            pid: u32,
216            status: ProcessStatus,
217            started_at: Option<u64>,
218        ) -> Self {
219            mutex_lock_or_recover(&self.process_states)
220                .insert(pid, MockProcessRecord { status, started_at });
221            self
222        }
223
224        pub fn with_check_error(self, error: std::io::Error) -> Self {
225            *mutex_lock_or_recover(&self.check_error) = Some(error);
226            self
227        }
228
229        pub fn with_signal_error(self, error: std::io::Error) -> Self {
230            *mutex_lock_or_recover(&self.signal_error) = Some(error);
231            self
232        }
233
234        pub fn signals_sent(&self) -> Vec<(u32, Signal)> {
235            mutex_lock_or_recover(&self.signals_sent).clone()
236        }
237    }
238
239    impl ProcessController for MockProcessController {
240        fn check_process(&self, pid: u32) -> std::io::Result<ProcessStatus> {
241            if let Some(err) = mutex_lock_or_recover(&self.check_error).take() {
242                return Err(err);
243            }
244            Ok(self
245                .process_states
246                .lock()
247                .unwrap_or_else(std::sync::PoisonError::into_inner)
248                .get(&pid)
249                .map(|record| record.status)
250                .unwrap_or(ProcessStatus::NotFound))
251        }
252
253        fn send_signal(&self, pid: u32, signal: Signal) -> std::io::Result<()> {
254            if let Some(err) = mutex_lock_or_recover(&self.signal_error).take() {
255                return Err(err);
256            }
257            mutex_lock_or_recover(&self.signals_sent).push((pid, signal));
258            if self.signal_kills_process_on == Some(signal) {
259                mutex_lock_or_recover(&self.process_states).remove(&pid);
260            }
261            Ok(())
262        }
263
264        fn process_identity(&self, pid: u32) -> std::io::Result<Option<ProcessIdentity>> {
265            let record = self
266                .process_states
267                .lock()
268                .unwrap_or_else(std::sync::PoisonError::into_inner)
269                .get(&pid)
270                .copied();
271            Ok(record.map(|record| ProcessIdentity {
272                pid,
273                started_at: record.started_at,
274            }))
275        }
276    }
277}
278
279#[cfg(test)]
280#[path = "process_tests.rs"]
281mod tests;
agent_tui_infra/infra/ipc/process.rs

agent_tui_infra/infra/ipc/
process.rs
