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

rigtop_core/
agent.rs
