1use crate::types::{WorkerConfig, WorkerId, declared_os};
9use anyhow::{Context, Result};
10use openssh::{Session, Stdio};
11use std::collections::HashMap;
12use std::ffi::OsString;
13use std::num::NonZeroUsize;
14use std::path::{Path, PathBuf};
15use std::process::ExitStatus;
16use std::sync::Arc;
17use std::time::Duration;
18use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
19use tokio::sync::{RwLock, mpsc, watch};
20use tokio::time::Instant;
21use tracing::{debug, error, warn};
22
23pub use crate::ssh_utils::{
25 CommandResult, EnvPrefix, build_env_prefix, is_retryable_transport_error,
26 is_retryable_transport_error_text, is_valid_env_key, shell_escape_value,
27};
28
29const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
31
32const DEFAULT_COMMAND_TIMEOUT: Duration = Duration::from_secs(300);
34
35const MAX_OUTPUT_SIZE: u64 = 10 * 1024 * 1024;
37
38const HEALTH_CHECK_COMMAND: &str = "echo ok";
39
40const POOL_LIVENESS_PROBE_TIMEOUT: Duration = Duration::from_secs(10);
44
45const SSH_CLEANUP_TIMEOUT: Duration = Duration::from_secs(1);
47const SSH_READY_POLL_INTERVAL: Duration = Duration::from_millis(20);
48
49#[derive(Clone)]
50struct SshProcessState {
51 active_commands: usize,
52 idle_since: Instant,
53 draining: bool,
54 stopping: bool,
55 outcome: Option<std::result::Result<ExitStatus, String>>,
56}
57
58struct OwnedSshProcess {
62 state: watch::Sender<SshProcessState>,
63 supervisor: tokio::task::JoinHandle<()>,
64}
65
66impl OwnedSshProcess {
67 fn spawn(
68 command: &mut tokio::process::Command,
69 control_dir: Option<Arc<tempfile::TempDir>>,
70 idle_timeout: Option<Duration>,
71 ) -> Result<Self> {
72 let mut child = command.kill_on_drop(true).spawn()?;
73 let (state, mut changes) = watch::channel(SshProcessState {
74 active_commands: 0,
75 idle_since: Instant::now(),
76 draining: false,
77 stopping: false,
78 outcome: None,
79 });
80 let updates = state.clone();
81 let supervisor = tokio::spawn(async move {
82 let _control_dir = control_dir;
85 let outcome = loop {
86 let current = changes.borrow_and_update().clone();
87 if current.stopping {
88 if let Err(error) = child.start_kill() {
89 debug!("Failed to signal owned SSH child: {error}");
90 }
91 break child.wait().await;
92 }
93 let idle_deadline = idle_timeout
94 .filter(|_| current.active_commands == 0)
95 .and_then(|idle| current.idle_since.checked_add(idle));
96 tokio::select! {
97 outcome = child.wait() => break outcome,
98 _ = changes.changed() => {},
99 _ = async {
100 if let Some(deadline) = idle_deadline {
101 tokio::time::sleep_until(deadline).await;
102 } else {
103 std::future::pending::<()>().await;
104 }
105 } => {
106 updates.send_if_modified(|state| {
110 if !state.stopping && state.active_commands == 0
111 && idle_timeout.is_some_and(|idle| {
112 state.idle_since.elapsed() >= idle
113 })
114 {
115 state.stopping = true;
116 true
117 } else {
118 false
119 }
120 });
121 }
122 }
123 };
124 updates.send_modify(|state| {
125 state.stopping = true;
126 state.outcome = Some(outcome.map_err(|error| error.to_string()));
127 });
128 });
129 Ok(Self { state, supervisor })
130 }
131
132 fn is_running(&self) -> bool {
133 let state = self.state.borrow();
134 !state.draining
135 && !state.stopping
136 && state.outcome.is_none()
137 && !self.supervisor.is_finished()
138 }
139
140 fn begin_use(&self) -> Result<SshProcessUse> {
141 let mut admitted = false;
142 self.state.send_if_modified(|state| {
143 if !state.draining
144 && !state.stopping
145 && state.outcome.is_none()
146 && let Some(active) = state.active_commands.checked_add(1)
147 {
148 state.active_commands = active;
149 admitted = true;
150 }
151 admitted
152 });
153 anyhow::ensure!(admitted, "SSH connection is closing or disconnected");
154 Ok(SshProcessUse {
155 state: self.state.clone(),
156 completed: false,
157 })
158 }
159
160 fn request_stop(&self) {
161 self.state.send_modify(|state| state.stopping = true);
162 }
163
164 async fn wait_until(&self, deadline: Instant) -> Result<ExitStatus> {
165 let mut changes = self.state.subscribe();
166 loop {
167 if let Some(outcome) = changes.borrow_and_update().outcome.clone() {
168 return outcome.map_err(anyhow::Error::msg);
169 }
170 tokio::time::timeout_at(deadline, changes.changed())
171 .await
172 .context(
173 "SSH connection timed out during authentication or control socket readiness",
174 )?
175 .context("SSH child supervisor stopped unexpectedly")?;
176 }
177 }
178
179 async fn stop_before(&mut self, deadline: Instant) -> Result<()> {
180 self.request_stop();
181 tokio::time::timeout_at(deadline, &mut self.supervisor)
182 .await
183 .context("Timed out reaping owned SSH child")?
184 .context("SSH child supervisor failed")
185 }
186}
187
188impl Drop for OwnedSshProcess {
189 fn drop(&mut self) {
190 self.request_stop();
193 }
194}
195
196struct SshProcessUse {
197 state: watch::Sender<SshProcessState>,
198 completed: bool,
199}
200
201impl SshProcessUse {
202 fn complete(mut self) {
203 self.completed = true;
204 }
205}
206
207impl Drop for SshProcessUse {
208 fn drop(&mut self) {
209 self.state.send_modify(|state| {
210 state.draining |= !self.completed;
214 state.active_commands = state.active_commands.saturating_sub(1);
215 if state.active_commands == 0 {
216 state.idle_since = Instant::now();
217 state.stopping |= state.draining;
218 }
219 });
220 }
221}
222
223struct OwnedSshSession {
224 session: Session,
227 process: OwnedSshProcess,
228 _control_dir: Arc<tempfile::TempDir>,
229}
230
231fn is_expected_health_check_output(stdout: &str) -> bool {
232 stdout
233 .trim()
234 .lines()
235 .last()
236 .is_some_and(is_health_check_sentinel)
237}
238
239fn is_health_check_sentinel(line: &str) -> bool {
240 matches!(line.trim(), "ok")
241}
242
243#[derive(Debug, Clone)]
245pub struct SshOptions {
246 pub connect_timeout: Duration,
249 pub command_timeout: Duration,
251 pub server_alive_interval: Option<Duration>,
255 pub control_persist_idle: Option<Duration>,
264 pub control_master: bool,
266 pub known_hosts: KnownHostsPolicy,
268}
269
270impl Default for SshOptions {
271 fn default() -> Self {
272 Self {
273 connect_timeout: DEFAULT_CONNECT_TIMEOUT,
274 command_timeout: DEFAULT_COMMAND_TIMEOUT,
275 server_alive_interval: None,
276 control_persist_idle: None,
277 control_master: false,
281 known_hosts: KnownHostsPolicy::Add,
282 }
283 }
284}
285
286#[derive(Debug, Clone, Copy, PartialEq, Eq)]
288pub enum KnownHostsPolicy {
289 Strict,
291 Add,
293 AcceptAll,
295}
296
297#[cfg(test)]
298mod retry_tests {
299 use super::*;
300 use crate::test_guard;
301
302 #[test]
303 fn test_retryable_transport_error_text() {
304 let _guard = test_guard!();
305 assert!(is_retryable_transport_error_text(
306 "ssh: connect to host 1.2.3.4 port 22: Connection timed out"
307 ));
308 assert!(is_retryable_transport_error_text(
309 "kex_exchange_identification: Connection reset by peer"
310 ));
311 assert!(is_retryable_transport_error_text("Broken pipe"));
312 assert!(is_retryable_transport_error_text("Network is unreachable"));
313 }
314
315 #[test]
316 fn test_non_retryable_transport_error_text() {
317 let _guard = test_guard!();
318 assert!(!is_retryable_transport_error_text(
319 "Permission denied (publickey)."
320 ));
321 assert!(!is_retryable_transport_error_text(
322 "Host key verification failed."
323 ));
324 assert!(!is_retryable_transport_error_text(
325 "Could not resolve hostname worker.example.com: Name or service not known"
326 ));
327 assert!(!is_retryable_transport_error_text(
328 "Identity file /nope/id_rsa not accessible: No such file or directory"
329 ));
330 }
331}
332
333pub struct SshClient {
335 config: WorkerConfig,
337 options: SshOptions,
339 session: Option<OwnedSshSession>,
341}
342
343impl SshClient {
344 pub fn new(config: WorkerConfig, options: SshOptions) -> Self {
346 Self {
347 config,
348 options,
349 session: None,
350 }
351 }
352
353 pub fn worker_id(&self) -> &WorkerId {
355 &self.config.id
356 }
357
358 pub fn is_connected(&self) -> bool {
360 self.session
361 .as_ref()
362 .is_some_and(|session| session.process.is_running())
363 }
364
365 #[cfg(test)]
366 fn is_configured_for(&self, config: &WorkerConfig) -> bool {
367 same_ssh_endpoint(&self.config, config)
368 }
369
370 pub async fn connect(&mut self) -> Result<()> {
372 self.connect_using(Path::new("ssh"), None).await
373 }
374
375 async fn connect_using(
376 &mut self,
377 ssh_program: &Path,
378 control_directory: Option<&Path>,
379 ) -> Result<()> {
380 if self.is_connected() {
381 debug!("Already connected to {}", self.config.id);
382 return Ok(());
383 }
384 self.session.take();
386 let deadline = Instant::now()
387 .checked_add(self.options.connect_timeout)
388 .context("SSH connect_timeout exceeds the supported clock range")?;
389
390 let destination = format!("{}@{}", self.config.user, self.config.host);
391 debug!("Connecting to {} via SSH...", destination);
392
393 let session = match self
394 .connect_with_mode(
395 &destination,
396 self.options.control_master,
397 ssh_program,
398 control_directory,
399 deadline,
400 )
401 .await
402 {
403 Ok(session) => session,
404 Err(primary_error) if self.options.control_master && Instant::now() < deadline => {
405 warn!(
406 "SSH ControlMaster connection to {} failed ({}). Retrying without ControlMaster.",
407 destination, primary_error
408 );
409 self.connect_with_mode(
410 &destination,
411 false,
412 ssh_program,
413 control_directory,
414 deadline,
415 )
416 .await
417 .with_context(|| {
418 format!(
419 "Failed to connect to {} after retrying without ControlMaster",
420 destination
421 )
422 })?
423 }
424 Err(primary_error) => {
425 return Err(primary_error)
426 .with_context(|| format!("Failed to connect to {}", destination));
427 }
428 };
429
430 debug!("Connected to {} ({})", self.config.id, self.config.host);
433 self.session = Some(session);
434 Ok(())
435 }
436
437 async fn connect_with_mode(
438 &self,
439 destination: &str,
440 control_master: bool,
441 ssh_program: &Path,
442 control_directory: Option<&Path>,
443 deadline: Instant,
444 ) -> Result<OwnedSshSession> {
445 anyhow::ensure!(
446 Instant::now() < deadline,
447 "SSH connection timed out before spawn"
448 );
449 let directory = control_directory
450 .map(Path::to_path_buf)
451 .unwrap_or_else(ssh_control_directory);
452 std::fs::create_dir_all(&directory)
453 .with_context(|| format!("Failed to create SSH control directory {directory:?}"))?;
454 {
455 use std::os::unix::fs::PermissionsExt;
456 std::fs::set_permissions(&directory, std::fs::Permissions::from_mode(0o700))
457 .context("Failed to protect SSH control directory")?;
458 }
459 let control_dir = Arc::new(
460 tempfile::Builder::new()
461 .prefix("rch-")
462 .tempdir_in(directory)
463 .context("Failed to create private SSH control directory")?,
464 );
465 let socket = control_dir.path().join("master");
466 let log = control_dir.path().join("log");
467 let idle_timeout =
468 match control_persist_mode(control_master, self.options.control_persist_idle) {
469 ControlPersistMode::IdleFor(seconds) => {
470 Some(Duration::from_secs(seconds.get() as u64))
471 }
472 ControlPersistMode::TooLarge(seconds) => {
473 warn!("control_persist_idle too large ({seconds}s); closing with client");
474 None
475 }
476 ControlPersistMode::Closed => None,
477 };
478 if let Some(idle) = idle_timeout {
479 anyhow::ensure!(
480 Instant::now().checked_add(idle).is_some(),
481 "SSH idle timeout exceeds the supported clock range"
482 );
483 }
484 let can_disable_background =
488 ssh_can_disable_background(ssh_program, control_dir.clone(), deadline).await?;
489 let mut command = self.master_command(
490 ssh_program,
491 destination,
492 &socket,
493 &log,
494 can_disable_background,
495 );
496 let mut process =
497 OwnedSshProcess::spawn(&mut command, Some(control_dir.clone()), idle_timeout)
498 .with_context(|| format!("Failed to spawn SSH connection to {destination}"))?;
499 let connecting = process.begin_use()?;
502 let ready = wait_for_ssh_master(&process, ssh_program, &socket, &log, deadline).await;
503 if let Err(error) = ready {
504 let _ = process.stop_before(deadline).await;
508 return Err(error).with_context(|| format!("Failed to connect to {destination}"));
509 }
510 connecting.complete();
511 Ok(OwnedSshSession {
512 session: Session::resume(socket.into_boxed_path(), Some(log.into_boxed_path())),
513 process,
514 _control_dir: control_dir,
515 })
516 }
517
518 fn master_command(
519 &self,
520 ssh_program: &Path,
521 destination: &str,
522 socket: &Path,
523 log: &Path,
524 can_disable_background: bool,
525 ) -> tokio::process::Command {
526 let mut command = tokio::process::Command::new(ssh_program);
527 command
528 .stdin(std::process::Stdio::null())
529 .stdout(std::process::Stdio::null())
530 .stderr(std::process::Stdio::null())
531 .arg("-E")
532 .arg(log)
533 .arg("-S")
534 .arg(socket)
535 .args(["-M", "-N", "-T", "-o", "BatchMode=yes"])
536 .args(["-o", "ControlPersist=no"])
539 .arg("-o")
540 .arg(match self.options.known_hosts {
541 KnownHostsPolicy::Strict => "StrictHostKeyChecking=yes",
542 KnownHostsPolicy::Add => "StrictHostKeyChecking=accept-new",
543 KnownHostsPolicy::AcceptAll => "StrictHostKeyChecking=no",
544 })
545 .arg("-o")
546 .arg(format!(
547 "ConnectTimeout={}",
548 self.options.connect_timeout.as_secs()
549 ));
550 if can_disable_background {
551 command.args(["-o", "ForkAfterAuthentication=no"]);
552 }
553 if let Some(interval) = self.options.server_alive_interval {
554 command
555 .arg("-o")
556 .arg(format!("ServerAliveInterval={}", interval.as_secs()));
557 }
558 let identity = shellexpand::tilde(&self.config.identity_file);
559 if Path::new(identity.as_ref()).exists() {
560 command
561 .arg("-i")
562 .arg(identity.as_ref())
563 .args(["-o", "IdentitiesOnly=yes"]);
564 }
565 command.arg(destination);
566 command
567 }
568
569 pub async fn disconnect(&mut self) -> Result<()> {
571 if let Some(mut session) = self.session.take() {
572 debug!("Disconnecting from {}", self.config.id);
573 session
574 .process
575 .stop_before(Instant::now() + SSH_CLEANUP_TIMEOUT)
576 .await?;
577 debug!("Disconnected from {}", self.config.id);
579 }
580 Ok(())
581 }
582
583 pub async fn reconnect(&mut self) -> Result<()> {
588 if let Err(e) = self.disconnect().await {
589 debug!(
590 "Ignoring error closing stale session to {} before reconnect: {}",
591 self.config.id, e
592 );
593 }
594 self.connect().await
595 }
596
597 pub async fn execute(&self, command: &str) -> Result<CommandResult> {
600 self.execute_with_timeout(command, self.options.command_timeout)
601 .await
602 }
603
604 pub async fn execute_with_timeout(
610 &self,
611 command: &str,
612 command_timeout: Duration,
613 ) -> Result<CommandResult> {
614 if prefers_system_ssh(&self.config) {
621 return system_ssh_execute(
622 &self.config,
623 command,
624 command_timeout,
625 self.options
626 .server_alive_interval
627 .unwrap_or(Duration::from_secs(2)),
628 )
629 .await;
630 }
631
632 let session = self.session.as_ref().context("Not connected to worker")?;
633 let activity = session.process.begin_use()?;
634
635 let start = std::time::Instant::now();
636 debug!(
637 "Executing on {}: {}",
638 self.config.id,
639 crate::util::mask_sensitive_command(command)
640 );
641
642 let mut child = session
643 .session
644 .command("sh")
645 .arg("-c")
646 .arg(command)
647 .stdout(Stdio::piped())
648 .stderr(Stdio::piped())
649 .spawn()
650 .await
651 .with_context(|| format!("Failed to spawn command on {}", self.config.id))?;
652
653 let execution_future = async {
654 let stdout_handle = child.stdout().take();
656 let stderr_handle = child.stderr().take();
657
658 let stdout_fut = async {
659 if let Some(out) = stdout_handle {
660 let reader = BufReader::new(out);
661 let mut take = reader.take(MAX_OUTPUT_SIZE);
662 let mut buf = String::new();
663 take.read_to_string(&mut buf).await?;
664 let mut reader = take.into_inner();
666 let mut sink = tokio::io::sink();
667 tokio::io::copy(&mut reader, &mut sink).await?;
668 if buf.len() >= MAX_OUTPUT_SIZE as usize {
669 buf.push_str("\n...[output truncated]...\n");
670 }
671 Ok::<String, anyhow::Error>(buf)
672 } else {
673 Ok(String::new())
674 }
675 };
676
677 let stderr_fut = async {
678 if let Some(err) = stderr_handle {
679 let reader = BufReader::new(err);
680 let mut take = reader.take(MAX_OUTPUT_SIZE);
681 let mut buf = String::new();
682 take.read_to_string(&mut buf).await?;
683 let mut reader = take.into_inner();
685 let mut sink = tokio::io::sink();
686 tokio::io::copy(&mut reader, &mut sink).await?;
687 if buf.len() >= MAX_OUTPUT_SIZE as usize {
688 buf.push_str("\n...[output truncated]...\n");
689 }
690 Ok::<String, anyhow::Error>(buf)
691 } else {
692 Ok(String::new())
693 }
694 };
695
696 let (stdout, stderr) = tokio::try_join!(stdout_fut, stderr_fut)?;
697
698 let status = child
699 .wait()
700 .await
701 .with_context(|| "Failed to wait for command completion")?;
702
703 Ok::<_, anyhow::Error>((status, stdout, stderr))
704 };
705
706 match tokio::time::timeout(command_timeout, execution_future).await {
707 Ok(result) => {
708 let (status, stdout, stderr) = result?;
709 let duration = start.elapsed();
710 let exit_code = status.code().unwrap_or(-1);
711
712 debug!(
713 "Command completed on {} (exit={}, duration={}ms)",
714 self.config.id,
715 exit_code,
716 duration.as_millis()
717 );
718
719 activity.complete();
720 Ok(CommandResult {
721 exit_code,
722 stdout,
723 stderr,
724 duration_ms: duration.as_millis() as u64,
725 })
726 }
727 Err(_) => {
728 warn!(
733 "Command timed out on {} after {:?}",
734 self.config.id, command_timeout
735 );
736 anyhow::bail!("Command timed out after {:?}", command_timeout);
737 }
738 }
739 }
740
741 pub async fn execute_streaming<F, G>(
743 &self,
744 command: &str,
745 mut on_stdout: F,
746 mut on_stderr: G,
747 ) -> Result<CommandResult>
748 where
749 F: FnMut(&str),
750 G: FnMut(&str),
751 {
752 let session = self.session.as_ref().context("Not connected to worker")?;
753 let activity = session.process.begin_use()?;
754
755 let start = std::time::Instant::now();
756 debug!(
757 "Executing (streaming) on {}: {}",
758 self.config.id,
759 crate::util::mask_sensitive_command(command)
760 );
761
762 let mut child = session
763 .session
764 .command("sh")
765 .arg("-c")
766 .arg(command)
767 .stdout(Stdio::piped())
768 .stderr(Stdio::piped())
769 .spawn()
770 .await
771 .with_context(|| format!("Failed to spawn command on {}", self.config.id))?;
772
773 let stdout = child.stdout().take();
774 let stderr = child.stderr().take();
775
776 let (tx, mut rx) = mpsc::channel(100);
779
780 if let Some(out) = stdout {
782 let tx = tx.clone();
783 tokio::spawn(async move {
784 let mut reader = BufReader::new(out);
785 let mut line = String::new();
786 loop {
787 line.clear();
788 match reader.read_line(&mut line).await {
789 Ok(0) => break, Ok(_) => {
791 if tx.send(StreamEvent::Stdout(line.clone())).await.is_err() {
792 break; }
794 }
795 Err(_) => break, }
797 }
798 });
799 }
800
801 if let Some(err) = stderr {
803 let tx = tx.clone();
804 tokio::spawn(async move {
805 let mut reader = BufReader::new(err);
806 let mut line = String::new();
807 loop {
808 line.clear();
809 match reader.read_line(&mut line).await {
810 Ok(0) => break, Ok(_) => {
812 if tx.send(StreamEvent::Stderr(line.clone())).await.is_err() {
813 break; }
815 }
816 Err(_) => break, }
818 }
819 });
820 }
821
822 drop(tx);
824
825 let mut stdout_acc = String::new();
826 let mut stderr_acc = String::new();
827
828 enum StreamEvent {
829 Stdout(String),
830 Stderr(String),
831 }
832
833 let streaming_future = async {
834 while let Some(event) = rx.recv().await {
836 match event {
837 StreamEvent::Stdout(line) => {
838 on_stdout(&line);
839 if stdout_acc.len() < MAX_OUTPUT_SIZE as usize {
840 stdout_acc.push_str(&line);
841 if stdout_acc.len() >= MAX_OUTPUT_SIZE as usize {
842 stdout_acc.push_str("\n...[output truncated]...\n");
843 }
844 }
845 }
846 StreamEvent::Stderr(line) => {
847 on_stderr(&line);
848 if stderr_acc.len() < MAX_OUTPUT_SIZE as usize {
849 stderr_acc.push_str(&line);
850 if stderr_acc.len() >= MAX_OUTPUT_SIZE as usize {
851 stderr_acc.push_str("\n...[output truncated]...\n");
852 }
853 }
854 }
855 }
856 }
857
858 let status = child.wait().await?;
859 Ok::<_, anyhow::Error>(status)
860 };
861
862 match tokio::time::timeout(self.options.command_timeout, streaming_future).await {
863 Ok(result) => {
864 let status = result?;
865 let duration = start.elapsed();
866 let exit_code = status.code().unwrap_or(-1);
867
868 activity.complete();
869 Ok(CommandResult {
870 exit_code,
871 stdout: stdout_acc,
872 stderr: stderr_acc,
873 duration_ms: duration.as_millis() as u64,
874 })
875 }
876 Err(_) => {
877 warn!(
881 "Command (streaming) timed out on {} after {:?}, cleaning up",
882 self.config.id, self.options.command_timeout
883 );
884 anyhow::bail!("Command timed out after {:?}", self.options.command_timeout);
885 }
886 }
887 }
888
889 pub async fn health_check(&self) -> Result<bool> {
891 self.health_check_with_timeout(self.options.command_timeout)
892 .await
893 }
894
895 pub async fn health_check_with_timeout(&self, timeout: Duration) -> Result<bool> {
898 match self
899 .execute_with_timeout(HEALTH_CHECK_COMMAND, timeout)
900 .await
901 {
902 Ok(result) => Ok(result.success() && is_expected_health_check_output(&result.stdout)),
903 Err(e) => {
904 warn!("Health check failed for {}: {}", self.config.id, e);
905 Ok(false)
906 }
907 }
908 }
909}
910
911fn ssh_control_directory() -> PathBuf {
912 if let Some(home) = dirs::home_dir() {
915 home.join(".ssh").join("rch")
916 } else if let Some(runtime) = dirs::runtime_dir() {
917 runtime.join("rch-ssh")
918 } else {
919 let username = whoami::username().unwrap_or_else(|_| "unknown".to_owned());
920 std::env::temp_dir().join(format!("rch-ssh-{username}"))
921 }
922}
923
924async fn ssh_can_disable_background(
925 ssh_program: &Path,
926 control_dir: Arc<tempfile::TempDir>,
927 deadline: Instant,
928) -> Result<bool> {
929 use std::io::Read;
930
931 anyhow::ensure!(
932 Instant::now() < deadline,
933 "SSH connection timed out before capability probe"
934 );
935 let configuration = control_dir.path().join("client-config");
936 let mut command = tokio::process::Command::new(ssh_program);
937 command
938 .args(["-G", "-F", "/dev/null", "-T", "none"])
939 .stdin(std::process::Stdio::null())
940 .stdout(std::fs::File::create(&configuration)?)
941 .stderr(std::process::Stdio::null());
942 let mut probe = OwnedSshProcess::spawn(&mut command, Some(control_dir.clone()), None)?;
943 let checked = probe.wait_until(deadline).await;
944 let _ = probe.stop_before(deadline).await;
945 anyhow::ensure!(checked?.success(), "SSH client capability probe failed");
946 let mut output = String::new();
947 std::fs::File::open(configuration)?
948 .take(65_537)
949 .read_to_string(&mut output)?;
950 anyhow::ensure!(
951 output.len() <= 65_536,
952 "SSH client capability response is too large"
953 );
954 anyhow::ensure!(
955 output.lines().any(|line| {
956 let mut words = line.split_whitespace();
957 words.next() == Some("hostname") && words.next() == Some("none")
958 }),
959 "SSH client returned an invalid capability response"
960 );
961 Ok(output.lines().any(|line| {
962 line.split_whitespace()
963 .next()
964 .is_some_and(|key| key.eq_ignore_ascii_case("forkafterauthentication"))
965 }))
966}
967
968async fn wait_for_ssh_master(
969 process: &OwnedSshProcess,
970 ssh_program: &Path,
971 socket: &Path,
972 log: &Path,
973 deadline: Instant,
974) -> Result<()> {
975 loop {
976 let outcome = process.state.borrow().outcome.clone();
977 if let Some(outcome) = outcome {
978 let status = outcome.map_err(anyhow::Error::msg)?;
979 let diagnostic = ssh_connect_diagnostic(log);
980 anyhow::bail!("SSH master exited before becoming ready ({status}): {diagnostic}");
981 }
982 anyhow::ensure!(
983 Instant::now() < deadline,
984 "SSH connection timed out during authentication or control socket readiness"
985 );
986 if socket.exists() {
987 let mut command = tokio::process::Command::new(ssh_program);
991 command
992 .stdin(std::process::Stdio::null())
993 .stdout(std::process::Stdio::null())
994 .stderr(std::process::Stdio::null())
995 .arg("-S")
996 .arg(socket)
997 .args(["-o", "BatchMode=yes", "-O", "check", "none"]);
998 let mut probe = OwnedSshProcess::spawn(&mut command, None, None)?;
999 let checked = probe.wait_until(deadline).await;
1000 let _ = probe.stop_before(deadline).await;
1001 if checked?.success() && process.is_running() {
1002 return Ok(());
1003 }
1004 }
1005 tokio::time::sleep_until(deadline.min(Instant::now() + SSH_READY_POLL_INTERVAL)).await;
1006 }
1007}
1008
1009fn ssh_connect_diagnostic(log: &Path) -> String {
1010 use std::io::Read;
1011
1012 let mut diagnostic = Vec::new();
1013 if let Ok(file) = std::fs::File::open(log) {
1014 let _ = file.take(16 * 1024).read_to_end(&mut diagnostic);
1015 }
1016 String::from_utf8_lossy(&diagnostic).trim().to_owned()
1017}
1018
1019fn pooled_client_options(pool_options: &SshOptions, config: &WorkerConfig) -> SshOptions {
1041 let mut options = pool_options.clone();
1042 if declared_os(&config.tags).as_deref() == Some("windows") {
1043 options.control_master = false;
1044 }
1045 options
1046}
1047
1048pub(crate) fn prefers_system_ssh(config: &WorkerConfig) -> bool {
1070 declared_os(&config.tags).as_deref() == Some("windows")
1071}
1072
1073#[derive(Debug, Clone, PartialEq, Eq)]
1075pub struct RemoteShellCommand {
1076 pub command: String,
1078 pub stdin_script: Option<String>,
1081}
1082
1083#[must_use]
1088pub fn remote_shell_command(config: &WorkerConfig, command: &str) -> RemoteShellCommand {
1089 if prefers_system_ssh(config) {
1090 RemoteShellCommand {
1091 command: "sh -s".into(),
1092 stdin_script: (command != "sh -s").then(|| {
1093 format!(
1094 "exec sh -c {} </dev/null\n",
1095 shell_escape::escape(command.into())
1096 )
1097 }),
1098 }
1099 } else {
1100 RemoteShellCommand {
1101 command: command.into(),
1102 stdin_script: None,
1103 }
1104 }
1105}
1106
1107pub use crate::ssh_utils::identities_only_args;
1108
1109pub(crate) fn system_ssh_argv(
1133 config: &WorkerConfig,
1134 command: &RemoteShellCommand,
1135 server_alive_interval: Duration,
1136) -> Vec<OsString> {
1137 let identity_path = shellexpand::tilde(&config.identity_file);
1138 let destination = format!("{}@{}", config.user, config.host);
1139
1140 let mut argv: Vec<OsString> = Vec::with_capacity(12);
1141 argv.push(OsString::from("ssh"));
1142 argv.push(OsString::from("-i"));
1143 argv.push(OsString::from(identity_path.as_ref()));
1144 if let Some(opts) = identities_only_args(&config.identity_file) {
1145 argv.extend(opts.map(OsString::from));
1146 }
1147 argv.push(OsString::from("-o"));
1148 argv.push(OsString::from("BatchMode=yes"));
1149 argv.push(OsString::from("-o"));
1150 argv.push(OsString::from("ConnectTimeout=8"));
1151 argv.push(OsString::from("-o"));
1152 argv.push(OsString::from("StrictHostKeyChecking=accept-new"));
1153 if !server_alive_interval.is_zero() {
1154 let secs = server_alive_interval.as_secs().max(1);
1158 argv.push(OsString::from("-o"));
1159 argv.push(OsString::from(format!("ServerAliveInterval={secs}")));
1160 }
1161 argv.push(OsString::from(destination));
1162 argv.push(OsString::from(&command.command));
1163 argv
1164}
1165
1166pub async fn write_ssh_command_stdin(
1173 mut stdin: tokio::process::ChildStdin,
1174 input: &str,
1175) -> std::io::Result<()> {
1176 use tokio::io::AsyncWriteExt;
1177
1178 let result = async {
1179 stdin.write_all(input.as_bytes()).await?;
1180 stdin.shutdown().await
1181 }
1182 .await;
1183 match result {
1184 Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),
1187 result => result,
1188 }
1189}
1190
1191pub(crate) async fn system_ssh_execute(
1206 config: &WorkerConfig,
1207 command: &str,
1208 command_timeout: Duration,
1209 server_alive_interval: Duration,
1210) -> Result<CommandResult> {
1211 use std::process::Stdio;
1212 use tokio::io::AsyncReadExt;
1213 use tokio::process::Command;
1214
1215 let prepared = remote_shell_command(config, command);
1216 let argv = system_ssh_argv(config, &prepared, server_alive_interval);
1217 let input = prepared.stdin_script;
1218 let args = argv.into_iter().skip(1);
1220
1221 let start = std::time::Instant::now();
1222 debug!(
1223 "system-ssh executing on {}: {}",
1224 config.id,
1225 crate::util::mask_sensitive_command(command)
1226 );
1227
1228 let mut child = Command::new("ssh")
1229 .args(args)
1230 .env_clear()
1231 .stdin(if input.is_some() {
1232 Stdio::piped()
1233 } else {
1234 Stdio::null()
1235 })
1236 .stdout(Stdio::piped())
1237 .stderr(Stdio::piped())
1238 .kill_on_drop(true)
1239 .spawn()
1240 .with_context(|| {
1241 format!(
1242 "failed to spawn system ssh for {} (Windows fallback)",
1243 config.id
1244 )
1245 })?;
1246
1247 let execution = async {
1248 let stdout = child.stdout.take();
1249 let stderr = child.stderr.take();
1250 let stdin = child.stdin.take();
1251
1252 let stdin_fut = async {
1253 if let (Some(stdin), Some(input)) = (stdin, input) {
1254 write_ssh_command_stdin(stdin, &input).await?;
1255 }
1256 Ok::<(), anyhow::Error>(())
1257 };
1258
1259 let stdout_fut = async {
1260 let mut buf = String::new();
1261 if let Some(out) = stdout {
1262 let mut handle = out.take(MAX_OUTPUT_SIZE);
1263 handle.read_to_string(&mut buf).await?;
1264 }
1265 Ok::<String, anyhow::Error>(buf)
1266 };
1267 let stderr_fut = async {
1268 let mut buf = String::new();
1269 if let Some(err) = stderr {
1270 let mut handle = err.take(MAX_OUTPUT_SIZE);
1271 handle.read_to_string(&mut buf).await?;
1272 }
1273 Ok::<String, anyhow::Error>(buf)
1274 };
1275
1276 let ((), stdout, stderr) = tokio::try_join!(stdin_fut, stdout_fut, stderr_fut)?;
1277 let status = child
1278 .wait()
1279 .await
1280 .with_context(|| "system ssh command failed to wait for completion")?;
1281 Ok::<_, anyhow::Error>((status, stdout, stderr))
1282 };
1283
1284 match tokio::time::timeout(command_timeout, execution).await {
1285 Ok(result) => {
1286 let (status, stdout, stderr) = result?;
1287 let duration = start.elapsed();
1288 let exit_code = status.code().unwrap_or(-1);
1289 debug!(
1290 "system-ssh command completed on {} (exit={}, duration={}ms)",
1291 config.id,
1292 exit_code,
1293 duration.as_millis()
1294 );
1295 Ok(CommandResult {
1296 exit_code,
1297 stdout,
1298 stderr,
1299 duration_ms: duration.as_millis() as u64,
1300 })
1301 }
1302 Err(_) => {
1303 warn!(
1308 "system-ssh command timed out on {} after {:?}",
1309 config.id, command_timeout
1310 );
1311 anyhow::bail!("Command timed out after {:?}", command_timeout);
1312 }
1313 }
1314}
1315
1316fn same_ssh_endpoint(left: &WorkerConfig, right: &WorkerConfig) -> bool {
1317 left.id == right.id
1318 && left.host == right.host
1319 && left.user == right.user
1320 && left.identity_file == right.identity_file
1321 && declared_os(&left.tags) == declared_os(&right.tags)
1322}
1323
1324#[derive(Clone)]
1325struct PooledSshClient {
1326 config: WorkerConfig,
1330 client: Arc<RwLock<SshClient>>,
1331}
1332
1333pub struct SshPool {
1335 connections: Arc<RwLock<HashMap<WorkerId, PooledSshClient>>>,
1337 options: SshOptions,
1339}
1340
1341impl SshPool {
1342 pub fn new(options: SshOptions) -> Self {
1344 Self {
1345 connections: Arc::new(RwLock::new(HashMap::new())),
1346 options,
1347 }
1348 }
1349
1350 pub async fn get_or_connect(&self, config: &WorkerConfig) -> Result<Arc<RwLock<SshClient>>> {
1357 self.get_or_connect_probing_within(config, POOL_LIVENESS_PROBE_TIMEOUT)
1358 .await
1359 }
1360
1361 async fn get_or_connect_probing_within(
1371 &self,
1372 config: &WorkerConfig,
1373 probe_timeout: Duration,
1374 ) -> Result<Arc<RwLock<SshClient>>> {
1375 let shared_client = self.get_or_create_client_entry(config).await;
1376
1377 let reusable = {
1381 let guard = shared_client.read().await;
1382 guard.is_connected()
1383 && guard
1384 .health_check_with_timeout(probe_timeout)
1385 .await
1386 .unwrap_or(false)
1387 };
1388 if reusable {
1389 debug!("Reusing live connection to {}", config.id);
1390 return Ok(shared_client);
1391 }
1392
1393 let mut client_guard = shared_client.write().await;
1398 if !client_guard.is_connected() {
1399 client_guard.connect().await?;
1400 } else if !client_guard
1401 .health_check_with_timeout(probe_timeout)
1402 .await
1403 .unwrap_or(false)
1404 {
1405 warn!(
1406 "Pooled SSH connection to {} failed liveness probe; reconnecting",
1407 config.id
1408 );
1409 client_guard.reconnect().await?;
1410 }
1411 drop(client_guard);
1413
1414 Ok(shared_client)
1415 }
1416
1417 async fn get_or_create_client_entry(&self, config: &WorkerConfig) -> Arc<RwLock<SshClient>> {
1418 let worker_id = config.id.clone();
1419
1420 loop {
1421 let existing_entry = {
1422 let connections = self.connections.read().await;
1423 connections.get(&worker_id).cloned()
1424 };
1425
1426 if let Some(entry) = existing_entry {
1427 if same_ssh_endpoint(&entry.config, config) {
1428 return entry.client;
1429 }
1430
1431 let replacement = Arc::new(RwLock::new(SshClient::new(
1432 config.clone(),
1433 pooled_client_options(&self.options, config),
1434 )));
1435 let replaced = {
1436 let mut connections = self.connections.write().await;
1437 if connections
1438 .get(&worker_id)
1439 .is_some_and(|current| Arc::ptr_eq(¤t.client, &entry.client))
1440 {
1441 connections.insert(
1442 worker_id.clone(),
1443 PooledSshClient {
1444 config: config.clone(),
1445 client: replacement.clone(),
1446 },
1447 );
1448 true
1449 } else {
1450 false
1451 }
1452 };
1453
1454 if replaced {
1455 debug!(
1456 "Replaced SSH connection entry for {} after endpoint config changed",
1457 worker_id
1458 );
1459 return replacement;
1460 }
1461
1462 continue;
1463 }
1464
1465 let new_client = Arc::new(RwLock::new(SshClient::new(
1466 config.clone(),
1467 pooled_client_options(&self.options, config),
1468 )));
1469 let inserted = {
1470 let mut connections = self.connections.write().await;
1471 if connections.contains_key(&worker_id) {
1472 false
1473 } else {
1474 connections.insert(
1475 worker_id.clone(),
1476 PooledSshClient {
1477 config: config.clone(),
1478 client: new_client.clone(),
1479 },
1480 );
1481 true
1482 }
1483 };
1484
1485 if inserted {
1486 return new_client;
1487 }
1488 }
1489 }
1490
1491 pub async fn close(&self, worker_id: &WorkerId) -> Result<()> {
1493 let client = {
1494 let mut connections = self.connections.write().await;
1495 connections.remove(worker_id).map(|entry| entry.client)
1496 };
1497
1498 if let Some(client) = client {
1499 let mut client = client.write().await;
1500 client.disconnect().await?;
1501 }
1502
1503 Ok(())
1504 }
1505
1506 pub async fn close_all(&self) -> Result<()> {
1508 let clients: Vec<_> = {
1509 let mut connections = self.connections.write().await;
1510 connections.drain().map(|(_, entry)| entry.client).collect()
1511 };
1512
1513 for client in clients {
1514 let mut client = client.write().await;
1515 if let Err(e) = client.disconnect().await {
1516 error!("Error closing connection: {}", e);
1517 }
1518 }
1519
1520 Ok(())
1521 }
1522
1523 pub async fn active_connections(&self) -> usize {
1525 self.connections.read().await.len()
1526 }
1527
1528 pub async fn run_with_timeout(
1548 &self,
1549 config: &WorkerConfig,
1550 command: &str,
1551 command_timeout: Duration,
1552 ) -> Result<CommandResult> {
1553 let client = self
1554 .get_or_connect_probing_within(config, command_timeout.min(POOL_LIVENESS_PROBE_TIMEOUT))
1555 .await?;
1556 let guard = client.read().await;
1560 guard.execute_with_timeout(command, command_timeout).await
1561 }
1562}
1563
1564impl Default for SshPool {
1565 fn default() -> Self {
1566 Self::new(SshOptions::default())
1567 }
1568}
1569
1570#[cfg(test)]
1571mod tests {
1572 use super::*;
1573 use crate::test_guard;
1574
1575 #[test]
1576 fn test_command_result_success() {
1577 let _guard = test_guard!();
1578 let result = CommandResult {
1579 exit_code: 0,
1580 stdout: "output".to_string(),
1581 stderr: String::new(),
1582 duration_ms: 100,
1583 };
1584 assert!(result.success());
1585
1586 let failed = CommandResult {
1587 exit_code: 1,
1588 stdout: String::new(),
1589 stderr: "error".to_string(),
1590 duration_ms: 50,
1591 };
1592 assert!(!failed.success());
1593 }
1594
1595 #[test]
1596 fn test_ssh_options_default() {
1597 let _guard = test_guard!();
1598 let options = SshOptions::default();
1599 assert_eq!(options.connect_timeout, Duration::from_secs(10));
1600 assert_eq!(options.command_timeout, Duration::from_secs(300));
1601 assert!(options.server_alive_interval.is_none());
1602 assert!(options.control_persist_idle.is_none());
1603 assert!(!options.control_master);
1604 }
1605
1606 #[test]
1607 fn test_ssh_client_creation() {
1608 let _guard = test_guard!();
1609 let config = WorkerConfig {
1610 id: WorkerId::new("test-worker"),
1611 host: "192.168.1.100".to_string(),
1612 user: "ubuntu".to_string(),
1613 identity_file: "~/.ssh/id_rsa".to_string(),
1614 total_slots: 8,
1615 priority: 100,
1616 tags: vec!["rust".to_string()],
1617 tools: Vec::new(),
1618 };
1619
1620 let client = SshClient::new(config.clone(), SshOptions::default());
1621 assert_eq!(client.worker_id().as_str(), "test-worker");
1622 assert!(!client.is_connected());
1623 }
1624
1625 #[test]
1626 fn test_expected_health_check_output_accepts_sentinel_as_last_line() {
1627 let _guard = test_guard!();
1628
1629 assert!(is_expected_health_check_output("ok\n"));
1630 assert!(is_expected_health_check_output("login banner\nok\n"));
1631 assert!(!is_expected_health_check_output(""));
1632 assert!(!is_expected_health_check_output("not ok\n"));
1633 assert!(!is_expected_health_check_output("ok\npost-command noise\n"));
1634 }
1635
1636 fn worker_config(id: &str, host: &str, user: &str, identity_file: &str) -> WorkerConfig {
1637 WorkerConfig {
1638 id: WorkerId::new(id),
1639 host: host.to_string(),
1640 user: user.to_string(),
1641 identity_file: identity_file.to_string(),
1642 total_slots: 8,
1643 priority: 100,
1644 tags: vec!["rust".to_string()],
1645 tools: Vec::new(),
1646 }
1647 }
1648
1649 #[test]
1650 fn test_ssh_client_configured_for_ignores_scheduling_fields() {
1651 let _guard = test_guard!();
1652 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1653 let client = SshClient::new(config.clone(), SshOptions::default());
1654
1655 let mut scheduling_only_change = config;
1656 scheduling_only_change.total_slots = 16;
1657 scheduling_only_change.priority = 250;
1658 scheduling_only_change.tags = vec!["rust".to_string(), "gpu".to_string()];
1659
1660 assert!(client.is_configured_for(&scheduling_only_change));
1661 }
1662
1663 #[test]
1664 fn test_ssh_client_configured_for_detects_endpoint_changes() {
1665 let _guard = test_guard!();
1666 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1667 let client = SshClient::new(config, SshOptions::default());
1668
1669 assert!(!client.is_configured_for(&worker_config(
1670 "worker-a",
1671 "192.168.1.101",
1672 "ubuntu",
1673 "~/.ssh/id_rsa",
1674 )));
1675 assert!(!client.is_configured_for(&worker_config(
1676 "worker-a",
1677 "192.168.1.100",
1678 "admin",
1679 "~/.ssh/id_rsa",
1680 )));
1681 assert!(!client.is_configured_for(&worker_config(
1682 "worker-a",
1683 "192.168.1.100",
1684 "ubuntu",
1685 "~/.ssh/other_key",
1686 )));
1687 }
1688
1689 #[tokio::test]
1690 async fn test_ssh_pool_reuses_matching_disconnected_entry() {
1691 let _guard = test_guard!();
1692 let pool = SshPool::default();
1693 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1694
1695 let first = pool.get_or_create_client_entry(&config).await;
1696 let second = pool.get_or_create_client_entry(&config).await;
1697
1698 assert!(Arc::ptr_eq(&first, &second));
1699 assert_eq!(pool.active_connections().await, 1);
1700 }
1701
1702 #[tokio::test]
1703 async fn test_ssh_pool_reuses_single_entry_and_close_all_drops_to_zero() {
1704 let _guard = test_guard!();
1710 let options = SshOptions {
1711 control_master: true,
1712 control_persist_idle: Some(Duration::from_secs(60)),
1713 ..Default::default()
1714 };
1715 let pool = SshPool::new(options);
1716 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1717
1718 for _ in 0..5 {
1719 let _entry = pool.get_or_create_client_entry(&config).await;
1720 }
1721 assert_eq!(
1722 pool.active_connections().await,
1723 1,
1724 "repeated borrows of one worker must reuse a single pool entry"
1725 );
1726
1727 pool.close_all().await.expect("close_all should succeed");
1728 assert_eq!(
1729 pool.active_connections().await,
1730 0,
1731 "close_all must drop all pooled connections"
1732 );
1733 }
1734
1735 #[test]
1736 fn test_pool_options_map_to_bounded_idle_persist() {
1737 let _guard = test_guard!();
1741 assert_eq!(
1742 control_persist_mode(true, Some(Duration::from_secs(60))),
1743 ControlPersistMode::IdleFor(NonZeroUsize::new(60).unwrap()),
1744 "pool mux options must keep a warm master for a bounded idle window"
1745 );
1746 }
1747
1748 fn windows_worker_config(id: &str) -> WorkerConfig {
1753 let mut config = worker_config(id, "100.68.2.11", "jeffr", "~/.ssh/surfacebookje_key");
1754 config.tags = vec!["rust".to_string(), crate::types::os_tag("windows")];
1755 config
1756 }
1757
1758 #[test]
1759 fn test_pooled_options_disable_control_master_for_windows_workers() {
1760 let _guard = test_guard!();
1766 let pool_mux = SshOptions {
1767 control_master: true,
1768 control_persist_idle: Some(Duration::from_secs(60)),
1769 ..Default::default()
1770 };
1771 let config = windows_worker_config("wsurf");
1772
1773 let effective = pooled_client_options(&pool_mux, &config);
1774 assert!(
1775 !effective.control_master,
1776 "a Windows worker must never get a pooled ControlMaster"
1777 );
1778 assert_eq!(effective.connect_timeout, pool_mux.connect_timeout);
1782 assert_eq!(effective.command_timeout, pool_mux.command_timeout);
1783 assert_eq!(
1784 effective.server_alive_interval,
1785 pool_mux.server_alive_interval
1786 );
1787 assert_eq!(
1788 effective.control_persist_idle,
1789 pool_mux.control_persist_idle
1790 );
1791 assert_eq!(effective.known_hosts, pool_mux.known_hosts);
1792
1793 let pool_plain = SshOptions {
1795 control_master: false,
1796 ..pool_mux
1797 };
1798 assert!(!pooled_client_options(&pool_plain, &config).control_master);
1799
1800 let mut mixed_case = windows_worker_config("wsurf-2");
1803 mixed_case.tags = vec![crate::types::os_tag("Windows")];
1804 assert!(!pooled_client_options(&pool_mux, &mixed_case).control_master);
1805 }
1806
1807 #[test]
1808 fn test_pooled_options_keep_pool_options_verbatim_for_non_windows_workers() {
1809 let _guard = test_guard!();
1813 let pool_mux = SshOptions {
1814 control_master: true,
1815 control_persist_idle: Some(Duration::from_secs(60)),
1816 ..Default::default()
1817 };
1818 let linux = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1819 let unlabelled = worker_config("contabo-b", "1.2.3.5", "root", "~/.ssh/id_rsa");
1820
1821 for config in [&linux, &unlabelled] {
1822 let effective = pooled_client_options(&pool_mux, config);
1823 assert!(
1824 effective.control_master,
1825 "non-Windows workers keep the pool's mux setting"
1826 );
1827 assert_eq!(
1828 effective.control_persist_idle,
1829 pool_mux.control_persist_idle
1830 );
1831 }
1832
1833 let pool_plain = SshOptions {
1834 control_master: false,
1835 ..pool_mux
1836 };
1837 for config in [&linux, &unlabelled] {
1838 assert!(!pooled_client_options(&pool_plain, config).control_master);
1839 }
1840 }
1841
1842 #[tokio::test]
1843 async fn test_pool_entry_for_windows_worker_is_never_mux() {
1844 let _guard = test_guard!();
1850 let pool = SshPool::new(SshOptions {
1851 control_master: true,
1852 control_persist_idle: Some(Duration::from_secs(60)),
1853 ..Default::default()
1854 });
1855 let config = windows_worker_config("wsurf");
1856
1857 let entry = pool.get_or_create_client_entry(&config).await;
1858 let guard = entry.read().await;
1859 assert!(
1860 !guard.options.control_master,
1861 "pooled entry for a Windows worker must be non-mux"
1862 );
1863 assert_eq!(
1866 guard.options.control_persist_idle,
1867 Some(Duration::from_secs(60))
1868 );
1869 }
1870
1871 #[tokio::test]
1872 async fn test_pool_entry_for_linux_worker_keeps_mux() {
1873 let _guard = test_guard!();
1874 let pool = SshPool::new(SshOptions {
1875 control_master: true,
1876 control_persist_idle: Some(Duration::from_secs(60)),
1877 ..Default::default()
1878 });
1879 let config = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1880
1881 let entry = pool.get_or_create_client_entry(&config).await;
1882 let guard = entry.read().await;
1883 assert!(
1884 guard.options.control_master,
1885 "pooled entry for a Linux worker must keep the pool's mux setting"
1886 );
1887 }
1888
1889 #[tokio::test]
1890 async fn test_ssh_pool_replaces_stale_entry_when_endpoint_changes() {
1891 let _guard = test_guard!();
1892 let pool = SshPool::default();
1893 let old_config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1894 let new_config = worker_config("worker-a", "192.168.1.101", "admin", "~/.ssh/new_key");
1895
1896 let stale = pool.get_or_create_client_entry(&old_config).await;
1897 let replacement = pool.get_or_create_client_entry(&new_config).await;
1898
1899 assert!(!Arc::ptr_eq(&stale, &replacement));
1900 assert_eq!(pool.active_connections().await, 1);
1901
1902 let replacement_guard = replacement.read().await;
1903 assert!(replacement_guard.is_configured_for(&new_config));
1904 }
1905
1906 #[tokio::test]
1907 async fn test_ssh_pool_retarget_does_not_wait_for_old_authentication_lock() {
1908 let _guard = test_guard!();
1909 let pool = SshPool::default();
1910 let old_config = worker_config("worker-a", "old-endpoint", "builder", "~/.ssh/id_rsa");
1911 let new_config = worker_config("worker-a", "new-endpoint", "builder", "~/.ssh/id_rsa");
1912 let old = pool.get_or_create_client_entry(&old_config).await;
1913 let authenticating = old.write().await;
1914 let unchanged = tokio::time::timeout(
1915 Duration::from_millis(100),
1916 pool.get_or_create_client_entry(&old_config),
1917 )
1918 .await
1919 .expect("immutable identity lookup must not wait on authentication");
1920 assert!(Arc::ptr_eq(&old, &unchanged));
1921 let replacement = tokio::time::timeout(
1922 Duration::from_millis(100),
1923 pool.get_or_create_client_entry(&new_config),
1924 )
1925 .await
1926 .expect("new endpoint must not wait on the old endpoint's authentication");
1927 assert!(!Arc::ptr_eq(&old, &replacement));
1928 assert_eq!(pool.active_connections().await, 1);
1929 assert!(replacement.read().await.is_configured_for(&new_config));
1930 drop(authenticating);
1931 }
1932
1933 #[tokio::test]
1934 async fn test_ssh_pool_os_transition_replaces_transport_but_case_and_scheduling_reuse() {
1935 let _guard = test_guard!();
1936 let pool = SshPool::new(SshOptions {
1937 control_master: true,
1938 ..Default::default()
1939 });
1940 let mut config = worker_config("worker-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1941 let unlabelled = pool.get_or_create_client_entry(&config).await;
1942 config.tags.push(crate::types::os_tag("windows"));
1943 let windows = pool.get_or_create_client_entry(&config).await;
1944 assert!(!Arc::ptr_eq(&unlabelled, &windows));
1945 {
1946 let client = windows.read().await;
1947 assert!(prefers_system_ssh(&client.config));
1948 assert!(!client.options.control_master);
1949 }
1950
1951 config.tags = vec!["gpu".into(), "os: WiNdOwS ".into()];
1952 config.total_slots += 1;
1953 config.priority += 1;
1954 let reused = pool.get_or_create_client_entry(&config).await;
1955 assert!(Arc::ptr_eq(&windows, &reused));
1956
1957 config.tags = vec![crate::types::os_tag("linux")];
1958 let linux = pool.get_or_create_client_entry(&config).await;
1959 assert!(!Arc::ptr_eq(&windows, &linux));
1960 {
1961 let client = linux.read().await;
1962 assert!(!prefers_system_ssh(&client.config));
1963 assert!(client.options.control_master);
1964 }
1965 assert!(prefers_system_ssh(&windows.read().await.config));
1968 assert!(unlabelled.read().await.options.control_master);
1969 assert_eq!(pool.active_connections().await, 1);
1970
1971 config.tags.clear();
1972 let removed = pool.get_or_create_client_entry(&config).await;
1973 assert!(!Arc::ptr_eq(&linux, &removed));
1974 assert!(removed.read().await.options.control_master);
1975 assert_eq!(pool.active_connections().await, 1);
1976 }
1977
1978 #[tokio::test]
1979 async fn test_health_check_reports_not_alive_without_session() {
1980 let _guard = test_guard!();
1986 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1987 let client = SshClient::new(config, SshOptions::default());
1988 assert!(!client.is_connected());
1989 let alive = client
1990 .health_check()
1991 .await
1992 .expect("health_check maps execution errors to Ok(false), never Err");
1993 assert!(!alive, "a client with no session must not report as alive");
1994 }
1995
1996 #[test]
1997 fn test_build_env_prefix_quotes_and_rejects() {
1998 let _guard = test_guard!();
1999 let mut env = HashMap::new();
2000 env.insert("RUSTFLAGS".to_string(), "-C target-cpu=native".to_string());
2001 env.insert("QUOTED".to_string(), "a'b".to_string());
2002 env.insert("BADVAL".to_string(), "line1\nline2".to_string());
2003
2004 let allowlist = vec![
2005 "RUSTFLAGS".to_string(),
2006 "QUOTED".to_string(),
2007 "MISSING".to_string(),
2008 "BADVAL".to_string(),
2009 "BAD=KEY".to_string(),
2010 ];
2011
2012 let prefix = build_env_prefix(&allowlist, |key| env.get(key).cloned());
2013
2014 assert!(prefix.prefix.contains("RUSTFLAGS='-C target-cpu=native'"));
2015 assert!(prefix.prefix.contains("QUOTED='a'\\''b'"));
2017 assert!(!prefix.prefix.contains("MISSING="));
2018 assert!(!prefix.prefix.contains("BADVAL="));
2019 assert!(prefix.rejected.contains(&"BADVAL".to_string()));
2020 assert!(prefix.rejected.contains(&"BAD=KEY".to_string()));
2021 assert_eq!(
2022 prefix.applied,
2023 vec!["RUSTFLAGS".to_string(), "QUOTED".to_string()]
2024 );
2025 }
2026
2027 mod proptest_ssh_escaping {
2032 use super::*;
2033 use proptest::prelude::*;
2034 use std::collections::HashMap;
2035
2036 proptest! {
2037 #![proptest_config(ProptestConfig::with_cases(1000))]
2038
2039 #[test]
2041 fn test_is_valid_env_key_no_panic(s in ".*") {
2042 let _guard = test_guard!();
2043 let _ = is_valid_env_key(&s);
2044 }
2045
2046 #[test]
2048 fn test_is_valid_env_key_accepts_valid(
2049 first in "[a-zA-Z_]",
2050 rest in "[a-zA-Z0-9_]{0,50}"
2051 ) {
2052 let _guard = test_guard!();
2053 let key = format!("{}{}", first, rest);
2054 prop_assert!(is_valid_env_key(&key), "Should accept valid key: {}", key);
2055 }
2056
2057 #[test]
2059 fn test_is_valid_env_key_rejects_digit_start(
2060 digit in "[0-9]",
2061 rest in "[a-zA-Z0-9_]{0,20}"
2062 ) {
2063 let _guard = test_guard!();
2064 let key = format!("{}{}", digit, rest);
2065 prop_assert!(!is_valid_env_key(&key), "Should reject digit-start key: {}", key);
2066 }
2067
2068 #[test]
2070 fn test_shell_escape_value_no_panic(s in ".*") {
2071 let _guard = test_guard!();
2072 let _ = shell_escape_value(&s);
2073 }
2074
2075 #[test]
2077 fn test_shell_escape_value_rejects_unsafe(
2078 prefix in "[a-zA-Z0-9 ]{0,10}",
2079 bad_char in "[\n\r\0]",
2080 suffix in "[a-zA-Z0-9 ]{0,10}"
2081 ) {
2082 let _guard = test_guard!();
2083 let value = format!("{}{}{}", prefix, bad_char, suffix);
2084 prop_assert!(shell_escape_value(&value).is_none(),
2085 "Should reject value with unsafe char: {:?}", value);
2086 }
2087
2088 #[test]
2090 fn test_shell_escape_value_accepts_safe(s in "[a-zA-Z0-9 !@#$%^&*()_+=\\-\\[\\]{}|;:,./<>?]{0,100}") {
2091 let result = shell_escape_value(&s);
2093 prop_assert!(result.is_some(), "Should accept safe value: {:?}", s);
2094
2095 let escaped = match result {
2098 Some(escaped) => escaped,
2099 None => {
2100 prop_assert!(false, "Should accept safe value: {:?}", s);
2101 String::new()
2102 }
2103 };
2104 if s.chars().any(|c| !c.is_ascii_alphanumeric() && c != '_') {
2105 prop_assert!(escaped.starts_with('\'') || escaped.contains('\''),
2107 "Value with special chars should be quoted: {:?} -> {:?}", s, escaped);
2108 }
2109 }
2110
2111 #[test]
2113 fn test_shell_escape_value_escapes_quotes(
2114 prefix in "[a-zA-Z0-9]{0,10}",
2115 suffix in "[a-zA-Z0-9]{0,10}"
2116 ) {
2117 let _guard = test_guard!();
2118 let value = format!("{}'{}", prefix, suffix);
2119 let result = shell_escape_value(&value);
2120 prop_assert!(result.is_some());
2121
2122 let escaped = match result {
2123 Some(escaped) => escaped,
2124 None => {
2125 prop_assert!(false, "Should escape single quote: {}", value);
2126 String::new()
2127 }
2128 };
2129 prop_assert!(escaped.contains("'\\''"),
2131 "Should escape single quote: {} -> {}", value, escaped);
2132 }
2133
2134 #[test]
2136 fn test_build_env_prefix_no_panic(
2137 keys in prop::collection::vec("[a-zA-Z_][a-zA-Z0-9_]{0,10}", 0..10),
2138 values in prop::collection::vec(".*", 0..10)
2139 ) {
2140 let mut env = HashMap::new();
2141 for (i, key) in keys.iter().enumerate() {
2142 if let Some(val) = values.get(i) {
2143 env.insert(key.clone(), val.clone());
2144 }
2145 }
2146
2147 let allowlist: Vec<String> = keys;
2148 let _ = build_env_prefix(&allowlist, |k| env.get(k).cloned());
2149 }
2150
2151 #[test]
2153 fn test_build_env_prefix_rejects_invalid_keys(
2154 invalid_key in "[0-9][a-zA-Z0-9_]{0,10}" ) {
2157 let _guard = test_guard!();
2158 let mut env = HashMap::new();
2159 env.insert(invalid_key.clone(), "value".to_string());
2160
2161 let allowlist = vec![invalid_key.clone()];
2162 let prefix = build_env_prefix(&allowlist, |k| env.get(k).cloned());
2163
2164 prop_assert!(!is_valid_env_key(&invalid_key),
2166 "Key should be invalid: {}", invalid_key);
2167 prop_assert!(prefix.rejected.contains(&invalid_key),
2168 "Should reject invalid key: {}", invalid_key);
2169 prop_assert!(prefix.prefix.is_empty());
2170 }
2171
2172 #[test]
2174 fn test_build_env_prefix_missing_values(
2175 keys in prop::collection::vec("[A-Z_][A-Z0-9_]{0,10}", 1..5)
2176 ) {
2177 let env: HashMap<String, String> = HashMap::new();
2179 let prefix = build_env_prefix(&keys, |k| env.get(k).cloned());
2180
2181 prop_assert!(prefix.prefix.is_empty(), "Should be empty when no values");
2183 prop_assert!(prefix.applied.is_empty());
2184 prop_assert!(prefix.rejected.is_empty());
2186 }
2187 }
2188
2189 #[test]
2191 fn test_shell_escape_edge_cases() {
2192 let _guard = test_guard!();
2193 let result = shell_escape_value("");
2195 assert_eq!(result, Some("''".to_string()));
2196
2197 let result = shell_escape_value("'");
2199 assert_eq!(result, Some("''\\'''".to_string()));
2200
2201 let result = shell_escape_value("'''");
2203 assert_eq!(
2205 result
2206 .as_deref()
2207 .map(|escaped| escaped.matches("'\\''").count()),
2208 Some(3)
2209 );
2210
2211 let result = shell_escape_value("日本語");
2213 assert!(result.is_some());
2214
2215 let result = shell_escape_value("🔥🚀");
2217 assert!(result.is_some());
2218
2219 let result = shell_escape_value("it's a \"test\" with $vars");
2221 assert!(result.is_some());
2222 }
2223
2224 #[test]
2225 fn test_is_valid_env_key_edge_cases() {
2226 let _guard = test_guard!();
2227 assert!(!is_valid_env_key(""));
2229
2230 assert!(is_valid_env_key("_"));
2232
2233 assert!(is_valid_env_key("A"));
2235
2236 assert!(is_valid_env_key("PATH"));
2238 assert!(is_valid_env_key("HOME"));
2239 assert!(is_valid_env_key("RUSTFLAGS"));
2240 assert!(is_valid_env_key("CC"));
2241 assert!(is_valid_env_key("_PRIVATE"));
2242 assert!(is_valid_env_key("MY_VAR_123"));
2243
2244 assert!(!is_valid_env_key("1VAR"));
2246 assert!(!is_valid_env_key("123"));
2247
2248 assert!(!is_valid_env_key("MY-VAR"));
2250 assert!(!is_valid_env_key("MY.VAR"));
2251 assert!(!is_valid_env_key("MY VAR"));
2252 assert!(!is_valid_env_key("MY=VAR"));
2253
2254 assert!(!is_valid_env_key("日本語"));
2256 assert!(!is_valid_env_key("VAR🔥"));
2257 }
2258
2259 #[test]
2260 fn test_build_env_prefix_integration() {
2261 let _guard = test_guard!();
2262 let mut env = HashMap::new();
2264 env.insert("VALID".to_string(), "simple".to_string());
2265 env.insert("WITH_QUOTE".to_string(), "it's here".to_string());
2266 env.insert("NEWLINE".to_string(), "line1\nline2".to_string());
2267 env.insert("UNICODE".to_string(), "日本語".to_string());
2268 env.insert("EMPTY".to_string(), String::new());
2269 env.insert("123INVALID".to_string(), "value".to_string());
2270
2271 let allowlist = vec![
2272 "VALID".to_string(),
2273 "WITH_QUOTE".to_string(),
2274 "NEWLINE".to_string(),
2275 "UNICODE".to_string(),
2276 "EMPTY".to_string(),
2277 "123INVALID".to_string(),
2278 "MISSING".to_string(),
2279 ];
2280
2281 let prefix = build_env_prefix(&allowlist, |k| env.get(k).cloned());
2282
2283 assert!(prefix.applied.contains(&"VALID".to_string()));
2285 assert!(prefix.prefix.contains("VALID=simple"));
2287
2288 assert!(prefix.applied.contains(&"WITH_QUOTE".to_string()));
2290
2291 assert!(prefix.rejected.contains(&"NEWLINE".to_string()));
2293
2294 assert!(prefix.applied.contains(&"UNICODE".to_string()));
2296
2297 assert!(prefix.applied.contains(&"EMPTY".to_string()));
2299
2300 assert!(prefix.rejected.contains(&"123INVALID".to_string()));
2302
2303 assert!(!prefix.applied.contains(&"MISSING".to_string()));
2305 assert!(!prefix.rejected.contains(&"MISSING".to_string()));
2306 }
2307
2308 #[test]
2309 fn test_shell_escape_roundtrip_safety() {
2310 let _guard = test_guard!();
2311 let test_values = [
2313 "simple",
2314 "with spaces",
2315 "with\ttab",
2316 "special!@#$%^&*()",
2317 "quoted\"value",
2318 "path/to/file",
2319 "-flag",
2320 "--long-flag=value",
2321 "",
2322 ];
2323
2324 for value in &test_values {
2325 let escaped = shell_escape_value(value);
2326 assert!(escaped.is_some(), "Should escape: {:?}", value);
2327 }
2328 }
2329 }
2330
2331 fn windows_worker(id: &str) -> WorkerConfig {
2336 let mut config = worker_config(id, "100.68.2.11", "jeffr", "~/.ssh/surfacebookje_key");
2337 config.tags = vec!["rust".to_string(), crate::types::os_tag("windows")];
2338 config
2339 }
2340
2341 #[test]
2342 fn test_prefers_system_ssh_for_windows_tag() {
2343 let _guard = test_guard!();
2346 let cfg = windows_worker("wsurf");
2347 assert!(prefers_system_ssh(&cfg));
2348 }
2349
2350 #[test]
2351 fn test_prefers_system_ssh_case_normalized() {
2352 let _guard = test_guard!();
2355 let mut cfg = windows_worker("wsurf-mixed");
2356 cfg.tags = vec![crate::types::os_tag("Windows")];
2357 assert!(prefers_system_ssh(&cfg));
2358
2359 let mut upper = windows_worker("wsurf-upper");
2360 upper.tags = vec![crate::types::os_tag("WINDOWS")];
2361 assert!(prefers_system_ssh(&upper));
2362 }
2363
2364 #[test]
2365 fn test_prefers_system_ssh_false_for_linux() {
2366 let _guard = test_guard!();
2369 let mut linux = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
2370 linux.tags = vec!["rust".to_string(), crate::types::os_tag("linux")];
2371 assert!(!prefers_system_ssh(&linux));
2372
2373 let unlabelled = worker_config("contabo-b", "1.2.3.5", "root", "~/.ssh/id_rsa");
2374 assert!(!prefers_system_ssh(&unlabelled));
2375 }
2376
2377 #[test]
2378 fn test_system_ssh_argv_basic_windows() {
2379 let _guard = test_guard!();
2384 let cfg = windows_worker("wsurf");
2385 let command = remote_shell_command(&cfg, "uname -a");
2386 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(2));
2387 let s: Vec<String> = argv
2388 .iter()
2389 .map(|o| o.to_string_lossy().into_owned())
2390 .collect();
2391
2392 assert_eq!(s[0], "ssh", "argv[0] is the program name");
2393 assert_eq!(s[1], "-i");
2394 assert!(
2397 s[2].ends_with(".ssh/surfacebookje_key"),
2398 "identity file is the tilde-expanded path: {}",
2399 s[2]
2400 );
2401 assert!(s.contains(&"-o".to_string()));
2404 assert!(s.contains(&"BatchMode=yes".to_string()));
2405 assert!(s.contains(&"ConnectTimeout=8".to_string()));
2406 assert!(s.contains(&"StrictHostKeyChecking=accept-new".to_string()));
2407 assert!(s.contains(&"ServerAliveInterval=2".to_string()));
2408 assert!(!s.iter().any(|arg| arg == "-p"));
2410 assert!(s.contains(&"jeffr@100.68.2.11".to_string()));
2412 assert_eq!(s[s.len() - 1], "sh -s", "Windows gets only a fixed reader");
2413 }
2414
2415 #[test]
2416 fn test_system_ssh_script_transport_keeps_posix_argv_and_exact_stdin_reader() {
2417 let _guard = test_guard!();
2418 let linux = worker_config("linux", "1.2.3.4", "root", "~/.ssh/id_rsa");
2419 let windows = windows_worker("windows");
2420 let script = "printf '%s\\n' '%PATH% !VAR! & | < > ^ $()'\nprintf done";
2421 let prepared = remote_shell_command(&linux, script);
2422 let argv = system_ssh_argv(&linux, &prepared, Duration::ZERO);
2423 assert_eq!(argv.last().unwrap(), script);
2424 assert!(prepared.stdin_script.is_none());
2425 for command in [script, "sh -s", "", "sh -s\n"] {
2426 let prepared = remote_shell_command(&windows, command);
2427 let argv = system_ssh_argv(&windows, &prepared, Duration::ZERO);
2428 assert_eq!(argv.last().unwrap(), "sh -s");
2429 assert_eq!(prepared.stdin_script.is_none(), command == "sh -s");
2430 }
2431 }
2432
2433 #[tokio::test]
2434 async fn test_system_ssh_stdin_roundtrip_matches_null_stdin_shell() {
2435 use std::process::Stdio;
2436
2437 let _guard = test_guard!();
2438 let config = windows_worker("windows");
2439 for script in [
2442 "",
2443 "printf '%s\\n' \"apostrophe's\" 'double\"quote' '%PATH% !VAR! ^ & | < >' '$() `x` \\'",
2444 "cat\nprintf 'AFTER\\n'\nread value || printf 'EMPTY\\n'",
2445 "cat <<'END'\nheredoc ' \" % ! $() \\\nEND\nprintf 'done\\n'",
2446 "printf '%s\\n' 'first\nsecond' # trailing comment",
2447 "printf 'before\\n'; printf 'diagnostic\\n' >&2; exit 37",
2448 "# comment ending with a backslash \\",
2449 ] {
2450 let expected = tokio::process::Command::new("sh")
2451 .args(["-c", script])
2452 .stdin(Stdio::null())
2453 .output()
2454 .await
2455 .unwrap();
2456 let mut child = tokio::process::Command::new("sh")
2457 .arg("-s")
2458 .stdin(Stdio::piped())
2459 .stdout(Stdio::piped())
2460 .stderr(Stdio::piped())
2461 .kill_on_drop(true)
2462 .spawn()
2463 .unwrap();
2464 let stdin = child.stdin.take().unwrap();
2465 let payload = remote_shell_command(&config, script).stdin_script.unwrap();
2466 let (written, observed) = tokio::time::timeout(Duration::from_secs(5), async {
2467 tokio::join!(
2468 write_ssh_command_stdin(stdin, &payload),
2469 child.wait_with_output()
2470 )
2471 })
2472 .await
2473 .unwrap();
2474 written.unwrap();
2475 let observed = observed.unwrap();
2476 assert_eq!(observed.status.code(), expected.status.code(), "{script:?}");
2477 assert_eq!(observed.stdout, expected.stdout, "{script:?}");
2478 assert_eq!(observed.stderr, expected.stderr, "{script:?}");
2479 }
2480 }
2481
2482 #[tokio::test]
2483 async fn test_system_ssh_split_stdin_never_becomes_command_input() {
2484 use std::process::Stdio;
2485 use tokio::io::AsyncWriteExt;
2486
2487 let _guard = test_guard!();
2488 let script = "cat\nprintf 'AFTER\\n'\nread value || printf 'EMPTY\\n'";
2489 let payload = remote_shell_command(&windows_worker("windows"), script)
2490 .stdin_script
2491 .unwrap();
2492 let split = payload.find("cat\n").unwrap() + "cat\n".len();
2493 let mut child = tokio::process::Command::new("sh")
2494 .arg("-s")
2495 .stdin(Stdio::piped())
2496 .stdout(Stdio::piped())
2497 .stderr(Stdio::piped())
2498 .kill_on_drop(true)
2499 .spawn()
2500 .unwrap();
2501 let mut stdin = child.stdin.take().unwrap();
2502 let writer = async {
2503 stdin.write_all(&payload.as_bytes()[..split]).await.unwrap();
2504 tokio::time::sleep(Duration::from_millis(100)).await;
2505 stdin.write_all(&payload.as_bytes()[split..]).await.unwrap();
2506 stdin.shutdown().await.unwrap();
2507 };
2508 let ((), output) = tokio::time::timeout(Duration::from_secs(5), async {
2509 tokio::join!(writer, child.wait_with_output())
2510 })
2511 .await
2512 .unwrap();
2513 let output = output.unwrap();
2514 assert!(output.status.success(), "{:?}", output.stderr);
2515 assert_eq!(output.stdout, b"AFTER\nEMPTY\n");
2516 assert!(output.stderr.is_empty());
2517 }
2518
2519 #[tokio::test]
2520 async fn test_system_ssh_closed_stdin_preserves_early_exit_diagnostic() {
2521 use std::process::Stdio;
2522
2523 let _guard = test_guard!();
2524 let mut child = tokio::process::Command::new("sh")
2525 .args(["-c", "printf 'remote failure\\n' >&2; exit 37"])
2526 .stdin(Stdio::piped())
2527 .stdout(Stdio::piped())
2528 .stderr(Stdio::piped())
2529 .kill_on_drop(true)
2530 .spawn()
2531 .unwrap();
2532 let stdin = child.stdin.take().unwrap();
2533 child.wait().await.unwrap();
2536 write_ssh_command_stdin(stdin, &"x".repeat(128 * 1024))
2537 .await
2538 .unwrap();
2539 let output = child.wait_with_output().await.unwrap();
2540 assert_eq!(output.status.code(), Some(37));
2541 assert_eq!(output.stderr, b"remote failure\n");
2542 assert!(output.stdout.is_empty());
2543 }
2544
2545 #[test]
2546 fn test_system_ssh_argv_zero_server_alive_omits_keepalive() {
2547 let _guard = test_guard!();
2551 let cfg = windows_worker("wsurf");
2552 let command = remote_shell_command(&cfg, "echo hi");
2553 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(0));
2554 let s: Vec<String> = argv
2555 .iter()
2556 .map(|o| o.to_string_lossy().into_owned())
2557 .collect();
2558 assert!(
2559 !s.iter().any(|a| a.starts_with("ServerAliveInterval=")),
2560 "no ServerAliveInterval flag when interval is 0; got {:?}",
2561 s
2562 );
2563 }
2564
2565 #[test]
2566 fn test_system_ssh_argv_long_keepalive_preserved() {
2567 let _guard = test_guard!();
2571 let cfg = windows_worker("wsurf");
2572 let command = remote_shell_command(&cfg, "true");
2573 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(60));
2574 let s: Vec<String> = argv
2575 .iter()
2576 .map(|o| o.to_string_lossy().into_owned())
2577 .collect();
2578 assert!(s.contains(&"ServerAliveInterval=60".to_string()));
2579 }
2580
2581 #[test]
2582 fn test_system_ssh_argv_uses_tilde_expanded_identity() {
2583 let _guard = test_guard!();
2588 let mut cfg = windows_worker("wsurf");
2589 cfg.identity_file = "~/.ssh/operator_key".to_string();
2590 let command = remote_shell_command(&cfg, "true");
2591 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(2));
2592 let s: Vec<String> = argv
2593 .iter()
2594 .map(|o| o.to_string_lossy().into_owned())
2595 .collect();
2596 assert!(
2597 !s.contains(&"~/.ssh/operator_key".to_string()),
2598 "identity file must be tilde-expanded, got {:?}",
2599 s
2600 );
2601 assert!(
2602 s.contains(&"~/.ssh/operator_key".to_string())
2603 || s.iter().any(|a| a.ends_with(".ssh/operator_key")),
2604 "tilde-expanded path appears in argv: {:?}",
2605 s
2606 );
2607 }
2608
2609 #[test]
2610 fn test_identities_only_only_when_identity_file_exists() {
2611 let _guard = test_guard!();
2614 let dir = tempfile::tempdir().expect("tempdir");
2615 let key = dir.path().join("worker_key");
2616 std::fs::write(&key, b"not a real key").expect("write key");
2617 let key = key.to_string_lossy().into_owned();
2618 assert_eq!(
2619 identities_only_args(&key),
2620 Some(["-o", "IdentitiesOnly=yes"])
2621 );
2622 assert_eq!(identities_only_args("/nonexistent/rch/worker_key"), None);
2623
2624 let mut cfg = windows_worker("wsurf");
2625 cfg.identity_file = key;
2626 let command = remote_shell_command(&cfg, "true");
2627 let argv = system_ssh_argv(&cfg, &command, Duration::ZERO);
2628 assert!(argv.contains(&OsString::from("IdentitiesOnly=yes")));
2629
2630 cfg.identity_file = "/nonexistent/rch/worker_key".to_string();
2631 let argv = system_ssh_argv(&cfg, &command, Duration::ZERO);
2632 assert!(!argv.contains(&OsString::from("IdentitiesOnly=yes")));
2633 }
2634
2635 #[test]
2636 fn test_dispatch_helper_matches_declared_os_key() {
2637 let _guard = test_guard!();
2642 let cfg = windows_worker("wsurf");
2643 assert_eq!(declared_os(&cfg.tags).as_deref(), Some("windows"));
2644 assert!(prefers_system_ssh(&cfg));
2645 }
2646}
2647
2648#[cfg(test)]
2649mod ssh_connection_tests {
2650 use super::*;
2651 use std::os::unix::fs::PermissionsExt;
2652
2653 struct FakeSsh {
2654 directory: tempfile::TempDir,
2655 program: PathBuf,
2656 }
2657
2658 impl FakeSsh {
2659 fn new(mode: &str) -> Self {
2660 let directory = tempfile::tempdir().unwrap();
2661 let program = directory.path().join("fake-ssh");
2662 let script = format!(
2663 r#"#!/bin/sh
2664root={root}
2665mode={mode}
2666operation=
2667socket=
2668log=
2669while [ "$#" -gt 0 ]; do
2670 case "$1" in
2671 -G) operation=config; shift ;;
2672 -F) shift 2 ;;
2673 -S) socket=$2; shift 2 ;;
2674 -E) log=$2; shift 2 ;;
2675 -O) operation=$2; shift 2 ;;
2676 *) shift ;;
2677 esac
2678done
2679if [ "$operation" = config ]; then
2680 printf '%s\n' "$$" >> "$root/capability-pids"
2681 case "$mode" in
2682 capability-stall) exec sleep 30 ;;
2683 old-client) printf 'hostname none\n' ;;
2684 *) printf 'hostname none\nforkafterauthentication no\n' ;;
2685 esac
2686 exit 0
2687fi
2688if [ "$operation" = check ]; then
2689 printf '%s\n' "$$" >> "$root/check-pids"
2690 case "$mode" in
2691 probe-stall) exec sleep 30 ;;
2692 probe-reject) exit 37 ;;
2693 *) exit 0 ;;
2694 esac
2695fi
2696printf '%s\n' "$$" >> "$root/master-pids"
2697case "$mode" in
2698 ready|old-client|probe-stall|probe-reject) : > "$socket" ;;
2699 fail) printf 'Permission denied (publickey).\n' > "$log"; exit 255 ;;
2700esac
2701exec sleep 30
2702"#,
2703 root = shell_escape_value(&directory.path().to_string_lossy())
2704 .expect("fixture path must be shell-escapable"),
2705 mode = shell_escape_value(mode).expect("fixture mode must be shell-escapable"),
2706 );
2707 std::fs::write(&program, script).unwrap();
2708 std::fs::set_permissions(&program, std::fs::Permissions::from_mode(0o700)).unwrap();
2709 Self { directory, program }
2710 }
2711
2712 fn client(&self, timeout: Duration, reuse: bool) -> SshClient {
2713 SshClient::new(
2714 WorkerConfig {
2715 host: "auth-stalled.invalid".to_owned(),
2716 identity_file: "/nonexistent/rch/test-key".to_owned(),
2717 ..WorkerConfig::default()
2718 },
2719 SshOptions {
2720 connect_timeout: timeout,
2721 control_master: reuse,
2722 ..SshOptions::default()
2723 },
2724 )
2725 }
2726
2727 async fn connect(&self, client: &mut SshClient) -> Result<()> {
2728 client
2729 .connect_using(&self.program, Some(self.directory.path()))
2730 .await
2731 }
2732
2733 fn pids(&self, filename: &str) -> Vec<u32> {
2734 std::fs::read_to_string(self.directory.path().join(filename))
2735 .unwrap_or_default()
2736 .lines()
2737 .map(|line| line.parse().unwrap())
2738 .collect()
2739 }
2740
2741 async fn wait_for_pid(&self, filename: &str) {
2742 tokio::time::timeout(Duration::from_secs(3), async {
2743 while self.pids(filename).is_empty() {
2744 tokio::time::sleep(Duration::from_millis(10)).await;
2745 }
2746 })
2747 .await
2748 .expect("fake SSH process must start");
2749 }
2750
2751 async fn assert_reaped(&self) {
2752 let pids: Vec<_> = self
2753 .pids("master-pids")
2754 .into_iter()
2755 .chain(self.pids("check-pids"))
2756 .chain(self.pids("capability-pids"))
2757 .collect();
2758 assert!(!pids.is_empty(), "test must launch a real SSH stand-in");
2759 tokio::time::timeout(Duration::from_secs(3), async {
2760 loop {
2761 let any_alive = pids.iter().any(|pid| {
2762 std::process::Command::new("sh")
2766 .args(["-c", &format!("kill -0 {pid} 2>/dev/null")])
2767 .stdin(std::process::Stdio::null())
2768 .stdout(std::process::Stdio::null())
2769 .stderr(std::process::Stdio::null())
2770 .status()
2771 .unwrap()
2772 .success()
2773 });
2774 let control_dir_exists =
2775 std::fs::read_dir(self.directory.path())
2776 .unwrap()
2777 .any(|entry| {
2778 entry
2779 .unwrap()
2780 .file_name()
2781 .to_string_lossy()
2782 .starts_with("rch-")
2783 });
2784 if !any_alive && !control_dir_exists {
2785 break;
2786 }
2787 tokio::time::sleep(Duration::from_millis(10)).await;
2788 }
2789 })
2790 .await
2791 .expect("owned SSH children must be killed, reaped, and release their sockets");
2792 }
2793 }
2794
2795 #[test]
2796 fn foreground_master_keeps_host_identity_and_keepalive_policy() {
2797 let fake = FakeSsh::new("ready");
2798 let key = fake.directory.path().join("identity");
2799 std::fs::write(&key, "test identity").unwrap();
2800 for (policy, expected) in [
2801 (KnownHostsPolicy::Strict, "StrictHostKeyChecking=yes"),
2802 (KnownHostsPolicy::Add, "StrictHostKeyChecking=accept-new"),
2803 (KnownHostsPolicy::AcceptAll, "StrictHostKeyChecking=no"),
2804 ] {
2805 let mut client = fake.client(Duration::from_secs(7), true);
2806 client.config.identity_file = key.to_string_lossy().into_owned();
2807 client.options.known_hosts = policy;
2808 client.options.server_alive_interval = Some(Duration::from_secs(13));
2809 client.options.control_persist_idle = Some(Duration::from_secs(60));
2810 let command = client.master_command(
2811 &fake.program,
2812 "builder@worker",
2813 Path::new("/private/control/master"),
2814 Path::new("/private/control/log"),
2815 true,
2816 );
2817 let args: Vec<_> = command
2818 .as_std()
2819 .get_args()
2820 .map(|argument| argument.to_string_lossy().into_owned())
2821 .collect();
2822 for expected in [
2823 expected,
2824 "BatchMode=yes",
2825 "ConnectTimeout=7",
2826 "ServerAliveInterval=13",
2827 "IdentitiesOnly=yes",
2828 "ControlPersist=no",
2829 "ForkAfterAuthentication=no",
2830 "-M",
2831 "-N",
2832 "builder@worker",
2833 ] {
2834 assert!(args.iter().any(|argument| argument == expected), "{args:?}");
2835 }
2836 assert!(
2837 args.iter()
2838 .any(|argument| argument == key.to_str().unwrap())
2839 );
2840 assert!(!args.iter().any(|argument| argument == "-f"));
2841 assert!(
2842 !args
2843 .iter()
2844 .any(|argument| argument.starts_with("IgnoreUnknown="))
2845 );
2846 }
2847
2848 let client = fake.client(Duration::from_secs(1), false);
2849 let command = client.master_command(
2850 &fake.program,
2851 "builder@worker",
2852 Path::new("/private/master"),
2853 Path::new("/private/log"),
2854 false,
2855 );
2856 assert!(
2857 command
2858 .as_std()
2859 .get_args()
2860 .all(|arg| arg != "IdentitiesOnly=yes")
2861 );
2862 assert!(
2863 command
2864 .as_std()
2865 .get_args()
2866 .all(|arg| arg != "ForkAfterAuthentication=no")
2867 );
2868 }
2869
2870 #[tokio::test]
2871 async fn capability_probe_deadline_and_cancellation_reap_before_authentication() {
2872 for cancelled in [false, true] {
2873 let fake = FakeSsh::new("capability-stall");
2874 let timeout = if cancelled {
2875 Duration::from_secs(30)
2876 } else {
2877 Duration::from_millis(200)
2878 };
2879 let mut client = fake.client(timeout, true);
2880 let program = fake.program.clone();
2881 let directory = fake.directory.path().to_owned();
2882 let connecting =
2883 tokio::spawn(async move { client.connect_using(&program, Some(&directory)).await });
2884 if cancelled {
2885 fake.wait_for_pid("capability-pids").await;
2886 connecting.abort();
2887 assert!(connecting.await.unwrap_err().is_cancelled());
2888 } else {
2889 let error = connecting.await.unwrap().unwrap_err();
2890 assert!(format!("{error:#}").contains("timed out"), "{error:#}");
2891 }
2892 assert!(fake.pids("master-pids").is_empty());
2893 assert_eq!(fake.pids("capability-pids").len(), 1);
2894 fake.assert_reaped().await;
2895 }
2896 }
2897
2898 #[tokio::test]
2899 async fn client_without_fork_option_remains_supported() {
2900 let fake = FakeSsh::new("old-client");
2901 let mut client = fake.client(Duration::from_secs(3), true);
2902 fake.connect(&mut client).await.unwrap();
2903 assert!(client.is_connected());
2904 client.disconnect().await.unwrap();
2905 fake.assert_reaped().await;
2906 }
2907
2908 #[tokio::test]
2909 async fn foreground_options_preserve_user_ignore_unknown_configuration() {
2910 let directory = Arc::new(tempfile::tempdir().unwrap());
2911 let config = directory.path().join("user-config");
2912 std::fs::write(
2913 &config,
2914 "IgnoreUnknown RchFixtureExtension\nRchFixtureExtension yes\n",
2915 )
2916 .unwrap();
2917 let deadline = Instant::now() + Duration::from_secs(3);
2918 let capability = ssh_can_disable_background(Path::new("ssh"), directory.clone(), deadline)
2919 .await
2920 .unwrap();
2921 let client = SshClient::new(WorkerConfig::default(), SshOptions::default());
2922 let master = client.master_command(
2923 Path::new("ssh"),
2924 "none",
2925 &directory.path().join("master"),
2926 &directory.path().join("log"),
2927 capability,
2928 );
2929 let mut command = tokio::process::Command::new("ssh");
2932 command
2933 .args(["-G", "-F"])
2934 .arg(config)
2935 .args(master.as_std().get_args())
2936 .stdin(std::process::Stdio::null())
2937 .stdout(std::process::Stdio::null())
2938 .stderr(std::process::Stdio::null());
2939 let process = OwnedSshProcess::spawn(&mut command, Some(directory), None).unwrap();
2940 assert!(process.wait_until(deadline).await.unwrap().success());
2941 }
2942
2943 #[tokio::test]
2944 async fn authentication_deadline_kills_and_reaps_without_starting_fallback() {
2945 let fake = FakeSsh::new("auth-stall");
2946 let mut client = fake.client(Duration::from_millis(200), true);
2947 let start = Instant::now();
2948 let error = fake.connect(&mut client).await.unwrap_err();
2949 assert!(format!("{error:#}").contains("timed out"), "{error:#}");
2950 assert!(start.elapsed() < Duration::from_secs(2));
2951 assert!(!client.is_connected());
2952 assert_eq!(fake.pids("master-pids").len(), 1);
2953 fake.assert_reaped().await;
2954 }
2955
2956 #[tokio::test]
2957 async fn cancelling_authentication_kills_and_reaps_child() {
2958 let fake = FakeSsh::new("auth-stall");
2959 let mut client = fake.client(Duration::from_secs(30), true);
2960 let program = fake.program.clone();
2961 let directory = fake.directory.path().to_owned();
2962 let connecting =
2963 tokio::spawn(async move { client.connect_using(&program, Some(&directory)).await });
2964 fake.wait_for_pid("master-pids").await;
2965 connecting.abort();
2966 assert!(connecting.await.unwrap_err().is_cancelled());
2967 fake.assert_reaped().await;
2968 }
2969
2970 #[tokio::test]
2971 async fn stalled_readiness_probe_shares_deadline_and_is_reaped() {
2972 let fake = FakeSsh::new("probe-stall");
2973 let mut client = fake.client(Duration::from_millis(200), true);
2974 let error = fake.connect(&mut client).await.unwrap_err();
2975 assert!(format!("{error:#}").contains("timed out"), "{error:#}");
2976 assert_eq!(fake.pids("master-pids").len(), 1);
2977 assert_eq!(fake.pids("check-pids").len(), 1);
2978 fake.assert_reaped().await;
2979 }
2980
2981 #[tokio::test]
2982 async fn unsuccessful_readiness_probe_never_marks_client_connected() {
2983 let fake = FakeSsh::new("probe-reject");
2984 let mut client = fake.client(Duration::from_millis(200), false);
2985 assert!(fake.connect(&mut client).await.is_err());
2986 assert!(!client.is_connected());
2987 assert!(!fake.pids("check-pids").is_empty());
2988 fake.assert_reaped().await;
2989 }
2990
2991 #[tokio::test]
2992 async fn dropping_and_disconnecting_ready_client_reap_master() {
2993 let fake = FakeSsh::new("ready");
2994 let mut client = fake.client(Duration::from_secs(3), true);
2995 fake.connect(&mut client).await.unwrap();
2996 assert!(client.is_connected());
2997 client.disconnect().await.unwrap();
2998 assert!(!client.is_connected());
2999 fake.assert_reaped().await;
3000
3001 fake.connect(&mut client).await.unwrap();
3002 assert!(client.is_connected());
3003 drop(client);
3004 fake.assert_reaped().await;
3005 assert_eq!(fake.pids("master-pids").len(), 2);
3006 }
3007
3008 #[tokio::test]
3009 async fn failed_primary_retries_and_preserves_authentication_diagnostic() {
3010 let fake = FakeSsh::new("fail");
3011 let mut client = fake.client(Duration::from_secs(3), true);
3012 let error = fake.connect(&mut client).await.unwrap_err();
3013 assert!(
3014 format!("{error:#}").contains("Permission denied (publickey)."),
3015 "{error:#}"
3016 );
3017 assert_eq!(fake.pids("master-pids").len(), 2);
3018 fake.assert_reaped().await;
3019 }
3020
3021 #[tokio::test]
3022 async fn active_commands_prevent_idle_expiry_and_release_restarts_timer() {
3023 let mut command = tokio::process::Command::new("sleep");
3024 command.arg("30");
3025 let process =
3026 OwnedSshProcess::spawn(&mut command, None, Some(Duration::from_millis(100))).unwrap();
3027 let first = process.begin_use().unwrap();
3028 let second = process.begin_use().unwrap();
3029 tokio::time::sleep(Duration::from_millis(150)).await;
3030 first.complete();
3031 tokio::time::sleep(Duration::from_millis(150)).await;
3032 assert!(
3033 process.is_running(),
3034 "the remaining command still owns the master"
3035 );
3036 second.complete();
3037 assert!(
3038 process.is_running(),
3039 "expiry starts after the last command completes"
3040 );
3041 let status = process
3042 .wait_until(Instant::now() + Duration::from_secs(3))
3043 .await
3044 .unwrap();
3045 assert!(!status.success());
3046 assert!(!process.is_running());
3047 assert!(process.begin_use().is_err());
3048 }
3049
3050 #[tokio::test]
3051 async fn cancelled_command_retires_master_after_other_commands_finish() {
3052 let mut command = tokio::process::Command::new("sleep");
3053 command.arg("30");
3054 let process = OwnedSshProcess::spawn(&mut command, None, None).unwrap();
3055 let cancelled = process.begin_use().unwrap();
3056 let running = process.begin_use().unwrap();
3057 drop(cancelled);
3058 assert!(
3059 process.begin_use().is_err(),
3060 "a draining master rejects new work"
3061 );
3062 tokio::time::sleep(Duration::from_millis(100)).await;
3063 assert!(
3064 process.state.borrow().outcome.is_none(),
3065 "cancellation must not kill another active command's master"
3066 );
3067 running.complete();
3068 let status = process
3069 .wait_until(Instant::now() + Duration::from_secs(3))
3070 .await
3071 .unwrap();
3072 assert!(!status.success());
3073 }
3074
3075 #[tokio::test]
3076 async fn expired_pooled_connection_can_connect_again() {
3077 let fake = FakeSsh::new("ready");
3078 let mut client = fake.client(Duration::from_secs(3), true);
3079 client.options.control_persist_idle = Some(Duration::from_secs(1));
3080 fake.connect(&mut client).await.unwrap();
3081 client
3082 .session
3083 .as_ref()
3084 .unwrap()
3085 .process
3086 .wait_until(Instant::now() + Duration::from_secs(3))
3087 .await
3088 .unwrap();
3089 assert!(!client.is_connected());
3090 fake.connect(&mut client).await.unwrap();
3091 assert!(client.is_connected());
3092 client.disconnect().await.unwrap();
3093 fake.assert_reaped().await;
3094 assert_eq!(fake.pids("master-pids").len(), 2);
3095 }
3096}
3097
3098#[derive(Debug, Clone, PartialEq, Eq)]
3100enum ControlPersistMode {
3101 Closed,
3103 IdleFor(NonZeroUsize),
3105 TooLarge(u64),
3107}
3108
3109fn control_persist_mode(control_master: bool, idle: Option<Duration>) -> ControlPersistMode {
3114 match idle {
3115 Some(idle) if control_master && !idle.is_zero() => match usize::try_from(idle.as_secs()) {
3116 Ok(secs) => match NonZeroUsize::new(secs) {
3117 Some(nonzero) => ControlPersistMode::IdleFor(nonzero),
3118 None => ControlPersistMode::Closed,
3119 },
3120 Err(_) => ControlPersistMode::TooLarge(idle.as_secs()),
3121 },
3122 _ => ControlPersistMode::Closed,
3123 }
3124}
3125
3126#[cfg(test)]
3127mod control_persist_tests {
3128 use super::*;
3129 use std::time::Duration;
3130
3131 #[test]
3132 fn non_mux_sessions_never_persist_forever() {
3133 assert_eq!(
3135 control_persist_mode(false, None),
3136 ControlPersistMode::Closed
3137 );
3138 assert_eq!(
3139 control_persist_mode(false, Some(Duration::from_secs(60))),
3140 ControlPersistMode::Closed
3141 );
3142 }
3143
3144 #[test]
3145 fn mux_without_idle_closes() {
3146 assert_eq!(control_persist_mode(true, None), ControlPersistMode::Closed);
3147 assert_eq!(
3148 control_persist_mode(true, Some(Duration::from_secs(0))),
3149 ControlPersistMode::Closed
3150 );
3151 }
3152
3153 #[test]
3154 fn mux_with_idle_keeps_warm() {
3155 assert_eq!(
3156 control_persist_mode(true, Some(Duration::from_secs(60))),
3157 ControlPersistMode::IdleFor(NonZeroUsize::new(60).unwrap())
3158 );
3159 }
3160}