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