1use crate::types::{WorkerConfig, WorkerId, declared_os};
9use anyhow::{Context, Result};
10use openssh::{ControlPersist, KnownHosts, Session, SessionBuilder, Stdio};
11use std::collections::HashMap;
12use std::ffi::OsString;
13use std::num::NonZeroUsize;
14use std::path::Path;
15use std::sync::Arc;
16use std::time::Duration;
17use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
18use tokio::sync::{RwLock, mpsc};
19use tracing::{debug, error, warn};
20
21pub use crate::ssh_utils::{
23 CommandResult, EnvPrefix, build_env_prefix, is_retryable_transport_error,
24 is_retryable_transport_error_text, is_valid_env_key, shell_escape_value,
25};
26
27const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
29
30const DEFAULT_COMMAND_TIMEOUT: Duration = Duration::from_secs(300);
32
33const MAX_OUTPUT_SIZE: u64 = 10 * 1024 * 1024;
35
36const HEALTH_CHECK_COMMAND: &str = "echo ok";
37
38const POOL_LIVENESS_PROBE_TIMEOUT: Duration = Duration::from_secs(10);
42
43fn is_expected_health_check_output(stdout: &str) -> bool {
44 stdout
45 .trim()
46 .lines()
47 .last()
48 .is_some_and(is_health_check_sentinel)
49}
50
51fn is_health_check_sentinel(line: &str) -> bool {
52 matches!(line.trim(), "ok")
53}
54
55#[derive(Debug, Clone)]
57pub struct SshOptions {
58 pub connect_timeout: Duration,
60 pub command_timeout: Duration,
62 pub server_alive_interval: Option<Duration>,
66 pub control_persist_idle: Option<Duration>,
75 pub control_master: bool,
77 pub known_hosts: KnownHostsPolicy,
79}
80
81impl Default for SshOptions {
82 fn default() -> Self {
83 Self {
84 connect_timeout: DEFAULT_CONNECT_TIMEOUT,
85 command_timeout: DEFAULT_COMMAND_TIMEOUT,
86 server_alive_interval: None,
87 control_persist_idle: None,
88 control_master: false,
92 known_hosts: KnownHostsPolicy::Add,
93 }
94 }
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99pub enum KnownHostsPolicy {
100 Strict,
102 Add,
104 AcceptAll,
106}
107
108#[cfg(test)]
109mod retry_tests {
110 use super::*;
111 use crate::test_guard;
112
113 #[test]
114 fn test_retryable_transport_error_text() {
115 let _guard = test_guard!();
116 assert!(is_retryable_transport_error_text(
117 "ssh: connect to host 1.2.3.4 port 22: Connection timed out"
118 ));
119 assert!(is_retryable_transport_error_text(
120 "kex_exchange_identification: Connection reset by peer"
121 ));
122 assert!(is_retryable_transport_error_text("Broken pipe"));
123 assert!(is_retryable_transport_error_text("Network is unreachable"));
124 }
125
126 #[test]
127 fn test_non_retryable_transport_error_text() {
128 let _guard = test_guard!();
129 assert!(!is_retryable_transport_error_text(
130 "Permission denied (publickey)."
131 ));
132 assert!(!is_retryable_transport_error_text(
133 "Host key verification failed."
134 ));
135 assert!(!is_retryable_transport_error_text(
136 "Could not resolve hostname worker.example.com: Name or service not known"
137 ));
138 assert!(!is_retryable_transport_error_text(
139 "Identity file /nope/id_rsa not accessible: No such file or directory"
140 ));
141 }
142}
143
144pub struct SshClient {
146 config: WorkerConfig,
148 options: SshOptions,
150 session: Option<Session>,
152}
153
154impl SshClient {
155 pub fn new(config: WorkerConfig, options: SshOptions) -> Self {
157 Self {
158 config,
159 options,
160 session: None,
161 }
162 }
163
164 pub fn worker_id(&self) -> &WorkerId {
166 &self.config.id
167 }
168
169 pub fn is_connected(&self) -> bool {
171 self.session.is_some()
172 }
173
174 fn is_configured_for(&self, config: &WorkerConfig) -> bool {
175 self.config.id == config.id
176 && self.config.host == config.host
177 && self.config.user == config.user
178 && self.config.identity_file == config.identity_file
179 && declared_os(&self.config.tags) == declared_os(&config.tags)
180 }
181
182 pub async fn connect(&mut self) -> Result<()> {
184 if self.session.is_some() {
185 debug!("Already connected to {}", self.config.id);
186 return Ok(());
187 }
188
189 let destination = format!("{}@{}", self.config.user, self.config.host);
190 debug!("Connecting to {} via SSH...", destination);
191
192 let session = match self
193 .connect_with_mode(&destination, self.options.control_master)
194 .await
195 {
196 Ok(session) => session,
197 Err(primary_error) if self.options.control_master => {
198 warn!(
199 "SSH ControlMaster connection to {} failed ({}). Retrying without ControlMaster.",
200 destination, primary_error
201 );
202 self.connect_with_mode(&destination, false)
203 .await
204 .with_context(|| {
205 format!(
206 "Failed to connect to {} after retrying without ControlMaster",
207 destination
208 )
209 })?
210 }
211 Err(primary_error) => {
212 return Err(primary_error)
213 .with_context(|| format!("Failed to connect to {}", destination));
214 }
215 };
216
217 debug!("Connected to {} ({})", self.config.id, self.config.host);
220 self.session = Some(session);
221 Ok(())
222 }
223
224 async fn connect_with_mode(&self, destination: &str, control_master: bool) -> Result<Session> {
225 let mut builder = SessionBuilder::default();
226 self.configure_builder(&mut builder, control_master);
227
228 builder.connect(destination).await.with_context(|| {
229 if control_master {
230 format!(
231 "Failed to connect to {} with ControlMaster enabled",
232 destination
233 )
234 } else {
235 format!(
236 "Failed to connect to {} with ControlMaster disabled",
237 destination
238 )
239 }
240 })
241 }
242
243 fn configure_builder(&self, builder: &mut SessionBuilder, control_master: bool) {
244 let known_hosts = match self.options.known_hosts {
245 KnownHostsPolicy::Strict => KnownHosts::Strict,
246 KnownHostsPolicy::Add => KnownHosts::Add,
247 KnownHostsPolicy::AcceptAll => KnownHosts::Accept,
248 };
249
250 builder
251 .known_hosts_check(known_hosts)
252 .connect_timeout(self.options.connect_timeout);
253
254 if let Some(interval) = self.options.server_alive_interval {
255 builder.server_alive_interval(interval);
256 }
257
258 let identity_path = shellexpand::tilde(&self.config.identity_file);
260 if Path::new(identity_path.as_ref()).exists() {
261 builder.keyfile(identity_path.as_ref());
262 }
263
264 match control_persist_mode(control_master, self.options.control_persist_idle) {
272 ControlPersistMode::IdleFor(nonzero) => {
273 builder.control_persist(ControlPersist::IdleFor(nonzero));
274 }
275 ControlPersistMode::TooLarge(secs) => {
276 warn!("control_persist_idle too large ({secs}s); closing after initial connection");
277 builder.control_persist(ControlPersist::ClosedAfterInitialConnection);
278 }
279 ControlPersistMode::Closed => {
280 builder.control_persist(ControlPersist::ClosedAfterInitialConnection);
281 }
282 }
283
284 if control_master {
286 let control_dir = {
295 let home_ssh = dirs::home_dir().map(|h| h.join(".ssh").join("rch"));
296
297 if let Some(ref dir) = home_ssh {
298 dir.clone()
299 } else if let Some(runtime_dir) = dirs::runtime_dir() {
300 runtime_dir.join("rch-ssh")
301 } else {
302 let username = whoami::username().unwrap_or_else(|_| "unknown".to_string());
303 std::env::temp_dir().join(format!("rch-ssh-{}", username))
304 }
305 };
306
307 if let Err(e) = std::fs::create_dir_all(&control_dir) {
308 warn!(
309 "Failed to create SSH control directory {:?}: {}",
310 control_dir, e
311 );
312 } else {
313 #[cfg(unix)]
316 {
317 use std::os::unix::fs::PermissionsExt;
318 if let Err(e) = std::fs::set_permissions(
319 &control_dir,
320 std::fs::Permissions::from_mode(0o700),
321 ) {
322 warn!(
323 "Failed to set permissions on SSH control directory {:?}: {}",
324 control_dir, e
325 );
326 }
327 }
328 }
329 builder.control_directory(&control_dir);
330 }
331 }
332
333 pub async fn disconnect(&mut self) -> Result<()> {
335 if let Some(session) = self.session.take() {
336 debug!("Disconnecting from {}", self.config.id);
337 session.close().await?;
338 debug!("Disconnected from {}", self.config.id);
340 }
341 Ok(())
342 }
343
344 pub async fn reconnect(&mut self) -> Result<()> {
353 if let Err(e) = self.disconnect().await {
354 debug!(
355 "Ignoring error closing stale session to {} before reconnect: {}",
356 self.config.id, e
357 );
358 }
359 self.connect().await
360 }
361
362 pub async fn execute(&self, command: &str) -> Result<CommandResult> {
365 self.execute_with_timeout(command, self.options.command_timeout)
366 .await
367 }
368
369 pub async fn execute_with_timeout(
375 &self,
376 command: &str,
377 command_timeout: Duration,
378 ) -> Result<CommandResult> {
379 if prefers_system_ssh(&self.config) {
387 return system_ssh_execute(
388 &self.config,
389 command,
390 command_timeout,
391 self.options
392 .server_alive_interval
393 .unwrap_or(Duration::from_secs(2)),
394 )
395 .await;
396 }
397
398 let session = self.session.as_ref().context("Not connected to worker")?;
399
400 let start = std::time::Instant::now();
401 debug!(
402 "Executing on {}: {}",
403 self.config.id,
404 crate::util::mask_sensitive_command(command)
405 );
406
407 let mut child = session
408 .command("sh")
409 .arg("-c")
410 .arg(command)
411 .stdout(Stdio::piped())
412 .stderr(Stdio::piped())
413 .spawn()
414 .await
415 .with_context(|| format!("Failed to spawn command on {}", self.config.id))?;
416
417 let execution_future = async {
418 let stdout_handle = child.stdout().take();
420 let stderr_handle = child.stderr().take();
421
422 let stdout_fut = async {
423 if let Some(out) = stdout_handle {
424 let reader = BufReader::new(out);
425 let mut take = reader.take(MAX_OUTPUT_SIZE);
426 let mut buf = String::new();
427 take.read_to_string(&mut buf).await?;
428 let mut reader = take.into_inner();
430 let mut sink = tokio::io::sink();
431 tokio::io::copy(&mut reader, &mut sink).await?;
432 if buf.len() >= MAX_OUTPUT_SIZE as usize {
433 buf.push_str("\n...[output truncated]...\n");
434 }
435 Ok::<String, anyhow::Error>(buf)
436 } else {
437 Ok(String::new())
438 }
439 };
440
441 let stderr_fut = async {
442 if let Some(err) = stderr_handle {
443 let reader = BufReader::new(err);
444 let mut take = reader.take(MAX_OUTPUT_SIZE);
445 let mut buf = String::new();
446 take.read_to_string(&mut buf).await?;
447 let mut reader = take.into_inner();
449 let mut sink = tokio::io::sink();
450 tokio::io::copy(&mut reader, &mut sink).await?;
451 if buf.len() >= MAX_OUTPUT_SIZE as usize {
452 buf.push_str("\n...[output truncated]...\n");
453 }
454 Ok::<String, anyhow::Error>(buf)
455 } else {
456 Ok(String::new())
457 }
458 };
459
460 let (stdout, stderr) = tokio::try_join!(stdout_fut, stderr_fut)?;
461
462 let status = child
463 .wait()
464 .await
465 .with_context(|| "Failed to wait for command completion")?;
466
467 Ok::<_, anyhow::Error>((status, stdout, stderr))
468 };
469
470 match tokio::time::timeout(command_timeout, execution_future).await {
471 Ok(result) => {
472 let (status, stdout, stderr) = result?;
473 let duration = start.elapsed();
474 let exit_code = status.code().unwrap_or(-1);
475
476 debug!(
477 "Command completed on {} (exit={}, duration={}ms)",
478 self.config.id,
479 exit_code,
480 duration.as_millis()
481 );
482
483 Ok(CommandResult {
484 exit_code,
485 stdout,
486 stderr,
487 duration_ms: duration.as_millis() as u64,
488 })
489 }
490 Err(_) => {
491 warn!(
498 "Command timed out on {} after {:?}",
499 self.config.id, command_timeout
500 );
501 anyhow::bail!("Command timed out after {:?}", command_timeout);
502 }
503 }
504 }
505
506 pub async fn execute_streaming<F, G>(
508 &self,
509 command: &str,
510 mut on_stdout: F,
511 mut on_stderr: G,
512 ) -> Result<CommandResult>
513 where
514 F: FnMut(&str),
515 G: FnMut(&str),
516 {
517 let session = self.session.as_ref().context("Not connected to worker")?;
518
519 let start = std::time::Instant::now();
520 debug!(
521 "Executing (streaming) on {}: {}",
522 self.config.id,
523 crate::util::mask_sensitive_command(command)
524 );
525
526 let mut child = session
527 .command("sh")
528 .arg("-c")
529 .arg(command)
530 .stdout(Stdio::piped())
531 .stderr(Stdio::piped())
532 .spawn()
533 .await
534 .with_context(|| format!("Failed to spawn command on {}", self.config.id))?;
535
536 let stdout = child.stdout().take();
537 let stderr = child.stderr().take();
538
539 let (tx, mut rx) = mpsc::channel(100);
542
543 if let Some(out) = stdout {
545 let tx = tx.clone();
546 tokio::spawn(async move {
547 let mut reader = BufReader::new(out);
548 let mut line = String::new();
549 loop {
550 line.clear();
551 match reader.read_line(&mut line).await {
552 Ok(0) => break, Ok(_) => {
554 if tx.send(StreamEvent::Stdout(line.clone())).await.is_err() {
555 break; }
557 }
558 Err(_) => break, }
560 }
561 });
562 }
563
564 if let Some(err) = stderr {
566 let tx = tx.clone();
567 tokio::spawn(async move {
568 let mut reader = BufReader::new(err);
569 let mut line = String::new();
570 loop {
571 line.clear();
572 match reader.read_line(&mut line).await {
573 Ok(0) => break, Ok(_) => {
575 if tx.send(StreamEvent::Stderr(line.clone())).await.is_err() {
576 break; }
578 }
579 Err(_) => break, }
581 }
582 });
583 }
584
585 drop(tx);
587
588 let mut stdout_acc = String::new();
589 let mut stderr_acc = String::new();
590
591 enum StreamEvent {
592 Stdout(String),
593 Stderr(String),
594 }
595
596 let streaming_future = async {
597 while let Some(event) = rx.recv().await {
599 match event {
600 StreamEvent::Stdout(line) => {
601 on_stdout(&line);
602 if stdout_acc.len() < MAX_OUTPUT_SIZE as usize {
603 stdout_acc.push_str(&line);
604 if stdout_acc.len() >= MAX_OUTPUT_SIZE as usize {
605 stdout_acc.push_str("\n...[output truncated]...\n");
606 }
607 }
608 }
609 StreamEvent::Stderr(line) => {
610 on_stderr(&line);
611 if stderr_acc.len() < MAX_OUTPUT_SIZE as usize {
612 stderr_acc.push_str(&line);
613 if stderr_acc.len() >= MAX_OUTPUT_SIZE as usize {
614 stderr_acc.push_str("\n...[output truncated]...\n");
615 }
616 }
617 }
618 }
619 }
620
621 let status = child.wait().await?;
622 Ok::<_, anyhow::Error>(status)
623 };
624
625 match tokio::time::timeout(self.options.command_timeout, streaming_future).await {
626 Ok(result) => {
627 let status = result?;
628 let duration = start.elapsed();
629 let exit_code = status.code().unwrap_or(-1);
630
631 Ok(CommandResult {
632 exit_code,
633 stdout: stdout_acc,
634 stderr: stderr_acc,
635 duration_ms: duration.as_millis() as u64,
636 })
637 }
638 Err(_) => {
639 warn!(
649 "Command (streaming) timed out on {} after {:?}, cleaning up",
650 self.config.id, self.options.command_timeout
651 );
652 anyhow::bail!("Command timed out after {:?}", self.options.command_timeout);
655 }
656 }
657 }
658
659 pub async fn health_check(&self) -> Result<bool> {
661 self.health_check_with_timeout(self.options.command_timeout)
662 .await
663 }
664
665 pub async fn health_check_with_timeout(&self, timeout: Duration) -> Result<bool> {
668 match self
669 .execute_with_timeout(HEALTH_CHECK_COMMAND, timeout)
670 .await
671 {
672 Ok(result) => Ok(result.success() && is_expected_health_check_output(&result.stdout)),
673 Err(e) => {
674 warn!("Health check failed for {}: {}", self.config.id, e);
675 Ok(false)
676 }
677 }
678 }
679}
680
681fn pooled_client_options(pool_options: &SshOptions, config: &WorkerConfig) -> SshOptions {
703 let mut options = pool_options.clone();
704 if declared_os(&config.tags).as_deref() == Some("windows") {
705 options.control_master = false;
706 }
707 options
708}
709
710pub(crate) fn prefers_system_ssh(config: &WorkerConfig) -> bool {
732 declared_os(&config.tags).as_deref() == Some("windows")
733}
734
735#[derive(Debug, Clone, PartialEq, Eq)]
737pub struct RemoteShellCommand {
738 pub command: String,
740 pub stdin_script: Option<String>,
743}
744
745#[must_use]
750pub fn remote_shell_command(config: &WorkerConfig, command: &str) -> RemoteShellCommand {
751 if prefers_system_ssh(config) {
752 RemoteShellCommand {
753 command: "sh -s".into(),
754 stdin_script: (command != "sh -s").then(|| {
755 format!(
756 "exec sh -c {} </dev/null\n",
757 shell_escape::escape(command.into())
758 )
759 }),
760 }
761 } else {
762 RemoteShellCommand {
763 command: command.into(),
764 stdin_script: None,
765 }
766 }
767}
768
769pub(crate) fn system_ssh_argv(
793 config: &WorkerConfig,
794 command: &RemoteShellCommand,
795 server_alive_interval: Duration,
796) -> Vec<OsString> {
797 let identity_path = shellexpand::tilde(&config.identity_file);
798 let destination = format!("{}@{}", config.user, config.host);
799
800 let mut argv: Vec<OsString> = Vec::with_capacity(10);
801 argv.push(OsString::from("ssh"));
802 argv.push(OsString::from("-i"));
803 argv.push(OsString::from(identity_path.as_ref()));
804 argv.push(OsString::from("-o"));
805 argv.push(OsString::from("BatchMode=yes"));
806 argv.push(OsString::from("-o"));
807 argv.push(OsString::from("ConnectTimeout=8"));
808 argv.push(OsString::from("-o"));
809 argv.push(OsString::from("StrictHostKeyChecking=accept-new"));
810 if !server_alive_interval.is_zero() {
811 let secs = server_alive_interval.as_secs().max(1);
815 argv.push(OsString::from("-o"));
816 argv.push(OsString::from(format!("ServerAliveInterval={secs}")));
817 }
818 argv.push(OsString::from(destination));
819 argv.push(OsString::from(&command.command));
820 argv
821}
822
823pub async fn write_ssh_command_stdin(
830 mut stdin: tokio::process::ChildStdin,
831 input: &str,
832) -> std::io::Result<()> {
833 use tokio::io::AsyncWriteExt;
834
835 let result = async {
836 stdin.write_all(input.as_bytes()).await?;
837 stdin.shutdown().await
838 }
839 .await;
840 match result {
841 Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),
844 result => result,
845 }
846}
847
848pub(crate) async fn system_ssh_execute(
863 config: &WorkerConfig,
864 command: &str,
865 command_timeout: Duration,
866 server_alive_interval: Duration,
867) -> Result<CommandResult> {
868 use std::process::Stdio;
869 use tokio::io::AsyncReadExt;
870 use tokio::process::Command;
871
872 let prepared = remote_shell_command(config, command);
873 let argv = system_ssh_argv(config, &prepared, server_alive_interval);
874 let input = prepared.stdin_script;
875 let args = argv.into_iter().skip(1);
877
878 let start = std::time::Instant::now();
879 debug!(
880 "system-ssh executing on {}: {}",
881 config.id,
882 crate::util::mask_sensitive_command(command)
883 );
884
885 let mut child = Command::new("ssh")
886 .args(args)
887 .env_clear()
888 .stdin(if input.is_some() {
889 Stdio::piped()
890 } else {
891 Stdio::null()
892 })
893 .stdout(Stdio::piped())
894 .stderr(Stdio::piped())
895 .kill_on_drop(true)
896 .spawn()
897 .with_context(|| {
898 format!(
899 "failed to spawn system ssh for {} (Windows fallback)",
900 config.id
901 )
902 })?;
903
904 let execution = async {
905 let stdout = child.stdout.take();
906 let stderr = child.stderr.take();
907 let stdin = child.stdin.take();
908
909 let stdin_fut = async {
910 if let (Some(stdin), Some(input)) = (stdin, input) {
911 write_ssh_command_stdin(stdin, &input).await?;
912 }
913 Ok::<(), anyhow::Error>(())
914 };
915
916 let stdout_fut = async {
917 let mut buf = String::new();
918 if let Some(out) = stdout {
919 let mut handle = out.take(MAX_OUTPUT_SIZE);
920 handle.read_to_string(&mut buf).await?;
921 }
922 Ok::<String, anyhow::Error>(buf)
923 };
924 let stderr_fut = async {
925 let mut buf = String::new();
926 if let Some(err) = stderr {
927 let mut handle = err.take(MAX_OUTPUT_SIZE);
928 handle.read_to_string(&mut buf).await?;
929 }
930 Ok::<String, anyhow::Error>(buf)
931 };
932
933 let ((), stdout, stderr) = tokio::try_join!(stdin_fut, stdout_fut, stderr_fut)?;
934 let status = child
935 .wait()
936 .await
937 .with_context(|| "system ssh command failed to wait for completion")?;
938 Ok::<_, anyhow::Error>((status, stdout, stderr))
939 };
940
941 match tokio::time::timeout(command_timeout, execution).await {
942 Ok(result) => {
943 let (status, stdout, stderr) = result?;
944 let duration = start.elapsed();
945 let exit_code = status.code().unwrap_or(-1);
946 debug!(
947 "system-ssh command completed on {} (exit={}, duration={}ms)",
948 config.id,
949 exit_code,
950 duration.as_millis()
951 );
952 Ok(CommandResult {
953 exit_code,
954 stdout,
955 stderr,
956 duration_ms: duration.as_millis() as u64,
957 })
958 }
959 Err(_) => {
960 warn!(
965 "system-ssh command timed out on {} after {:?}",
966 config.id, command_timeout
967 );
968 anyhow::bail!("Command timed out after {:?}", command_timeout);
969 }
970 }
971}
972
973pub struct SshPool {
975 connections: Arc<RwLock<HashMap<WorkerId, Arc<RwLock<SshClient>>>>>,
977 options: SshOptions,
979}
980
981impl SshPool {
982 pub fn new(options: SshOptions) -> Self {
984 Self {
985 connections: Arc::new(RwLock::new(HashMap::new())),
986 options,
987 }
988 }
989
990 pub async fn get_or_connect(&self, config: &WorkerConfig) -> Result<Arc<RwLock<SshClient>>> {
1002 self.get_or_connect_probing_within(config, POOL_LIVENESS_PROBE_TIMEOUT)
1003 .await
1004 }
1005
1006 async fn get_or_connect_probing_within(
1016 &self,
1017 config: &WorkerConfig,
1018 probe_timeout: Duration,
1019 ) -> Result<Arc<RwLock<SshClient>>> {
1020 let shared_client = self.get_or_create_client_entry(config).await;
1021
1022 let reusable = {
1026 let guard = shared_client.read().await;
1027 guard.is_connected()
1028 && guard
1029 .health_check_with_timeout(probe_timeout)
1030 .await
1031 .unwrap_or(false)
1032 };
1033 if reusable {
1034 debug!("Reusing live connection to {}", config.id);
1035 return Ok(shared_client);
1036 }
1037
1038 let mut client_guard = shared_client.write().await;
1043 if !client_guard.is_connected() {
1044 client_guard.connect().await?;
1045 } else if !client_guard
1046 .health_check_with_timeout(probe_timeout)
1047 .await
1048 .unwrap_or(false)
1049 {
1050 warn!(
1051 "Pooled SSH connection to {} failed liveness probe; reconnecting",
1052 config.id
1053 );
1054 client_guard.reconnect().await?;
1055 }
1056 drop(client_guard);
1058
1059 Ok(shared_client)
1060 }
1061
1062 async fn get_or_create_client_entry(&self, config: &WorkerConfig) -> Arc<RwLock<SshClient>> {
1063 let worker_id = config.id.clone();
1064
1065 loop {
1066 let existing_client = {
1067 let connections = self.connections.read().await;
1068 connections.get(&worker_id).cloned()
1069 };
1070
1071 if let Some(client) = existing_client {
1072 let is_configured_for_worker = {
1073 let guard = client.read().await;
1074 guard.is_configured_for(config)
1075 };
1076 if is_configured_for_worker {
1077 return client;
1078 }
1079
1080 let replacement = Arc::new(RwLock::new(SshClient::new(
1081 config.clone(),
1082 pooled_client_options(&self.options, config),
1083 )));
1084 let replaced = {
1085 let mut connections = self.connections.write().await;
1086 if connections
1087 .get(&worker_id)
1088 .is_some_and(|current| Arc::ptr_eq(current, &client))
1089 {
1090 connections.insert(worker_id.clone(), replacement.clone());
1091 true
1092 } else {
1093 false
1094 }
1095 };
1096
1097 if replaced {
1098 debug!(
1099 "Replaced SSH connection entry for {} after endpoint config changed",
1100 worker_id
1101 );
1102 return replacement;
1103 }
1104
1105 continue;
1106 }
1107
1108 let new_client = Arc::new(RwLock::new(SshClient::new(
1109 config.clone(),
1110 pooled_client_options(&self.options, config),
1111 )));
1112 let inserted = {
1113 let mut connections = self.connections.write().await;
1114 if connections.contains_key(&worker_id) {
1115 false
1116 } else {
1117 connections.insert(worker_id.clone(), new_client.clone());
1118 true
1119 }
1120 };
1121
1122 if inserted {
1123 return new_client;
1124 }
1125 }
1126 }
1127
1128 pub async fn close(&self, worker_id: &WorkerId) -> Result<()> {
1130 let client = {
1131 let mut connections = self.connections.write().await;
1132 connections.remove(worker_id)
1133 };
1134
1135 if let Some(client) = client {
1136 let mut client = client.write().await;
1137 client.disconnect().await?;
1138 }
1139
1140 Ok(())
1141 }
1142
1143 pub async fn close_all(&self) -> Result<()> {
1145 let clients: Vec<_> = {
1146 let mut connections = self.connections.write().await;
1147 connections.drain().map(|(_, v)| v).collect()
1148 };
1149
1150 for client in clients {
1151 let mut client = client.write().await;
1152 if let Err(e) = client.disconnect().await {
1153 error!("Error closing connection: {}", e);
1154 }
1155 }
1156
1157 Ok(())
1158 }
1159
1160 pub async fn active_connections(&self) -> usize {
1162 self.connections.read().await.len()
1163 }
1164
1165 pub async fn run_with_timeout(
1185 &self,
1186 config: &WorkerConfig,
1187 command: &str,
1188 command_timeout: Duration,
1189 ) -> Result<CommandResult> {
1190 let client = self
1191 .get_or_connect_probing_within(config, command_timeout.min(POOL_LIVENESS_PROBE_TIMEOUT))
1192 .await?;
1193 let guard = client.read().await;
1197 guard.execute_with_timeout(command, command_timeout).await
1198 }
1199}
1200
1201impl Default for SshPool {
1202 fn default() -> Self {
1203 Self::new(SshOptions::default())
1204 }
1205}
1206
1207#[cfg(test)]
1208mod tests {
1209 use super::*;
1210 use crate::test_guard;
1211
1212 #[test]
1213 fn test_command_result_success() {
1214 let _guard = test_guard!();
1215 let result = CommandResult {
1216 exit_code: 0,
1217 stdout: "output".to_string(),
1218 stderr: String::new(),
1219 duration_ms: 100,
1220 };
1221 assert!(result.success());
1222
1223 let failed = CommandResult {
1224 exit_code: 1,
1225 stdout: String::new(),
1226 stderr: "error".to_string(),
1227 duration_ms: 50,
1228 };
1229 assert!(!failed.success());
1230 }
1231
1232 #[test]
1233 fn test_ssh_options_default() {
1234 let _guard = test_guard!();
1235 let options = SshOptions::default();
1236 assert_eq!(options.connect_timeout, Duration::from_secs(10));
1237 assert_eq!(options.command_timeout, Duration::from_secs(300));
1238 assert!(options.server_alive_interval.is_none());
1239 assert!(options.control_persist_idle.is_none());
1240 assert!(!options.control_master);
1241 }
1242
1243 #[test]
1244 fn test_ssh_client_creation() {
1245 let _guard = test_guard!();
1246 let config = WorkerConfig {
1247 id: WorkerId::new("test-worker"),
1248 host: "192.168.1.100".to_string(),
1249 user: "ubuntu".to_string(),
1250 identity_file: "~/.ssh/id_rsa".to_string(),
1251 total_slots: 8,
1252 priority: 100,
1253 tags: vec!["rust".to_string()],
1254 tools: Vec::new(),
1255 };
1256
1257 let client = SshClient::new(config.clone(), SshOptions::default());
1258 assert_eq!(client.worker_id().as_str(), "test-worker");
1259 assert!(!client.is_connected());
1260 }
1261
1262 #[test]
1263 fn test_expected_health_check_output_accepts_sentinel_as_last_line() {
1264 let _guard = test_guard!();
1265
1266 assert!(is_expected_health_check_output("ok\n"));
1267 assert!(is_expected_health_check_output("login banner\nok\n"));
1268 assert!(!is_expected_health_check_output(""));
1269 assert!(!is_expected_health_check_output("not ok\n"));
1270 assert!(!is_expected_health_check_output("ok\npost-command noise\n"));
1271 }
1272
1273 fn worker_config(id: &str, host: &str, user: &str, identity_file: &str) -> WorkerConfig {
1274 WorkerConfig {
1275 id: WorkerId::new(id),
1276 host: host.to_string(),
1277 user: user.to_string(),
1278 identity_file: identity_file.to_string(),
1279 total_slots: 8,
1280 priority: 100,
1281 tags: vec!["rust".to_string()],
1282 tools: Vec::new(),
1283 }
1284 }
1285
1286 #[test]
1287 fn test_ssh_client_configured_for_ignores_scheduling_fields() {
1288 let _guard = test_guard!();
1289 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1290 let client = SshClient::new(config.clone(), SshOptions::default());
1291
1292 let mut scheduling_only_change = config;
1293 scheduling_only_change.total_slots = 16;
1294 scheduling_only_change.priority = 250;
1295 scheduling_only_change.tags = vec!["rust".to_string(), "gpu".to_string()];
1296
1297 assert!(client.is_configured_for(&scheduling_only_change));
1298 }
1299
1300 #[test]
1301 fn test_ssh_client_configured_for_detects_endpoint_changes() {
1302 let _guard = test_guard!();
1303 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1304 let client = SshClient::new(config, SshOptions::default());
1305
1306 assert!(!client.is_configured_for(&worker_config(
1307 "worker-a",
1308 "192.168.1.101",
1309 "ubuntu",
1310 "~/.ssh/id_rsa",
1311 )));
1312 assert!(!client.is_configured_for(&worker_config(
1313 "worker-a",
1314 "192.168.1.100",
1315 "admin",
1316 "~/.ssh/id_rsa",
1317 )));
1318 assert!(!client.is_configured_for(&worker_config(
1319 "worker-a",
1320 "192.168.1.100",
1321 "ubuntu",
1322 "~/.ssh/other_key",
1323 )));
1324 }
1325
1326 #[tokio::test]
1327 async fn test_ssh_pool_reuses_matching_disconnected_entry() {
1328 let _guard = test_guard!();
1329 let pool = SshPool::default();
1330 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1331
1332 let first = pool.get_or_create_client_entry(&config).await;
1333 let second = pool.get_or_create_client_entry(&config).await;
1334
1335 assert!(Arc::ptr_eq(&first, &second));
1336 assert_eq!(pool.active_connections().await, 1);
1337 }
1338
1339 #[tokio::test]
1340 async fn test_ssh_pool_reuses_single_entry_and_close_all_drops_to_zero() {
1341 let _guard = test_guard!();
1347 let options = SshOptions {
1348 control_master: true,
1349 control_persist_idle: Some(Duration::from_secs(60)),
1350 ..Default::default()
1351 };
1352 let pool = SshPool::new(options);
1353 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1354
1355 for _ in 0..5 {
1356 let _entry = pool.get_or_create_client_entry(&config).await;
1357 }
1358 assert_eq!(
1359 pool.active_connections().await,
1360 1,
1361 "repeated borrows of one worker must reuse a single pool entry"
1362 );
1363
1364 pool.close_all().await.expect("close_all should succeed");
1365 assert_eq!(
1366 pool.active_connections().await,
1367 0,
1368 "close_all must drop all pooled connections"
1369 );
1370 }
1371
1372 #[test]
1373 fn test_pool_options_map_to_bounded_idle_persist() {
1374 let _guard = test_guard!();
1378 assert_eq!(
1379 control_persist_mode(true, Some(Duration::from_secs(60))),
1380 ControlPersistMode::IdleFor(NonZeroUsize::new(60).unwrap()),
1381 "pool mux options must keep a warm master for a bounded idle window"
1382 );
1383 }
1384
1385 fn windows_worker_config(id: &str) -> WorkerConfig {
1390 let mut config = worker_config(id, "100.68.2.11", "jeffr", "~/.ssh/surfacebookje_key");
1391 config.tags = vec!["rust".to_string(), crate::types::os_tag("windows")];
1392 config
1393 }
1394
1395 #[test]
1396 fn test_pooled_options_disable_control_master_for_windows_workers() {
1397 let _guard = test_guard!();
1403 let pool_mux = SshOptions {
1404 control_master: true,
1405 control_persist_idle: Some(Duration::from_secs(60)),
1406 ..Default::default()
1407 };
1408 let config = windows_worker_config("wsurf");
1409
1410 let effective = pooled_client_options(&pool_mux, &config);
1411 assert!(
1412 !effective.control_master,
1413 "a Windows worker must never get a pooled ControlMaster"
1414 );
1415 assert_eq!(effective.connect_timeout, pool_mux.connect_timeout);
1419 assert_eq!(effective.command_timeout, pool_mux.command_timeout);
1420 assert_eq!(
1421 effective.server_alive_interval,
1422 pool_mux.server_alive_interval
1423 );
1424 assert_eq!(
1425 effective.control_persist_idle,
1426 pool_mux.control_persist_idle
1427 );
1428 assert_eq!(effective.known_hosts, pool_mux.known_hosts);
1429
1430 let pool_plain = SshOptions {
1432 control_master: false,
1433 ..pool_mux
1434 };
1435 assert!(!pooled_client_options(&pool_plain, &config).control_master);
1436
1437 let mut mixed_case = windows_worker_config("wsurf-2");
1440 mixed_case.tags = vec![crate::types::os_tag("Windows")];
1441 assert!(!pooled_client_options(&pool_mux, &mixed_case).control_master);
1442 }
1443
1444 #[test]
1445 fn test_pooled_options_keep_pool_options_verbatim_for_non_windows_workers() {
1446 let _guard = test_guard!();
1450 let pool_mux = SshOptions {
1451 control_master: true,
1452 control_persist_idle: Some(Duration::from_secs(60)),
1453 ..Default::default()
1454 };
1455 let linux = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1456 let unlabelled = worker_config("contabo-b", "1.2.3.5", "root", "~/.ssh/id_rsa");
1457
1458 for config in [&linux, &unlabelled] {
1459 let effective = pooled_client_options(&pool_mux, config);
1460 assert!(
1461 effective.control_master,
1462 "non-Windows workers keep the pool's mux setting"
1463 );
1464 assert_eq!(
1465 effective.control_persist_idle,
1466 pool_mux.control_persist_idle
1467 );
1468 }
1469
1470 let pool_plain = SshOptions {
1471 control_master: false,
1472 ..pool_mux
1473 };
1474 for config in [&linux, &unlabelled] {
1475 assert!(!pooled_client_options(&pool_plain, config).control_master);
1476 }
1477 }
1478
1479 #[tokio::test]
1480 async fn test_pool_entry_for_windows_worker_is_never_mux() {
1481 let _guard = test_guard!();
1487 let pool = SshPool::new(SshOptions {
1488 control_master: true,
1489 control_persist_idle: Some(Duration::from_secs(60)),
1490 ..Default::default()
1491 });
1492 let config = windows_worker_config("wsurf");
1493
1494 let entry = pool.get_or_create_client_entry(&config).await;
1495 let guard = entry.read().await;
1496 assert!(
1497 !guard.options.control_master,
1498 "pooled entry for a Windows worker must be non-mux"
1499 );
1500 assert_eq!(
1503 guard.options.control_persist_idle,
1504 Some(Duration::from_secs(60))
1505 );
1506 }
1507
1508 #[tokio::test]
1509 async fn test_pool_entry_for_linux_worker_keeps_mux() {
1510 let _guard = test_guard!();
1511 let pool = SshPool::new(SshOptions {
1512 control_master: true,
1513 control_persist_idle: Some(Duration::from_secs(60)),
1514 ..Default::default()
1515 });
1516 let config = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1517
1518 let entry = pool.get_or_create_client_entry(&config).await;
1519 let guard = entry.read().await;
1520 assert!(
1521 guard.options.control_master,
1522 "pooled entry for a Linux worker must keep the pool's mux setting"
1523 );
1524 }
1525
1526 #[tokio::test]
1527 async fn test_ssh_pool_replaces_stale_entry_when_endpoint_changes() {
1528 let _guard = test_guard!();
1529 let pool = SshPool::default();
1530 let old_config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1531 let new_config = worker_config("worker-a", "192.168.1.101", "admin", "~/.ssh/new_key");
1532
1533 let stale = pool.get_or_create_client_entry(&old_config).await;
1534 let replacement = pool.get_or_create_client_entry(&new_config).await;
1535
1536 assert!(!Arc::ptr_eq(&stale, &replacement));
1537 assert_eq!(pool.active_connections().await, 1);
1538
1539 let replacement_guard = replacement.read().await;
1540 assert!(replacement_guard.is_configured_for(&new_config));
1541 }
1542
1543 #[tokio::test]
1544 async fn test_ssh_pool_os_transition_replaces_transport_but_case_and_scheduling_reuse() {
1545 let _guard = test_guard!();
1546 let pool = SshPool::new(SshOptions {
1547 control_master: true,
1548 ..Default::default()
1549 });
1550 let mut config = worker_config("worker-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1551 let unlabelled = pool.get_or_create_client_entry(&config).await;
1552 config.tags.push(crate::types::os_tag("windows"));
1553 let windows = pool.get_or_create_client_entry(&config).await;
1554 assert!(!Arc::ptr_eq(&unlabelled, &windows));
1555 {
1556 let client = windows.read().await;
1557 assert!(prefers_system_ssh(&client.config));
1558 assert!(!client.options.control_master);
1559 }
1560
1561 config.tags = vec!["gpu".into(), "os: WiNdOwS ".into()];
1562 config.total_slots += 1;
1563 config.priority += 1;
1564 let reused = pool.get_or_create_client_entry(&config).await;
1565 assert!(Arc::ptr_eq(&windows, &reused));
1566
1567 config.tags = vec![crate::types::os_tag("linux")];
1568 let linux = pool.get_or_create_client_entry(&config).await;
1569 assert!(!Arc::ptr_eq(&windows, &linux));
1570 {
1571 let client = linux.read().await;
1572 assert!(!prefers_system_ssh(&client.config));
1573 assert!(client.options.control_master);
1574 }
1575 assert!(prefers_system_ssh(&windows.read().await.config));
1578 assert!(unlabelled.read().await.options.control_master);
1579 assert_eq!(pool.active_connections().await, 1);
1580
1581 config.tags.clear();
1582 let removed = pool.get_or_create_client_entry(&config).await;
1583 assert!(!Arc::ptr_eq(&linux, &removed));
1584 assert!(removed.read().await.options.control_master);
1585 assert_eq!(pool.active_connections().await, 1);
1586 }
1587
1588 #[tokio::test]
1589 async fn test_health_check_reports_not_alive_without_session() {
1590 let _guard = test_guard!();
1596 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1597 let client = SshClient::new(config, SshOptions::default());
1598 assert!(!client.is_connected());
1599 let alive = client
1600 .health_check()
1601 .await
1602 .expect("health_check maps execution errors to Ok(false), never Err");
1603 assert!(!alive, "a client with no session must not report as alive");
1604 }
1605
1606 #[test]
1607 fn test_build_env_prefix_quotes_and_rejects() {
1608 let _guard = test_guard!();
1609 let mut env = HashMap::new();
1610 env.insert("RUSTFLAGS".to_string(), "-C target-cpu=native".to_string());
1611 env.insert("QUOTED".to_string(), "a'b".to_string());
1612 env.insert("BADVAL".to_string(), "line1\nline2".to_string());
1613
1614 let allowlist = vec![
1615 "RUSTFLAGS".to_string(),
1616 "QUOTED".to_string(),
1617 "MISSING".to_string(),
1618 "BADVAL".to_string(),
1619 "BAD=KEY".to_string(),
1620 ];
1621
1622 let prefix = build_env_prefix(&allowlist, |key| env.get(key).cloned());
1623
1624 assert!(prefix.prefix.contains("RUSTFLAGS='-C target-cpu=native'"));
1625 assert!(prefix.prefix.contains("QUOTED='a'\\''b'"));
1627 assert!(!prefix.prefix.contains("MISSING="));
1628 assert!(!prefix.prefix.contains("BADVAL="));
1629 assert!(prefix.rejected.contains(&"BADVAL".to_string()));
1630 assert!(prefix.rejected.contains(&"BAD=KEY".to_string()));
1631 assert_eq!(
1632 prefix.applied,
1633 vec!["RUSTFLAGS".to_string(), "QUOTED".to_string()]
1634 );
1635 }
1636
1637 mod proptest_ssh_escaping {
1642 use super::*;
1643 use proptest::prelude::*;
1644 use std::collections::HashMap;
1645
1646 proptest! {
1647 #![proptest_config(ProptestConfig::with_cases(1000))]
1648
1649 #[test]
1651 fn test_is_valid_env_key_no_panic(s in ".*") {
1652 let _guard = test_guard!();
1653 let _ = is_valid_env_key(&s);
1654 }
1655
1656 #[test]
1658 fn test_is_valid_env_key_accepts_valid(
1659 first in "[a-zA-Z_]",
1660 rest in "[a-zA-Z0-9_]{0,50}"
1661 ) {
1662 let _guard = test_guard!();
1663 let key = format!("{}{}", first, rest);
1664 prop_assert!(is_valid_env_key(&key), "Should accept valid key: {}", key);
1665 }
1666
1667 #[test]
1669 fn test_is_valid_env_key_rejects_digit_start(
1670 digit in "[0-9]",
1671 rest in "[a-zA-Z0-9_]{0,20}"
1672 ) {
1673 let _guard = test_guard!();
1674 let key = format!("{}{}", digit, rest);
1675 prop_assert!(!is_valid_env_key(&key), "Should reject digit-start key: {}", key);
1676 }
1677
1678 #[test]
1680 fn test_shell_escape_value_no_panic(s in ".*") {
1681 let _guard = test_guard!();
1682 let _ = shell_escape_value(&s);
1683 }
1684
1685 #[test]
1687 fn test_shell_escape_value_rejects_unsafe(
1688 prefix in "[a-zA-Z0-9 ]{0,10}",
1689 bad_char in "[\n\r\0]",
1690 suffix in "[a-zA-Z0-9 ]{0,10}"
1691 ) {
1692 let _guard = test_guard!();
1693 let value = format!("{}{}{}", prefix, bad_char, suffix);
1694 prop_assert!(shell_escape_value(&value).is_none(),
1695 "Should reject value with unsafe char: {:?}", value);
1696 }
1697
1698 #[test]
1700 fn test_shell_escape_value_accepts_safe(s in "[a-zA-Z0-9 !@#$%^&*()_+=\\-\\[\\]{}|;:,./<>?]{0,100}") {
1701 let result = shell_escape_value(&s);
1703 prop_assert!(result.is_some(), "Should accept safe value: {:?}", s);
1704
1705 let escaped = match result {
1708 Some(escaped) => escaped,
1709 None => {
1710 prop_assert!(false, "Should accept safe value: {:?}", s);
1711 String::new()
1712 }
1713 };
1714 if s.chars().any(|c| !c.is_ascii_alphanumeric() && c != '_') {
1715 prop_assert!(escaped.starts_with('\'') || escaped.contains('\''),
1717 "Value with special chars should be quoted: {:?} -> {:?}", s, escaped);
1718 }
1719 }
1720
1721 #[test]
1723 fn test_shell_escape_value_escapes_quotes(
1724 prefix in "[a-zA-Z0-9]{0,10}",
1725 suffix in "[a-zA-Z0-9]{0,10}"
1726 ) {
1727 let _guard = test_guard!();
1728 let value = format!("{}'{}", prefix, suffix);
1729 let result = shell_escape_value(&value);
1730 prop_assert!(result.is_some());
1731
1732 let escaped = match result {
1733 Some(escaped) => escaped,
1734 None => {
1735 prop_assert!(false, "Should escape single quote: {}", value);
1736 String::new()
1737 }
1738 };
1739 prop_assert!(escaped.contains("'\\''"),
1741 "Should escape single quote: {} -> {}", value, escaped);
1742 }
1743
1744 #[test]
1746 fn test_build_env_prefix_no_panic(
1747 keys in prop::collection::vec("[a-zA-Z_][a-zA-Z0-9_]{0,10}", 0..10),
1748 values in prop::collection::vec(".*", 0..10)
1749 ) {
1750 let mut env = HashMap::new();
1751 for (i, key) in keys.iter().enumerate() {
1752 if let Some(val) = values.get(i) {
1753 env.insert(key.clone(), val.clone());
1754 }
1755 }
1756
1757 let allowlist: Vec<String> = keys;
1758 let _ = build_env_prefix(&allowlist, |k| env.get(k).cloned());
1759 }
1760
1761 #[test]
1763 fn test_build_env_prefix_rejects_invalid_keys(
1764 invalid_key in "[0-9][a-zA-Z0-9_]{0,10}" ) {
1767 let _guard = test_guard!();
1768 let mut env = HashMap::new();
1769 env.insert(invalid_key.clone(), "value".to_string());
1770
1771 let allowlist = vec![invalid_key.clone()];
1772 let prefix = build_env_prefix(&allowlist, |k| env.get(k).cloned());
1773
1774 prop_assert!(!is_valid_env_key(&invalid_key),
1776 "Key should be invalid: {}", invalid_key);
1777 prop_assert!(prefix.rejected.contains(&invalid_key),
1778 "Should reject invalid key: {}", invalid_key);
1779 prop_assert!(prefix.prefix.is_empty());
1780 }
1781
1782 #[test]
1784 fn test_build_env_prefix_missing_values(
1785 keys in prop::collection::vec("[A-Z_][A-Z0-9_]{0,10}", 1..5)
1786 ) {
1787 let env: HashMap<String, String> = HashMap::new();
1789 let prefix = build_env_prefix(&keys, |k| env.get(k).cloned());
1790
1791 prop_assert!(prefix.prefix.is_empty(), "Should be empty when no values");
1793 prop_assert!(prefix.applied.is_empty());
1794 prop_assert!(prefix.rejected.is_empty());
1796 }
1797 }
1798
1799 #[test]
1801 fn test_shell_escape_edge_cases() {
1802 let _guard = test_guard!();
1803 let result = shell_escape_value("");
1805 assert_eq!(result, Some("''".to_string()));
1806
1807 let result = shell_escape_value("'");
1809 assert_eq!(result, Some("''\\'''".to_string()));
1810
1811 let result = shell_escape_value("'''");
1813 assert_eq!(
1815 result
1816 .as_deref()
1817 .map(|escaped| escaped.matches("'\\''").count()),
1818 Some(3)
1819 );
1820
1821 let result = shell_escape_value("日本語");
1823 assert!(result.is_some());
1824
1825 let result = shell_escape_value("🔥🚀");
1827 assert!(result.is_some());
1828
1829 let result = shell_escape_value("it's a \"test\" with $vars");
1831 assert!(result.is_some());
1832 }
1833
1834 #[test]
1835 fn test_is_valid_env_key_edge_cases() {
1836 let _guard = test_guard!();
1837 assert!(!is_valid_env_key(""));
1839
1840 assert!(is_valid_env_key("_"));
1842
1843 assert!(is_valid_env_key("A"));
1845
1846 assert!(is_valid_env_key("PATH"));
1848 assert!(is_valid_env_key("HOME"));
1849 assert!(is_valid_env_key("RUSTFLAGS"));
1850 assert!(is_valid_env_key("CC"));
1851 assert!(is_valid_env_key("_PRIVATE"));
1852 assert!(is_valid_env_key("MY_VAR_123"));
1853
1854 assert!(!is_valid_env_key("1VAR"));
1856 assert!(!is_valid_env_key("123"));
1857
1858 assert!(!is_valid_env_key("MY-VAR"));
1860 assert!(!is_valid_env_key("MY.VAR"));
1861 assert!(!is_valid_env_key("MY VAR"));
1862 assert!(!is_valid_env_key("MY=VAR"));
1863
1864 assert!(!is_valid_env_key("日本語"));
1866 assert!(!is_valid_env_key("VAR🔥"));
1867 }
1868
1869 #[test]
1870 fn test_build_env_prefix_integration() {
1871 let _guard = test_guard!();
1872 let mut env = HashMap::new();
1874 env.insert("VALID".to_string(), "simple".to_string());
1875 env.insert("WITH_QUOTE".to_string(), "it's here".to_string());
1876 env.insert("NEWLINE".to_string(), "line1\nline2".to_string());
1877 env.insert("UNICODE".to_string(), "日本語".to_string());
1878 env.insert("EMPTY".to_string(), String::new());
1879 env.insert("123INVALID".to_string(), "value".to_string());
1880
1881 let allowlist = vec![
1882 "VALID".to_string(),
1883 "WITH_QUOTE".to_string(),
1884 "NEWLINE".to_string(),
1885 "UNICODE".to_string(),
1886 "EMPTY".to_string(),
1887 "123INVALID".to_string(),
1888 "MISSING".to_string(),
1889 ];
1890
1891 let prefix = build_env_prefix(&allowlist, |k| env.get(k).cloned());
1892
1893 assert!(prefix.applied.contains(&"VALID".to_string()));
1895 assert!(prefix.prefix.contains("VALID=simple"));
1897
1898 assert!(prefix.applied.contains(&"WITH_QUOTE".to_string()));
1900
1901 assert!(prefix.rejected.contains(&"NEWLINE".to_string()));
1903
1904 assert!(prefix.applied.contains(&"UNICODE".to_string()));
1906
1907 assert!(prefix.applied.contains(&"EMPTY".to_string()));
1909
1910 assert!(prefix.rejected.contains(&"123INVALID".to_string()));
1912
1913 assert!(!prefix.applied.contains(&"MISSING".to_string()));
1915 assert!(!prefix.rejected.contains(&"MISSING".to_string()));
1916 }
1917
1918 #[test]
1919 fn test_shell_escape_roundtrip_safety() {
1920 let _guard = test_guard!();
1921 let test_values = [
1923 "simple",
1924 "with spaces",
1925 "with\ttab",
1926 "special!@#$%^&*()",
1927 "quoted\"value",
1928 "path/to/file",
1929 "-flag",
1930 "--long-flag=value",
1931 "",
1932 ];
1933
1934 for value in &test_values {
1935 let escaped = shell_escape_value(value);
1936 assert!(escaped.is_some(), "Should escape: {:?}", value);
1937 }
1938 }
1939 }
1940
1941 fn windows_worker(id: &str) -> WorkerConfig {
1946 let mut config = worker_config(id, "100.68.2.11", "jeffr", "~/.ssh/surfacebookje_key");
1947 config.tags = vec!["rust".to_string(), crate::types::os_tag("windows")];
1948 config
1949 }
1950
1951 #[test]
1952 fn test_prefers_system_ssh_for_windows_tag() {
1953 let _guard = test_guard!();
1956 let cfg = windows_worker("wsurf");
1957 assert!(prefers_system_ssh(&cfg));
1958 }
1959
1960 #[test]
1961 fn test_prefers_system_ssh_case_normalized() {
1962 let _guard = test_guard!();
1965 let mut cfg = windows_worker("wsurf-mixed");
1966 cfg.tags = vec![crate::types::os_tag("Windows")];
1967 assert!(prefers_system_ssh(&cfg));
1968
1969 let mut upper = windows_worker("wsurf-upper");
1970 upper.tags = vec![crate::types::os_tag("WINDOWS")];
1971 assert!(prefers_system_ssh(&upper));
1972 }
1973
1974 #[test]
1975 fn test_prefers_system_ssh_false_for_linux() {
1976 let _guard = test_guard!();
1979 let mut linux = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1980 linux.tags = vec!["rust".to_string(), crate::types::os_tag("linux")];
1981 assert!(!prefers_system_ssh(&linux));
1982
1983 let unlabelled = worker_config("contabo-b", "1.2.3.5", "root", "~/.ssh/id_rsa");
1984 assert!(!prefers_system_ssh(&unlabelled));
1985 }
1986
1987 #[test]
1988 fn test_system_ssh_argv_basic_windows() {
1989 let _guard = test_guard!();
1994 let cfg = windows_worker("wsurf");
1995 let command = remote_shell_command(&cfg, "uname -a");
1996 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(2));
1997 let s: Vec<String> = argv
1998 .iter()
1999 .map(|o| o.to_string_lossy().into_owned())
2000 .collect();
2001
2002 assert_eq!(s[0], "ssh", "argv[0] is the program name");
2003 assert_eq!(s[1], "-i");
2004 assert!(
2007 s[2].ends_with(".ssh/surfacebookje_key"),
2008 "identity file is the tilde-expanded path: {}",
2009 s[2]
2010 );
2011 assert!(s.contains(&"-o".to_string()));
2014 assert!(s.contains(&"BatchMode=yes".to_string()));
2015 assert!(s.contains(&"ConnectTimeout=8".to_string()));
2016 assert!(s.contains(&"StrictHostKeyChecking=accept-new".to_string()));
2017 assert!(s.contains(&"ServerAliveInterval=2".to_string()));
2018 assert!(!s.iter().any(|arg| arg == "-p"));
2020 assert!(s.contains(&"jeffr@100.68.2.11".to_string()));
2022 assert_eq!(s[s.len() - 1], "sh -s", "Windows gets only a fixed reader");
2023 }
2024
2025 #[test]
2026 fn test_system_ssh_script_transport_keeps_posix_argv_and_exact_stdin_reader() {
2027 let _guard = test_guard!();
2028 let linux = worker_config("linux", "1.2.3.4", "root", "~/.ssh/id_rsa");
2029 let windows = windows_worker("windows");
2030 let script = "printf '%s\\n' '%PATH% !VAR! & | < > ^ $()'\nprintf done";
2031 let prepared = remote_shell_command(&linux, script);
2032 let argv = system_ssh_argv(&linux, &prepared, Duration::ZERO);
2033 assert_eq!(argv.last().unwrap(), script);
2034 assert!(prepared.stdin_script.is_none());
2035 for command in [script, "sh -s", "", "sh -s\n"] {
2036 let prepared = remote_shell_command(&windows, command);
2037 let argv = system_ssh_argv(&windows, &prepared, Duration::ZERO);
2038 assert_eq!(argv.last().unwrap(), "sh -s");
2039 assert_eq!(prepared.stdin_script.is_none(), command == "sh -s");
2040 }
2041 }
2042
2043 #[tokio::test]
2044 async fn test_system_ssh_stdin_roundtrip_matches_null_stdin_shell() {
2045 use std::process::Stdio;
2046
2047 let _guard = test_guard!();
2048 let config = windows_worker("windows");
2049 for script in [
2052 "",
2053 "printf '%s\\n' \"apostrophe's\" 'double\"quote' '%PATH% !VAR! ^ & | < >' '$() `x` \\'",
2054 "cat\nprintf 'AFTER\\n'\nread value || printf 'EMPTY\\n'",
2055 "cat <<'END'\nheredoc ' \" % ! $() \\\nEND\nprintf 'done\\n'",
2056 "printf '%s\\n' 'first\nsecond' # trailing comment",
2057 "printf 'before\\n'; printf 'diagnostic\\n' >&2; exit 37",
2058 "# comment ending with a backslash \\",
2059 ] {
2060 let expected = tokio::process::Command::new("sh")
2061 .args(["-c", script])
2062 .stdin(Stdio::null())
2063 .output()
2064 .await
2065 .unwrap();
2066 let mut child = tokio::process::Command::new("sh")
2067 .arg("-s")
2068 .stdin(Stdio::piped())
2069 .stdout(Stdio::piped())
2070 .stderr(Stdio::piped())
2071 .kill_on_drop(true)
2072 .spawn()
2073 .unwrap();
2074 let stdin = child.stdin.take().unwrap();
2075 let payload = remote_shell_command(&config, script).stdin_script.unwrap();
2076 let (written, observed) = tokio::time::timeout(Duration::from_secs(5), async {
2077 tokio::join!(
2078 write_ssh_command_stdin(stdin, &payload),
2079 child.wait_with_output()
2080 )
2081 })
2082 .await
2083 .unwrap();
2084 written.unwrap();
2085 let observed = observed.unwrap();
2086 assert_eq!(observed.status.code(), expected.status.code(), "{script:?}");
2087 assert_eq!(observed.stdout, expected.stdout, "{script:?}");
2088 assert_eq!(observed.stderr, expected.stderr, "{script:?}");
2089 }
2090 }
2091
2092 #[tokio::test]
2093 async fn test_system_ssh_split_stdin_never_becomes_command_input() {
2094 use std::process::Stdio;
2095 use tokio::io::AsyncWriteExt;
2096
2097 let _guard = test_guard!();
2098 let script = "cat\nprintf 'AFTER\\n'\nread value || printf 'EMPTY\\n'";
2099 let payload = remote_shell_command(&windows_worker("windows"), script)
2100 .stdin_script
2101 .unwrap();
2102 let split = payload.find("cat\n").unwrap() + "cat\n".len();
2103 let mut child = tokio::process::Command::new("sh")
2104 .arg("-s")
2105 .stdin(Stdio::piped())
2106 .stdout(Stdio::piped())
2107 .stderr(Stdio::piped())
2108 .kill_on_drop(true)
2109 .spawn()
2110 .unwrap();
2111 let mut stdin = child.stdin.take().unwrap();
2112 let writer = async {
2113 stdin.write_all(&payload.as_bytes()[..split]).await.unwrap();
2114 tokio::time::sleep(Duration::from_millis(100)).await;
2115 stdin.write_all(&payload.as_bytes()[split..]).await.unwrap();
2116 stdin.shutdown().await.unwrap();
2117 };
2118 let ((), output) = tokio::time::timeout(Duration::from_secs(5), async {
2119 tokio::join!(writer, child.wait_with_output())
2120 })
2121 .await
2122 .unwrap();
2123 let output = output.unwrap();
2124 assert!(output.status.success(), "{:?}", output.stderr);
2125 assert_eq!(output.stdout, b"AFTER\nEMPTY\n");
2126 assert!(output.stderr.is_empty());
2127 }
2128
2129 #[tokio::test]
2130 async fn test_system_ssh_closed_stdin_preserves_early_exit_diagnostic() {
2131 use std::process::Stdio;
2132
2133 let _guard = test_guard!();
2134 let mut child = tokio::process::Command::new("sh")
2135 .args(["-c", "printf 'remote failure\\n' >&2; exit 37"])
2136 .stdin(Stdio::piped())
2137 .stdout(Stdio::piped())
2138 .stderr(Stdio::piped())
2139 .kill_on_drop(true)
2140 .spawn()
2141 .unwrap();
2142 let stdin = child.stdin.take().unwrap();
2143 child.wait().await.unwrap();
2146 write_ssh_command_stdin(stdin, &"x".repeat(128 * 1024))
2147 .await
2148 .unwrap();
2149 let output = child.wait_with_output().await.unwrap();
2150 assert_eq!(output.status.code(), Some(37));
2151 assert_eq!(output.stderr, b"remote failure\n");
2152 assert!(output.stdout.is_empty());
2153 }
2154
2155 #[test]
2156 fn test_system_ssh_argv_zero_server_alive_omits_keepalive() {
2157 let _guard = test_guard!();
2161 let cfg = windows_worker("wsurf");
2162 let command = remote_shell_command(&cfg, "echo hi");
2163 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(0));
2164 let s: Vec<String> = argv
2165 .iter()
2166 .map(|o| o.to_string_lossy().into_owned())
2167 .collect();
2168 assert!(
2169 !s.iter().any(|a| a.starts_with("ServerAliveInterval=")),
2170 "no ServerAliveInterval flag when interval is 0; got {:?}",
2171 s
2172 );
2173 }
2174
2175 #[test]
2176 fn test_system_ssh_argv_long_keepalive_preserved() {
2177 let _guard = test_guard!();
2181 let cfg = windows_worker("wsurf");
2182 let command = remote_shell_command(&cfg, "true");
2183 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(60));
2184 let s: Vec<String> = argv
2185 .iter()
2186 .map(|o| o.to_string_lossy().into_owned())
2187 .collect();
2188 assert!(s.contains(&"ServerAliveInterval=60".to_string()));
2189 }
2190
2191 #[test]
2192 fn test_system_ssh_argv_uses_tilde_expanded_identity() {
2193 let _guard = test_guard!();
2198 let mut cfg = windows_worker("wsurf");
2199 cfg.identity_file = "~/.ssh/operator_key".to_string();
2200 let command = remote_shell_command(&cfg, "true");
2201 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(2));
2202 let s: Vec<String> = argv
2203 .iter()
2204 .map(|o| o.to_string_lossy().into_owned())
2205 .collect();
2206 assert!(
2207 !s.contains(&"~/.ssh/operator_key".to_string()),
2208 "identity file must be tilde-expanded, got {:?}",
2209 s
2210 );
2211 assert!(
2212 s.contains(&"~/.ssh/operator_key".to_string())
2213 || s.iter().any(|a| a.ends_with(".ssh/operator_key")),
2214 "tilde-expanded path appears in argv: {:?}",
2215 s
2216 );
2217 }
2218
2219 #[test]
2220 fn test_dispatch_helper_matches_declared_os_key() {
2221 let _guard = test_guard!();
2226 let cfg = windows_worker("wsurf");
2227 assert_eq!(declared_os(&cfg.tags).as_deref(), Some("windows"));
2228 assert!(prefers_system_ssh(&cfg));
2229 }
2230}
2231
2232#[derive(Debug, Clone, PartialEq, Eq)]
2234enum ControlPersistMode {
2235 Closed,
2237 IdleFor(NonZeroUsize),
2239 TooLarge(u64),
2241}
2242
2243fn control_persist_mode(control_master: bool, idle: Option<Duration>) -> ControlPersistMode {
2248 match idle {
2249 Some(idle) if control_master && !idle.is_zero() => match usize::try_from(idle.as_secs()) {
2250 Ok(secs) => match NonZeroUsize::new(secs) {
2251 Some(nonzero) => ControlPersistMode::IdleFor(nonzero),
2252 None => ControlPersistMode::Closed,
2253 },
2254 Err(_) => ControlPersistMode::TooLarge(idle.as_secs()),
2255 },
2256 _ => ControlPersistMode::Closed,
2257 }
2258}
2259
2260#[cfg(test)]
2261mod control_persist_tests {
2262 use super::*;
2263 use std::time::Duration;
2264
2265 #[test]
2266 fn non_mux_sessions_never_persist_forever() {
2267 assert_eq!(
2269 control_persist_mode(false, None),
2270 ControlPersistMode::Closed
2271 );
2272 assert_eq!(
2273 control_persist_mode(false, Some(Duration::from_secs(60))),
2274 ControlPersistMode::Closed
2275 );
2276 }
2277
2278 #[test]
2279 fn mux_without_idle_closes() {
2280 assert_eq!(control_persist_mode(true, None), ControlPersistMode::Closed);
2281 assert_eq!(
2282 control_persist_mode(true, Some(Duration::from_secs(0))),
2283 ControlPersistMode::Closed
2284 );
2285 }
2286
2287 #[test]
2288 fn mux_with_idle_keeps_warm() {
2289 assert_eq!(
2290 control_persist_mode(true, Some(Duration::from_secs(60))),
2291 ControlPersistMode::IdleFor(NonZeroUsize::new(60).unwrap())
2292 );
2293 }
2294}