1use std::collections::VecDeque;
9use std::io::{BufRead, BufReader};
10use std::process::{Child, ChildStdin, Command, Stdio};
11use std::sync::mpsc::{self, Receiver, RecvTimeoutError};
12use std::sync::{Arc, Mutex};
13use std::thread;
14use std::time::{Duration, Instant};
15
16use crate::agent::{
17 AgentMessage, ClientMessage, FEATURE_PING, FEATURE_PROCESS_CONTROL, FEATURE_WATCH_PROCESSES,
18 PROTOCOL_VERSION, write_message,
19};
20use crate::snapshot::Snapshot;
21use crate::source::{DataSource, ProcessSignal, SourceError};
22
23pub const DEFAULT_REMOTE_COMMAND: &str = "rigtop --agent";
25pub const FALLBACK_REMOTE_COMMANDS: [&str; 2] =
30 ["~/.local/bin/rigtop --agent", "~/.cargo/bin/rigtop --agent"];
31const NOT_FOUND_STATUS: i32 = 127;
33
34const RECONNECT_DELAY: Duration = Duration::from_secs(5);
36const CONNECT_TIMEOUT: Duration = Duration::from_secs(20);
39const SIGNAL_TIMEOUT: Duration = Duration::from_secs(10);
41const STDERR_LINES: usize = 20;
43
44const PING_EVERY: Duration = Duration::from_secs(5);
46
47const PING_TIMEOUT: Duration = Duration::from_secs(30);
49
50#[derive(Debug, Clone, PartialEq, Eq)]
52pub struct RemoteTarget {
53 pub name: String,
55 pub program: String,
57 pub args: Vec<String>,
59 pub fallback_args: Vec<Vec<String>>,
62}
63
64impl RemoteTarget {
65 pub fn ssh(destination: &str, remote_command: &str) -> Self {
68 let options = [
69 "BatchMode=yes",
71 "ConnectTimeout=10",
72 "ServerAliveInterval=15",
74 "ServerAliveCountMax=3",
75 ];
76 let mut args = vec!["-T".to_owned()];
77 for option in options {
78 args.push("-o".to_owned());
79 args.push(option.to_owned());
80 }
81 args.push("--".to_owned());
83 args.push(destination.to_owned());
84 args.push(remote_command.to_owned());
85 let fallback_args = if remote_command == DEFAULT_REMOTE_COMMAND {
88 FALLBACK_REMOTE_COMMANDS
89 .iter()
90 .map(|command| {
91 let mut fallback = args.clone();
92 *fallback.last_mut().expect("command pushed above") = (*command).to_owned();
93 fallback
94 })
95 .collect()
96 } else {
97 Vec::new()
98 };
99 Self {
100 name: destination.to_owned(),
101 program: "ssh".to_owned(),
102 args,
103 fallback_args,
104 }
105 }
106}
107
108enum Incoming {
110 Message(AgentMessage),
111 Pong {
114 id: u64,
115 at: Instant,
116 },
117 Closed,
119}
120
121struct Connection {
123 child: Child,
124 stdin: ChildStdin,
125 incoming: Receiver<Incoming>,
126 stderr: Arc<Mutex<VecDeque<String>>>,
128 established: bool,
130}
131
132impl Connection {
133 fn spawn(target: &RemoteTarget, args: &[String]) -> Result<Self, SourceError> {
134 let mut command = Command::new(&target.program);
135 command
136 .args(args)
137 .stdin(Stdio::piped())
138 .stdout(Stdio::piped())
139 .stderr(Stdio::piped());
140 #[cfg(windows)]
141 {
142 use std::os::windows::process::CommandExt;
143 command.creation_flags(0x0800_0000);
145 }
146 let mut child = command.spawn().map_err(|err| {
147 SourceError::Disconnected(format!("cannot start {}: {err}", target.program))
148 })?;
149 let stdin = child.stdin.take().expect("piped stdin");
150 let stdout = child.stdout.take().expect("piped stdout");
151 let stderr_pipe = child.stderr.take().expect("piped stderr");
152
153 let (tx, incoming) = mpsc::channel();
154 thread::Builder::new()
155 .name("rigtop-remote-reader".to_owned())
156 .spawn(move || {
157 for line in BufReader::new(stdout).lines() {
158 let Ok(line) = line else { break };
159 let incoming = match serde_json::from_str::<AgentMessage>(&line) {
162 Ok(AgentMessage::Pong { id }) => Incoming::Pong {
163 id,
164 at: Instant::now(),
165 },
166 Ok(message) => Incoming::Message(message),
167 Err(_) => continue,
168 };
169 if tx.send(incoming).is_err() {
170 return;
171 }
172 }
173 let _ = tx.send(Incoming::Closed);
174 })
175 .map_err(SourceError::Io)?;
176
177 let stderr = Arc::new(Mutex::new(VecDeque::new()));
178 let lines = Arc::clone(&stderr);
179 thread::Builder::new()
180 .name("rigtop-remote-stderr".to_owned())
181 .spawn(move || {
182 for line in BufReader::new(stderr_pipe).lines() {
183 let Ok(line) = line else { break };
184 let mut lines = lines.lock().unwrap_or_else(|e| e.into_inner());
185 if lines.len() == STDERR_LINES {
186 lines.pop_front();
187 }
188 lines.push_back(line);
189 }
190 })
191 .map_err(SourceError::Io)?;
192
193 Ok(Self {
194 child,
195 stdin,
196 incoming,
197 stderr,
198 established: false,
199 })
200 }
201
202 fn send(&mut self, message: &ClientMessage) -> Result<(), SourceError> {
203 write_message(&mut self.stdin, message).map_err(SourceError::Io)
204 }
205
206 fn failure(&mut self) -> (String, bool) {
209 thread::sleep(Duration::from_millis(100));
211 let status = self.child.try_wait().ok().flatten();
212 let lines = self.stderr.lock().unwrap_or_else(|e| e.into_inner());
213 let detail = lines
214 .iter()
215 .rev()
216 .find(|l| !l.trim().is_empty())
217 .cloned()
218 .unwrap_or_default();
219 let not_found = status.and_then(|s| s.code()) == Some(NOT_FOUND_STATUS)
221 || detail.contains("command not found")
222 || detail.contains("not recognized");
223 let reason = match (detail.is_empty(), status) {
224 (false, _) => detail,
225 (true, Some(status)) => format!("connection closed ({status})"),
226 (true, None) => "connection closed".to_owned(),
227 };
228 (reason, not_found)
229 }
230}
231
232impl Drop for Connection {
233 fn drop(&mut self) {
234 let _ = self.child.kill();
237 let _ = self.child.wait();
238 }
239}
240
241pub struct RemoteSource {
243 target: RemoteTarget,
244 connection: Option<Connection>,
245 last_attempt: Option<Instant>,
246 interval: Duration,
247 process_limit: Option<usize>,
248 watched_pids: Vec<u32>,
250 watched_names: Vec<String>,
252 agent_watches: bool,
254 agent_controls: bool,
256 agent_pings: bool,
258 ping: Option<(u64, Instant)>,
260 last_ping: Option<Instant>,
262 link_rtt: Option<Duration>,
263 pending: VecDeque<Snapshot>,
265 agent_version: String,
267 last_error: Option<String>,
269 attempt: usize,
272 retry_now: bool,
274}
275
276impl RemoteSource {
277 pub fn new(target: RemoteTarget) -> Self {
279 Self {
280 target,
281 connection: None,
282 last_attempt: None,
283 interval: Duration::from_secs(1),
284 process_limit: None,
285 watched_pids: Vec::new(),
286 watched_names: Vec::new(),
287 agent_watches: false,
288 agent_controls: false,
289 agent_pings: false,
290 ping: None,
291 last_ping: None,
292 link_rtt: None,
293 pending: VecDeque::new(),
294 agent_version: String::new(),
295 last_error: None,
296 attempt: 0,
297 retry_now: false,
298 }
299 }
300
301 pub fn target(&self) -> &RemoteTarget {
303 &self.target
304 }
305
306 fn connection(&mut self) -> Result<&mut Connection, SourceError> {
308 if self.connection.is_none() {
309 if let Some(at) = self.last_attempt {
310 let wait = RECONNECT_DELAY.saturating_sub(at.elapsed());
311 if !wait.is_zero() {
312 let reason = self.last_error.as_deref().unwrap_or("not connected");
313 return Err(SourceError::Disconnected(format!(
314 "{}: {reason} (retrying in {}s)",
315 self.target.name,
316 wait.as_secs() + 1
317 )));
318 }
319 }
320 self.last_attempt = Some(Instant::now());
321 let args = match self.attempt {
322 0 => &self.target.args,
323 n => self
324 .target
325 .fallback_args
326 .get(n - 1)
327 .unwrap_or(&self.target.args),
328 };
329 let mut connection = match Connection::spawn(&self.target, args) {
330 Ok(connection) => connection,
331 Err(err) => {
332 self.last_error = Some(match &err {
333 SourceError::Disconnected(reason) => reason.clone(),
334 other => other.to_string(),
335 });
336 return Err(err);
337 }
338 };
339 self.agent_watches = false;
341 self.agent_controls = false;
342 self.agent_pings = false;
343 self.ping = None;
344 self.link_rtt = None;
345 let sent = connection
347 .send(&ClientMessage::SetInterval {
348 ms: self.interval.as_millis() as u64,
349 })
350 .and_then(|()| {
351 connection.send(&ClientMessage::SetProcessLimit {
352 limit: self.process_limit,
353 })
354 });
355 self.connection = Some(connection);
356 if sent.is_err() {
357 return Err(self.disconnect(None));
360 }
361 }
362 Ok(self.connection.as_mut().expect("connected above"))
363 }
364
365 fn disconnect(&mut self, reason: Option<String>) -> SourceError {
367 let (mut reason, not_found) = match (reason, self.connection.as_mut()) {
368 (Some(reason), _) => (reason, false),
369 (None, Some(connection)) => connection.failure(),
370 (None, None) => ("connection closed".to_owned(), false),
371 };
372 self.connection = None;
373 if not_found && !self.target.fallback_args.is_empty() {
374 if self.attempt < self.target.fallback_args.len() {
375 self.attempt += 1;
376 self.retry_now = true;
377 self.last_attempt = None;
378 } else {
379 self.attempt = 0;
382 reason = format!(
383 "rigtop is not on the PATH, in ~/.local/bin or in ~/.cargo/bin ({reason})"
384 );
385 }
386 }
387 let error = SourceError::Disconnected(format!("{}: {reason}", self.target.name));
388 self.last_error = Some(reason);
389 error
390 }
391
392 fn next_message(&mut self, timeout: Duration) -> Result<Option<AgentMessage>, SourceError> {
394 let deadline = Instant::now() + timeout;
395 loop {
396 let remaining = deadline.saturating_duration_since(Instant::now());
397 let connection = self.connection()?;
398 match connection.incoming.recv_timeout(remaining) {
399 Ok(Incoming::Message(message)) => return Ok(Some(message)),
400 Ok(Incoming::Pong { id, at }) => {
402 if let Some((sent_id, sent)) = self.ping
403 && sent_id == id
404 {
405 self.link_rtt = Some(at.saturating_duration_since(sent));
406 self.ping = None;
407 }
408 }
409 Ok(Incoming::Closed) | Err(RecvTimeoutError::Disconnected) => {
410 return Err(self.disconnect(None));
411 }
412 Err(RecvTimeoutError::Timeout) => return Ok(None),
413 }
414 }
415 }
416
417 fn maybe_ping(&mut self) {
419 if !self.agent_pings {
420 return;
421 }
422 let now = Instant::now();
423 let due = self
424 .last_ping
425 .is_none_or(|at| now.duration_since(at) >= PING_EVERY);
426 let waiting = self
427 .ping
428 .is_some_and(|(_, sent)| now.duration_since(sent) < PING_TIMEOUT);
429 if !due || waiting {
430 return;
431 }
432 let id = self.ping.map_or(1, |(id, _)| id + 1);
433 if let Some(connection) = self.connection.as_mut()
434 && connection.send(&ClientMessage::Ping { id }).is_ok()
435 {
436 self.ping = Some((id, now));
437 self.last_ping = Some(now);
438 }
439 }
440
441 fn send_watched(&mut self) {
445 if let Some(connection) = self.connection.as_mut() {
446 let _ = connection.send(&ClientMessage::WatchProcesses {
447 pids: self.watched_pids.clone(),
448 names: self.watched_names.clone(),
450 });
451 }
452 }
453
454 fn absorb(&mut self, message: AgentMessage) -> Result<(), SourceError> {
455 match message {
456 AgentMessage::Hello {
457 protocol, version, ..
458 } if protocol != PROTOCOL_VERSION => {
459 let reason = format!(
460 "agent speaks protocol {protocol} (rigtop {version}), this client speaks {PROTOCOL_VERSION}; \
461 install the same rigtop version on both sides"
462 );
463 self.last_attempt = Some(Instant::now());
465 Err(self.disconnect(Some(reason)))
466 }
467 AgentMessage::Hello {
468 version, features, ..
469 } => {
470 self.agent_version = version;
471 self.agent_watches = features.iter().any(|f| f == FEATURE_WATCH_PROCESSES);
472 self.agent_controls = features.iter().any(|f| f == FEATURE_PROCESS_CONTROL);
473 self.agent_pings = features.iter().any(|f| f == FEATURE_PING);
474 if self.agent_watches
475 && !(self.watched_pids.is_empty() && self.watched_names.is_empty())
476 {
477 self.send_watched();
478 }
479 Ok(())
480 }
481 AgentMessage::SignalResult { .. } | AgentMessage::Pong { .. } => Ok(()),
482 AgentMessage::Snapshot { mut snapshot } => {
483 if snapshot.host.rigtop_version.is_empty() {
485 snapshot.host.rigtop_version = self.agent_version.clone();
486 }
487 if let Some(connection) = self.connection.as_mut() {
488 connection.established = true;
489 }
490 self.last_error = None;
491 self.pending.push_back(*snapshot);
492 Ok(())
493 }
494 AgentMessage::Error { message } => Err(SourceError::Remote(message)),
495 }
496 }
497}
498
499impl RemoteSource {
500 fn poll_once(&mut self) -> Result<Snapshot, SourceError> {
502 let established = self.connection.as_ref().is_some_and(|c| c.established);
504 let timeout = if established {
505 self.interval * 3 + Duration::from_secs(5)
506 } else {
507 CONNECT_TIMEOUT
508 };
509 let deadline = Instant::now() + timeout;
510 loop {
511 while let Some(message) = self.next_message(Duration::ZERO)? {
513 self.absorb(message)?;
514 }
515 if let Some(mut newest) = self.pending.pop_back() {
516 self.pending.clear();
517 self.maybe_ping();
518 newest.link_rtt_ms = self.link_rtt.map(|rtt| rtt.as_secs_f32() * 1000.0);
520 return Ok(newest);
521 }
522 let remaining = deadline.saturating_duration_since(Instant::now());
523 if remaining.is_zero() {
524 return Err(self.disconnect(Some("no data from the agent (timed out)".to_owned())));
525 }
526 if let Some(message) = self.next_message(remaining)? {
527 self.absorb(message)?;
528 }
529 }
530 }
531}
532
533impl DataSource for RemoteSource {
534 fn poll(&mut self) -> Result<Snapshot, SourceError> {
537 loop {
538 match self.poll_once() {
539 Err(_) if std::mem::take(&mut self.retry_now) => continue,
542 result => return result,
543 }
544 }
545 }
546
547 fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError> {
548 if signal.is_control() {
549 self.connection()?;
551 if !self.agent_controls {
552 return Err(SourceError::Remote(format!(
553 "rigtop {} on the host cannot do this; it needs 0.22 or later",
554 if self.agent_version.is_empty() {
555 "(unknown)"
556 } else {
557 &self.agent_version
558 }
559 )));
560 }
561 }
562 if self
563 .connection()?
564 .send(&ClientMessage::Signal { pid, signal })
565 .is_err()
566 {
567 return Err(self.disconnect(None));
568 }
569 let deadline = Instant::now() + SIGNAL_TIMEOUT;
570 loop {
571 let remaining = deadline.saturating_duration_since(Instant::now());
572 if remaining.is_zero() {
573 return Err(SourceError::Remote(format!(
574 "no answer to the signal for {pid}"
575 )));
576 }
577 match self.next_message(remaining)? {
578 Some(AgentMessage::SignalResult {
579 pid: answered,
580 error,
581 ..
582 }) if answered == pid => {
583 return match error {
584 None => Ok(()),
585 Some(error) => Err(SourceError::Remote(error)),
586 };
587 }
588 Some(other) => self.absorb(other)?,
589 None => {}
590 }
591 }
592 }
593
594 fn set_interval(&mut self, interval: Duration) {
595 self.interval = interval;
596 if let Some(connection) = self.connection.as_mut() {
597 let _ = connection.send(&ClientMessage::SetInterval {
598 ms: interval.as_millis() as u64,
599 });
600 }
601 }
602
603 fn set_process_limit(&mut self, limit: Option<usize>) {
604 self.process_limit = limit;
605 if let Some(connection) = self.connection.as_mut() {
606 let _ = connection.send(&ClientMessage::SetProcessLimit { limit });
607 }
608 }
609
610 fn watch_processes(&mut self, pids: Vec<u32>) {
611 self.watched_pids = pids;
612 if self.agent_watches {
615 self.send_watched();
616 }
617 }
618
619 fn watch_process_names(&mut self, names: Vec<String>) {
620 self.watched_names = names;
621 if self.agent_watches {
622 self.send_watched();
623 }
624 }
625}
626
627#[cfg(test)]
628mod tests {
629 use super::*;
630
631 #[test]
632 fn ssh_target_uses_batch_mode_and_keeps_the_destination_safe() {
633 let target = RemoteTarget::ssh("zefat", DEFAULT_REMOTE_COMMAND);
634 assert_eq!(target.program, "ssh");
635 assert_eq!(target.name, "zefat");
636 assert!(target.args.windows(2).any(|w| w == ["-o", "BatchMode=yes"]));
637 let n = target.args.len();
638 assert_eq!(&target.args[n - 3..], ["--", "zefat", "rigtop --agent"]);
639 let fallbacks: Vec<&str> = target
641 .fallback_args
642 .iter()
643 .map(|args| args[n - 1].as_str())
644 .collect();
645 assert_eq!(fallbacks, FALLBACK_REMOTE_COMMANDS);
646 assert_eq!(&target.fallback_args[0][n - 3..n - 1], ["--", "zefat"]);
647 let custom = RemoteTarget::ssh("zefat", "/opt/rigtop/bin/rigtop --agent");
649 assert!(custom.fallback_args.is_empty());
650 }
651
652 #[cfg(unix)]
654 fn sh_target(script: &str, fallbacks: &[&str]) -> RemoteTarget {
655 let args = |script: &str| vec!["-c".to_owned(), script.to_owned()];
656 RemoteTarget {
657 name: "fake".to_owned(),
658 program: "sh".to_owned(),
659 args: args(script),
660 fallback_args: fallbacks.iter().map(|script| args(script)).collect(),
661 }
662 }
663
664 #[cfg(unix)]
665 #[test]
666 fn command_not_found_tries_the_fallback_at_once() {
667 let missing = "echo 'zsh:1: command not found: rigtop' >&2; exit 127";
669 let mut source = RemoteSource::new(sh_target(missing, &[missing, missing]));
670 let err = source.poll().unwrap_err().to_string();
671 assert!(
672 err.contains("not on the PATH, in ~/.local/bin or in ~/.cargo/bin"),
673 "{err}"
674 );
675 assert_eq!(source.attempt, 0, "the PATH comes first again next time");
676
677 let mut source = RemoteSource::new(sh_target(
679 "echo 'Permission denied (publickey)' >&2; exit 255",
680 &[missing],
681 ));
682 let err = source.poll().unwrap_err().to_string();
683 assert!(err.contains("Permission denied"), "{err}");
684 assert_eq!(source.attempt, 0);
685 }
686
687 #[test]
688 fn missing_program_is_reported_and_retried_later() {
689 let mut source = RemoteSource::new(RemoteTarget {
690 name: "nowhere".to_owned(),
691 program: "/nonexistent/ssh".to_owned(),
692 args: Vec::new(),
693 fallback_args: Vec::new(),
694 });
695 let err = source.poll().unwrap_err().to_string();
696 assert!(err.contains("cannot start"), "{err}");
697 let err = source.poll().unwrap_err().to_string();
699 assert!(
700 err.contains("retrying") && err.contains("cannot start"),
701 "{err}"
702 );
703 }
704
705 #[cfg(unix)]
706 #[test]
707 fn stderr_explains_an_agent_that_exits() {
708 let mut source = RemoteSource::new(RemoteTarget {
710 name: "fake".to_owned(),
711 program: "sh".to_owned(),
712 args: vec![
713 "-c".to_owned(),
714 "echo 'rigtop: command not found' >&2; exit 127".to_owned(),
715 ],
716 fallback_args: Vec::new(),
717 });
718 let err = source.poll().unwrap_err().to_string();
719 assert!(err.contains("fake: rigtop: command not found"), "{err}");
720 }
721
722 #[cfg(unix)]
723 #[test]
724 fn incompatible_protocol_is_rejected() {
725 let mut source = RemoteSource::new(RemoteTarget {
726 name: "old".to_owned(),
727 program: "sh".to_owned(),
728 args: vec![
729 "-c".to_owned(),
730 r#"echo '{"type":"hello","protocol":999,"version":"9.9.9"}'; sleep 5"#.to_owned(),
731 ],
732 fallback_args: Vec::new(),
733 });
734 let err = source.poll().unwrap_err().to_string();
735 assert!(err.contains("protocol 999"), "{err}");
736 }
737
738 #[cfg(unix)]
739 #[test]
740 fn old_agents_get_their_version_from_the_greeting() {
741 let mut source = RemoteSource::new(RemoteTarget {
743 name: "old".to_owned(),
744 program: "sh".to_owned(),
745 args: vec![
746 "-c".to_owned(),
747 r#"echo '{"type":"hello","protocol":1,"version":"0.8.2"}'
748 echo '{"type":"snapshot","snapshot":{"host":{"hostname":"nas"}}}'
749 sleep 5"#
750 .to_owned(),
751 ],
752 fallback_args: Vec::new(),
753 });
754 let snapshot = source.poll().unwrap();
755 assert_eq!(snapshot.host.hostname, "nas");
756 assert_eq!(snapshot.host.rigtop_version, "0.8.2");
757 assert_eq!(
758 crate::view::version_note(&snapshot).as_deref(),
759 Some("agent 0.8.2")
760 );
761 }
762}