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

rigtop_core/
agent.rs
