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

rigtop_core/
agent.rs
