rigtop_core/
agent.rs

1//! The agent protocol for remote monitoring, and the agent loop itself.
2//!
3//! A client runs `ssh host rigtop --agent`; the agent writes snapshots to
4//! stdout and reads commands from stdin. Both directions use JSON Lines: one
5//! JSON object per line, tagged by `"type"`:
6//!
7//! ```text
8//! agent  -> client  {"type":"hello","protocol":1,"version":"0.1.0","features":["watch_processes"]}
9//! agent  -> client  {"type":"snapshot","snapshot":{...}}
10//! client -> agent   {"type":"set_interval","ms":2000}
11//! client -> agent   {"type":"set_process_limit","limit":10}
12//! client -> agent   {"type":"watch_processes","pids":[1234]}
13//! client -> agent   {"type":"signal","pid":1234,"signal":"Terminate"}
14//! agent  -> client  {"type":"signal_result","pid":1234,"signal":"Terminate","error":null}
15//! ```
16//!
17//! The agent exits when stdin reaches end of file, so it never outlives the
18//! SSH connection that started it.
19
20use std::io::{self, BufRead, Write};
21use std::sync::mpsc::{self, RecvTimeoutError};
22use std::thread;
23use std::time::{Duration, Instant};
24
25use serde::{Deserialize, Serialize};
26
27use crate::snapshot::Snapshot;
28use crate::source::{DataSource, MIN_POLL_INTERVAL, ProcessSignal};
29
30/// Version of the wire format. Bumped on incompatible changes; adding
31/// fields to snapshots does not need a bump (all snapshot types default
32/// missing fields).
33pub const PROTOCOL_VERSION: u32 = 1;
34
35/// [`AgentMessage::Hello`] feature: the agent understands
36/// [`ClientMessage::WatchProcesses`] (0.19 and later). Older agents answer
37/// commands they do not know with an error, so clients check first.
38pub const FEATURE_WATCH_PROCESSES: &str = "watch_processes";
39
40/// The agent finds processes by name for [`ClientMessage::WatchProcesses`]
41/// (`names`); agents before 0.21 ignore them.
42pub const FEATURE_WATCH_PROCESS_NAMES: &str = "watch_process_names";
43
44/// The agent pauses, resumes and renices processes
45/// ([`ProcessSignal::is_control`]; 0.22 and later); older agents cannot
46/// read those signals.
47pub const FEATURE_PROCESS_CONTROL: &str = "process_control";
48
49/// Messages from the agent to the client.
50#[derive(Debug, Clone, Serialize, Deserialize)]
51#[serde(tag = "type", rename_all = "snake_case")]
52pub enum AgentMessage {
53    /// First message after start-up.
54    Hello {
55        /// [`PROTOCOL_VERSION`] of the agent.
56        protocol: u32,
57        /// rigtop version of the agent.
58        version: String,
59        /// Commands added since protocol 1 that the agent understands
60        /// (e.g. [`FEATURE_WATCH_PROCESSES`]); empty before 0.19.
61        #[serde(default)]
62        features: Vec<String>,
63    },
64    /// A new measurement.
65    Snapshot {
66        /// The measurement (processes possibly limited, see
67        /// [`ClientMessage::SetProcessLimit`]).
68        snapshot: Box<Snapshot>,
69    },
70    /// Outcome of a [`ClientMessage::Signal`].
71    SignalResult {
72        pid: u32,
73        signal: ProcessSignal,
74        /// `None` on success.
75        error: Option<String>,
76    },
77    /// A problem the agent wants to report without exiting.
78    Error { message: String },
79}
80
81/// Messages from the client to the agent.
82#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
83#[serde(tag = "type", rename_all = "snake_case")]
84pub enum ClientMessage {
85    /// Change the sampling interval.
86    SetInterval { ms: u64 },
87    /// Send only the top `limit` processes by CPU (`None`: all). Used for
88    /// hosts shown in an overview rather than in detail.
89    SetProcessLimit { limit: Option<usize> },
90    /// Always send these processes, and list them in
91    /// [`Snapshot::watched_pids`], so the client can tell when they exit.
92    /// Only for agents with [`FEATURE_WATCH_PROCESSES`]. `names`: also
93    /// watch every process they match (see
94    /// [`crate::DataSource::watch_process_names`]), for agents with
95    /// [`FEATURE_WATCH_PROCESS_NAMES`].
96    WatchProcesses {
97        pids: Vec<u32>,
98        #[serde(default, skip_serializing_if = "Vec::is_empty")]
99        names: Vec<String>,
100    },
101    /// Send a signal to a process on the agent's host.
102    Signal { pid: u32, signal: ProcessSignal },
103}
104
105/// Writes one message as a JSON line and flushes it.
106pub fn write_message<W: Write, M: Serialize>(writer: &mut W, message: &M) -> io::Result<()> {
107    serde_json::to_writer(&mut *writer, message)?;
108    writer.write_all(b"\n")?;
109    writer.flush()
110}
111
112/// What the stdin reader thread hands to the agent loop.
113enum Input {
114    Command(ClientMessage),
115    Invalid(String),
116    Closed,
117}
118
119/// Runs the agent until `input` is closed or `output` fails.
120///
121/// `source` is polled every `interval` (changeable by the client); each
122/// snapshot is written to `output`. Commands are read from `input` on a
123/// separate thread.
124pub fn run_agent<S, R, W>(
125    mut source: S,
126    mut interval: Duration,
127    input: R,
128    mut output: W,
129) -> io::Result<()>
130where
131    S: DataSource,
132    R: BufRead + Send + 'static,
133    W: Write,
134{
135    let (tx, rx) = mpsc::channel();
136    thread::Builder::new()
137        .name("rigtop-agent-input".to_owned())
138        .spawn(move || {
139            for line in input.lines() {
140                let Ok(line) = line else { break };
141                if line.trim().is_empty() {
142                    continue;
143                }
144                let message = match serde_json::from_str::<ClientMessage>(&line) {
145                    Ok(command) => Input::Command(command),
146                    Err(err) => Input::Invalid(format!("invalid command {line:?}: {err}")),
147                };
148                if tx.send(message).is_err() {
149                    return;
150                }
151            }
152            let _ = tx.send(Input::Closed);
153        })?;
154
155    let hello = AgentMessage::Hello {
156        protocol: PROTOCOL_VERSION,
157        version: env!("CARGO_PKG_VERSION").to_owned(),
158        features: vec![
159            FEATURE_WATCH_PROCESSES.to_owned(),
160            FEATURE_WATCH_PROCESS_NAMES.to_owned(),
161            FEATURE_PROCESS_CONTROL.to_owned(),
162        ],
163    };
164    // A failed write means the client went away: stop quietly.
165    if write_message(&mut output, &hello).is_err() {
166        return Ok(());
167    }
168
169    let mut process_limit: Option<usize> = None;
170    let mut next_poll = Instant::now() + MIN_POLL_INTERVAL;
171    loop {
172        let timeout = next_poll.saturating_duration_since(Instant::now());
173        let reply = match rx.recv_timeout(timeout) {
174            Ok(Input::Command(ClientMessage::SetInterval { ms })) => {
175                interval = Duration::from_millis(ms).max(MIN_POLL_INTERVAL);
176                next_poll = Instant::now() + interval;
177                None
178            }
179            Ok(Input::Command(ClientMessage::SetProcessLimit { limit })) => {
180                process_limit = limit;
181                // The source can skip work only the full list needs.
182                source.set_process_limit(limit);
183                None
184            }
185            Ok(Input::Command(ClientMessage::WatchProcesses { pids, names })) => {
186                source.watch_processes(pids);
187                source.watch_process_names(names);
188                None
189            }
190            Ok(Input::Command(ClientMessage::Signal { pid, signal })) => {
191                let error = source
192                    .signal_process(pid, signal)
193                    .err()
194                    .map(|e| e.to_string());
195                Some(AgentMessage::SignalResult { pid, signal, error })
196            }
197            Ok(Input::Invalid(message)) => Some(AgentMessage::Error { message }),
198            Ok(Input::Closed) | Err(RecvTimeoutError::Disconnected) => return Ok(()),
199            Err(RecvTimeoutError::Timeout) => {
200                let message = match source.poll() {
201                    Ok(mut snapshot) => {
202                        if let Some(limit) = process_limit {
203                            snapshot.truncate_processes(limit);
204                        }
205                        AgentMessage::Snapshot {
206                            snapshot: Box::new(snapshot),
207                        }
208                    }
209                    Err(err) => AgentMessage::Error {
210                        message: err.to_string(),
211                    },
212                };
213                next_poll = (next_poll + interval).max(Instant::now() + MIN_POLL_INTERVAL);
214                Some(message)
215            }
216        };
217        if let Some(message) = reply
218            && write_message(&mut output, &message).is_err()
219        {
220            return Ok(());
221        }
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use crate::snapshot::ProcessInfo;
229    use crate::source::SourceError;
230    use std::io::Cursor;
231    use std::sync::{Arc, Mutex};
232
233    #[derive(Default)]
234    struct FakeSource {
235        watched_pids: Vec<u32>,
236    }
237
238    impl DataSource for FakeSource {
239        fn poll(&mut self) -> Result<Snapshot, SourceError> {
240            Ok(Snapshot {
241                watched_pids: self.watched_pids.clone(),
242                processes: (1..=20)
243                    .map(|pid| ProcessInfo {
244                        pid,
245                        cpu_percent: pid as f32,
246                        ..Default::default()
247                    })
248                    .collect(),
249                process_count: 20,
250                ..Default::default()
251            })
252        }
253
254        fn signal_process(&mut self, pid: u32, _signal: ProcessSignal) -> Result<(), SourceError> {
255            if pid == 1 {
256                Err(SourceError::Disconnected("denied".to_owned()))
257            } else {
258                Ok(())
259            }
260        }
261
262        fn watch_processes(&mut self, pids: Vec<u32>) {
263            self.watched_pids = pids;
264        }
265    }
266
267    /// A writer whose contents can be read after the agent has finished.
268    #[derive(Clone, Default)]
269    struct SharedBuffer(Arc<Mutex<Vec<u8>>>);
270
271    impl Write for SharedBuffer {
272        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
273            self.0.lock().unwrap().write(buf)
274        }
275        fn flush(&mut self) -> io::Result<()> {
276            Ok(())
277        }
278    }
279
280    #[test]
281    fn messages_round_trip_as_tagged_json() {
282        let command = ClientMessage::Signal {
283            pid: 42,
284            signal: ProcessSignal::Kill,
285        };
286        let line = serde_json::to_string(&command).unwrap();
287        assert_eq!(line, r#"{"type":"signal","pid":42,"signal":"Kill"}"#);
288        assert_eq!(
289            serde_json::from_str::<ClientMessage>(&line).unwrap(),
290            command
291        );
292        let limit: ClientMessage =
293            serde_json::from_str(r#"{"type":"set_process_limit","limit":null}"#).unwrap();
294        assert_eq!(limit, ClientMessage::SetProcessLimit { limit: None });
295
296        // Names are left out when empty, as 0.19 and 0.20 clients send it.
297        let watch = ClientMessage::WatchProcesses {
298            pids: vec![7],
299            names: Vec::new(),
300        };
301        let line = serde_json::to_string(&watch).unwrap();
302        assert_eq!(line, r#"{"type":"watch_processes","pids":[7]}"#);
303        assert_eq!(serde_json::from_str::<ClientMessage>(&line).unwrap(), watch);
304        let named: ClientMessage =
305            serde_json::from_str(r#"{"type":"watch_processes","pids":[],"names":["rsync"]}"#)
306                .unwrap();
307        assert_eq!(
308            named,
309            ClientMessage::WatchProcesses {
310                pids: Vec::new(),
311                names: vec!["rsync".to_owned()]
312            }
313        );
314    }
315
316    #[test]
317    fn agent_answers_commands_and_stops_at_end_of_input() {
318        // Commands are processed before the first poll (which waits
319        // MIN_POLL_INTERVAL); end of input stops the agent.
320        let input = Cursor::new(
321            "{\"type\":\"set_process_limit\",\"limit\":3}\n\
322             {\"type\":\"signal\",\"pid\":1,\"signal\":\"Terminate\"}\n\
323             not json\n",
324        );
325        let output = SharedBuffer::default();
326        run_agent(
327            FakeSource::default(),
328            Duration::from_secs(1),
329            input,
330            output.clone(),
331        )
332        .unwrap();
333
334        let text = String::from_utf8(output.0.lock().unwrap().clone()).unwrap();
335        let messages: Vec<AgentMessage> = text
336            .lines()
337            .map(|l| serde_json::from_str(l).unwrap())
338            .collect();
339        assert!(matches!(
340            &messages[0],
341            AgentMessage::Hello {
342                protocol: PROTOCOL_VERSION,
343                features,
344                ..
345            } if features.iter().any(|f| f == FEATURE_WATCH_PROCESSES)
346        ));
347        assert!(matches!(
348            &messages[1],
349            AgentMessage::SignalResult { pid: 1, error: Some(e), .. } if e.contains("denied")
350        ));
351        assert!(
352            matches!(&messages[2], AgentMessage::Error { message } if message.contains("invalid command"))
353        );
354        assert_eq!(
355            messages.len(),
356            3,
357            "no snapshot before the first interval: {text}"
358        );
359    }
360
361    #[test]
362    fn snapshots_respect_the_process_limit() {
363        // Keep the input open (a pipe) until a snapshot has arrived.
364        let (reader, mut writer) = std::io::pipe().unwrap();
365        let output = SharedBuffer::default();
366        let out = output.clone();
367        let agent = thread::spawn(move || {
368            run_agent(
369                FakeSource::default(),
370                MIN_POLL_INTERVAL,
371                io::BufReader::new(reader),
372                out,
373            )
374        });
375        writeln!(writer, r#"{{"type":"set_process_limit","limit":3}}"#).unwrap();
376        // pid 2 uses little CPU but is watched: sent anyway.
377        writeln!(writer, r#"{{"type":"watch_processes","pids":[2]}}"#).unwrap();
378        let deadline = Instant::now() + Duration::from_secs(5);
379        let snapshot = loop {
380            let text = String::from_utf8(output.0.lock().unwrap().clone()).unwrap();
381            let found = text.lines().find_map(|l| match serde_json::from_str(l) {
382                Ok(AgentMessage::Snapshot { snapshot }) => Some(snapshot),
383                _ => None,
384            });
385            if let Some(snapshot) = found {
386                break snapshot;
387            }
388            assert!(Instant::now() < deadline, "no snapshot: {text}");
389            thread::sleep(Duration::from_millis(20));
390        };
391        let pids: Vec<u32> = snapshot.processes.iter().map(|p| p.pid).collect();
392        assert_eq!(pids, [20, 19, 18, 2], "top CPU first, then watched ones");
393        assert_eq!(snapshot.watched_pids, [2]);
394        assert_eq!(snapshot.process_count, 20);
395        drop(writer);
396        agent.join().unwrap().unwrap();
397    }
398}
rigtop_core/agent.rs

rigtop_core/
agent.rs
