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
38fn is_expected_health_check_output(stdout: &str) -> bool {
39 stdout
40 .trim()
41 .lines()
42 .last()
43 .is_some_and(is_health_check_sentinel)
44}
45
46fn is_health_check_sentinel(line: &str) -> bool {
47 matches!(line.trim(), "ok")
48}
49
50#[derive(Debug, Clone)]
52pub struct SshOptions {
53 pub connect_timeout: Duration,
55 pub command_timeout: Duration,
57 pub server_alive_interval: Option<Duration>,
61 pub control_persist_idle: Option<Duration>,
70 pub control_master: bool,
72 pub known_hosts: KnownHostsPolicy,
74}
75
76impl Default for SshOptions {
77 fn default() -> Self {
78 Self {
79 connect_timeout: DEFAULT_CONNECT_TIMEOUT,
80 command_timeout: DEFAULT_COMMAND_TIMEOUT,
81 server_alive_interval: None,
82 control_persist_idle: None,
83 control_master: false,
87 known_hosts: KnownHostsPolicy::Add,
88 }
89 }
90}
91
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94pub enum KnownHostsPolicy {
95 Strict,
97 Add,
99 AcceptAll,
101}
102
103#[cfg(test)]
104mod retry_tests {
105 use super::*;
106 use crate::test_guard;
107
108 #[test]
109 fn test_retryable_transport_error_text() {
110 let _guard = test_guard!();
111 assert!(is_retryable_transport_error_text(
112 "ssh: connect to host 1.2.3.4 port 22: Connection timed out"
113 ));
114 assert!(is_retryable_transport_error_text(
115 "kex_exchange_identification: Connection reset by peer"
116 ));
117 assert!(is_retryable_transport_error_text("Broken pipe"));
118 assert!(is_retryable_transport_error_text("Network is unreachable"));
119 }
120
121 #[test]
122 fn test_non_retryable_transport_error_text() {
123 let _guard = test_guard!();
124 assert!(!is_retryable_transport_error_text(
125 "Permission denied (publickey)."
126 ));
127 assert!(!is_retryable_transport_error_text(
128 "Host key verification failed."
129 ));
130 assert!(!is_retryable_transport_error_text(
131 "Could not resolve hostname worker.example.com: Name or service not known"
132 ));
133 assert!(!is_retryable_transport_error_text(
134 "Identity file /nope/id_rsa not accessible: No such file or directory"
135 ));
136 }
137}
138
139pub struct SshClient {
141 config: WorkerConfig,
143 options: SshOptions,
145 session: Option<Session>,
147}
148
149impl SshClient {
150 pub fn new(config: WorkerConfig, options: SshOptions) -> Self {
152 Self {
153 config,
154 options,
155 session: None,
156 }
157 }
158
159 pub fn worker_id(&self) -> &WorkerId {
161 &self.config.id
162 }
163
164 pub fn is_connected(&self) -> bool {
166 self.session.is_some()
167 }
168
169 fn is_configured_for(&self, config: &WorkerConfig) -> bool {
170 self.config.id == config.id
171 && self.config.host == config.host
172 && self.config.user == config.user
173 && self.config.identity_file == config.identity_file
174 && declared_os(&self.config.tags) == declared_os(&config.tags)
175 }
176
177 pub async fn connect(&mut self) -> Result<()> {
179 if self.session.is_some() {
180 debug!("Already connected to {}", self.config.id);
181 return Ok(());
182 }
183
184 let destination = format!("{}@{}", self.config.user, self.config.host);
185 debug!("Connecting to {} via SSH...", destination);
186
187 let session = match self
188 .connect_with_mode(&destination, self.options.control_master)
189 .await
190 {
191 Ok(session) => session,
192 Err(primary_error) if self.options.control_master => {
193 warn!(
194 "SSH ControlMaster connection to {} failed ({}). Retrying without ControlMaster.",
195 destination, primary_error
196 );
197 self.connect_with_mode(&destination, false)
198 .await
199 .with_context(|| {
200 format!(
201 "Failed to connect to {} after retrying without ControlMaster",
202 destination
203 )
204 })?
205 }
206 Err(primary_error) => {
207 return Err(primary_error)
208 .with_context(|| format!("Failed to connect to {}", destination));
209 }
210 };
211
212 debug!("Connected to {} ({})", self.config.id, self.config.host);
215 self.session = Some(session);
216 Ok(())
217 }
218
219 async fn connect_with_mode(&self, destination: &str, control_master: bool) -> Result<Session> {
220 let mut builder = SessionBuilder::default();
221 self.configure_builder(&mut builder, control_master);
222
223 builder.connect(destination).await.with_context(|| {
224 if control_master {
225 format!(
226 "Failed to connect to {} with ControlMaster enabled",
227 destination
228 )
229 } else {
230 format!(
231 "Failed to connect to {} with ControlMaster disabled",
232 destination
233 )
234 }
235 })
236 }
237
238 fn configure_builder(&self, builder: &mut SessionBuilder, control_master: bool) {
239 let known_hosts = match self.options.known_hosts {
240 KnownHostsPolicy::Strict => KnownHosts::Strict,
241 KnownHostsPolicy::Add => KnownHosts::Add,
242 KnownHostsPolicy::AcceptAll => KnownHosts::Accept,
243 };
244
245 builder
246 .known_hosts_check(known_hosts)
247 .connect_timeout(self.options.connect_timeout);
248
249 if let Some(interval) = self.options.server_alive_interval {
250 builder.server_alive_interval(interval);
251 }
252
253 let identity_path = shellexpand::tilde(&self.config.identity_file);
255 if Path::new(identity_path.as_ref()).exists() {
256 builder.keyfile(identity_path.as_ref());
257 }
258
259 match control_persist_mode(control_master, self.options.control_persist_idle) {
267 ControlPersistMode::IdleFor(nonzero) => {
268 builder.control_persist(ControlPersist::IdleFor(nonzero));
269 }
270 ControlPersistMode::TooLarge(secs) => {
271 warn!("control_persist_idle too large ({secs}s); closing after initial connection");
272 builder.control_persist(ControlPersist::ClosedAfterInitialConnection);
273 }
274 ControlPersistMode::Closed => {
275 builder.control_persist(ControlPersist::ClosedAfterInitialConnection);
276 }
277 }
278
279 if control_master {
281 let control_dir = {
290 let home_ssh = dirs::home_dir().map(|h| h.join(".ssh").join("rch"));
291
292 if let Some(ref dir) = home_ssh {
293 dir.clone()
294 } else if let Some(runtime_dir) = dirs::runtime_dir() {
295 runtime_dir.join("rch-ssh")
296 } else {
297 let username = whoami::username().unwrap_or_else(|_| "unknown".to_string());
298 std::env::temp_dir().join(format!("rch-ssh-{}", username))
299 }
300 };
301
302 if let Err(e) = std::fs::create_dir_all(&control_dir) {
303 warn!(
304 "Failed to create SSH control directory {:?}: {}",
305 control_dir, e
306 );
307 } else {
308 #[cfg(unix)]
311 {
312 use std::os::unix::fs::PermissionsExt;
313 if let Err(e) = std::fs::set_permissions(
314 &control_dir,
315 std::fs::Permissions::from_mode(0o700),
316 ) {
317 warn!(
318 "Failed to set permissions on SSH control directory {:?}: {}",
319 control_dir, e
320 );
321 }
322 }
323 }
324 builder.control_directory(&control_dir);
325 }
326 }
327
328 pub async fn disconnect(&mut self) -> Result<()> {
330 if let Some(session) = self.session.take() {
331 debug!("Disconnecting from {}", self.config.id);
332 session.close().await?;
333 debug!("Disconnected from {}", self.config.id);
335 }
336 Ok(())
337 }
338
339 pub async fn reconnect(&mut self) -> Result<()> {
348 if let Err(e) = self.disconnect().await {
349 debug!(
350 "Ignoring error closing stale session to {} before reconnect: {}",
351 self.config.id, e
352 );
353 }
354 self.connect().await
355 }
356
357 pub async fn execute(&self, command: &str) -> Result<CommandResult> {
360 self.execute_with_timeout(command, self.options.command_timeout)
361 .await
362 }
363
364 pub async fn execute_with_timeout(
370 &self,
371 command: &str,
372 command_timeout: Duration,
373 ) -> Result<CommandResult> {
374 if prefers_system_ssh(&self.config) {
382 return system_ssh_execute(
383 &self.config,
384 command,
385 command_timeout,
386 self.options
387 .server_alive_interval
388 .unwrap_or(Duration::from_secs(2)),
389 )
390 .await;
391 }
392
393 let session = self.session.as_ref().context("Not connected to worker")?;
394
395 let start = std::time::Instant::now();
396 debug!(
397 "Executing on {}: {}",
398 self.config.id,
399 crate::util::mask_sensitive_command(command)
400 );
401
402 let mut child = session
403 .command("sh")
404 .arg("-c")
405 .arg(command)
406 .stdout(Stdio::piped())
407 .stderr(Stdio::piped())
408 .spawn()
409 .await
410 .with_context(|| format!("Failed to spawn command on {}", self.config.id))?;
411
412 let execution_future = async {
413 let stdout_handle = child.stdout().take();
415 let stderr_handle = child.stderr().take();
416
417 let stdout_fut = async {
418 if let Some(out) = stdout_handle {
419 let reader = BufReader::new(out);
420 let mut take = reader.take(MAX_OUTPUT_SIZE);
421 let mut buf = String::new();
422 take.read_to_string(&mut buf).await?;
423 let mut reader = take.into_inner();
425 let mut sink = tokio::io::sink();
426 tokio::io::copy(&mut reader, &mut sink).await?;
427 if buf.len() >= MAX_OUTPUT_SIZE as usize {
428 buf.push_str("\n...[output truncated]...\n");
429 }
430 Ok::<String, anyhow::Error>(buf)
431 } else {
432 Ok(String::new())
433 }
434 };
435
436 let stderr_fut = async {
437 if let Some(err) = stderr_handle {
438 let reader = BufReader::new(err);
439 let mut take = reader.take(MAX_OUTPUT_SIZE);
440 let mut buf = String::new();
441 take.read_to_string(&mut buf).await?;
442 let mut reader = take.into_inner();
444 let mut sink = tokio::io::sink();
445 tokio::io::copy(&mut reader, &mut sink).await?;
446 if buf.len() >= MAX_OUTPUT_SIZE as usize {
447 buf.push_str("\n...[output truncated]...\n");
448 }
449 Ok::<String, anyhow::Error>(buf)
450 } else {
451 Ok(String::new())
452 }
453 };
454
455 let (stdout, stderr) = tokio::try_join!(stdout_fut, stderr_fut)?;
456
457 let status = child
458 .wait()
459 .await
460 .with_context(|| "Failed to wait for command completion")?;
461
462 Ok::<_, anyhow::Error>((status, stdout, stderr))
463 };
464
465 match tokio::time::timeout(command_timeout, execution_future).await {
466 Ok(result) => {
467 let (status, stdout, stderr) = result?;
468 let duration = start.elapsed();
469 let exit_code = status.code().unwrap_or(-1);
470
471 debug!(
472 "Command completed on {} (exit={}, duration={}ms)",
473 self.config.id,
474 exit_code,
475 duration.as_millis()
476 );
477
478 Ok(CommandResult {
479 exit_code,
480 stdout,
481 stderr,
482 duration_ms: duration.as_millis() as u64,
483 })
484 }
485 Err(_) => {
486 warn!(
493 "Command timed out on {} after {:?}",
494 self.config.id, command_timeout
495 );
496 anyhow::bail!("Command timed out after {:?}", command_timeout);
497 }
498 }
499 }
500
501 pub async fn execute_streaming<F, G>(
503 &self,
504 command: &str,
505 mut on_stdout: F,
506 mut on_stderr: G,
507 ) -> Result<CommandResult>
508 where
509 F: FnMut(&str),
510 G: FnMut(&str),
511 {
512 let session = self.session.as_ref().context("Not connected to worker")?;
513
514 let start = std::time::Instant::now();
515 debug!(
516 "Executing (streaming) on {}: {}",
517 self.config.id,
518 crate::util::mask_sensitive_command(command)
519 );
520
521 let mut child = session
522 .command("sh")
523 .arg("-c")
524 .arg(command)
525 .stdout(Stdio::piped())
526 .stderr(Stdio::piped())
527 .spawn()
528 .await
529 .with_context(|| format!("Failed to spawn command on {}", self.config.id))?;
530
531 let stdout = child.stdout().take();
532 let stderr = child.stderr().take();
533
534 let (tx, mut rx) = mpsc::channel(100);
537
538 if let Some(out) = stdout {
540 let tx = tx.clone();
541 tokio::spawn(async move {
542 let mut reader = BufReader::new(out);
543 let mut line = String::new();
544 loop {
545 line.clear();
546 match reader.read_line(&mut line).await {
547 Ok(0) => break, Ok(_) => {
549 if tx.send(StreamEvent::Stdout(line.clone())).await.is_err() {
550 break; }
552 }
553 Err(_) => break, }
555 }
556 });
557 }
558
559 if let Some(err) = stderr {
561 let tx = tx.clone();
562 tokio::spawn(async move {
563 let mut reader = BufReader::new(err);
564 let mut line = String::new();
565 loop {
566 line.clear();
567 match reader.read_line(&mut line).await {
568 Ok(0) => break, Ok(_) => {
570 if tx.send(StreamEvent::Stderr(line.clone())).await.is_err() {
571 break; }
573 }
574 Err(_) => break, }
576 }
577 });
578 }
579
580 drop(tx);
582
583 let mut stdout_acc = String::new();
584 let mut stderr_acc = String::new();
585
586 enum StreamEvent {
587 Stdout(String),
588 Stderr(String),
589 }
590
591 let streaming_future = async {
592 while let Some(event) = rx.recv().await {
594 match event {
595 StreamEvent::Stdout(line) => {
596 on_stdout(&line);
597 if stdout_acc.len() < MAX_OUTPUT_SIZE as usize {
598 stdout_acc.push_str(&line);
599 if stdout_acc.len() >= MAX_OUTPUT_SIZE as usize {
600 stdout_acc.push_str("\n...[output truncated]...\n");
601 }
602 }
603 }
604 StreamEvent::Stderr(line) => {
605 on_stderr(&line);
606 if stderr_acc.len() < MAX_OUTPUT_SIZE as usize {
607 stderr_acc.push_str(&line);
608 if stderr_acc.len() >= MAX_OUTPUT_SIZE as usize {
609 stderr_acc.push_str("\n...[output truncated]...\n");
610 }
611 }
612 }
613 }
614 }
615
616 let status = child.wait().await?;
617 Ok::<_, anyhow::Error>(status)
618 };
619
620 match tokio::time::timeout(self.options.command_timeout, streaming_future).await {
621 Ok(result) => {
622 let status = result?;
623 let duration = start.elapsed();
624 let exit_code = status.code().unwrap_or(-1);
625
626 Ok(CommandResult {
627 exit_code,
628 stdout: stdout_acc,
629 stderr: stderr_acc,
630 duration_ms: duration.as_millis() as u64,
631 })
632 }
633 Err(_) => {
634 warn!(
644 "Command (streaming) timed out on {} after {:?}, cleaning up",
645 self.config.id, self.options.command_timeout
646 );
647 anyhow::bail!("Command timed out after {:?}", self.options.command_timeout);
650 }
651 }
652 }
653
654 pub async fn health_check(&self) -> Result<bool> {
656 match self.execute(HEALTH_CHECK_COMMAND).await {
657 Ok(result) => Ok(result.success() && is_expected_health_check_output(&result.stdout)),
658 Err(e) => {
659 warn!("Health check failed for {}: {}", self.config.id, e);
660 Ok(false)
661 }
662 }
663 }
664}
665
666fn pooled_client_options(pool_options: &SshOptions, config: &WorkerConfig) -> SshOptions {
688 let mut options = pool_options.clone();
689 if declared_os(&config.tags).as_deref() == Some("windows") {
690 options.control_master = false;
691 }
692 options
693}
694
695pub(crate) fn prefers_system_ssh(config: &WorkerConfig) -> bool {
717 declared_os(&config.tags).as_deref() == Some("windows")
718}
719
720#[derive(Debug, Clone, PartialEq, Eq)]
722pub struct RemoteShellCommand {
723 pub command: String,
725 pub stdin_script: Option<String>,
728}
729
730#[must_use]
735pub fn remote_shell_command(config: &WorkerConfig, command: &str) -> RemoteShellCommand {
736 if prefers_system_ssh(config) {
737 RemoteShellCommand {
738 command: "sh -s".into(),
739 stdin_script: (command != "sh -s").then(|| {
740 format!(
741 "exec sh -c {} </dev/null\n",
742 shell_escape::escape(command.into())
743 )
744 }),
745 }
746 } else {
747 RemoteShellCommand {
748 command: command.into(),
749 stdin_script: None,
750 }
751 }
752}
753
754pub(crate) fn system_ssh_argv(
778 config: &WorkerConfig,
779 command: &RemoteShellCommand,
780 server_alive_interval: Duration,
781) -> Vec<OsString> {
782 let identity_path = shellexpand::tilde(&config.identity_file);
783 let destination = format!("{}@{}", config.user, config.host);
784
785 let mut argv: Vec<OsString> = Vec::with_capacity(10);
786 argv.push(OsString::from("ssh"));
787 argv.push(OsString::from("-i"));
788 argv.push(OsString::from(identity_path.as_ref()));
789 argv.push(OsString::from("-o"));
790 argv.push(OsString::from("BatchMode=yes"));
791 argv.push(OsString::from("-o"));
792 argv.push(OsString::from("ConnectTimeout=8"));
793 argv.push(OsString::from("-o"));
794 argv.push(OsString::from("StrictHostKeyChecking=accept-new"));
795 if !server_alive_interval.is_zero() {
796 let secs = server_alive_interval.as_secs().max(1);
800 argv.push(OsString::from("-o"));
801 argv.push(OsString::from(format!("ServerAliveInterval={secs}")));
802 }
803 argv.push(OsString::from(destination));
804 argv.push(OsString::from(&command.command));
805 argv
806}
807
808pub async fn write_ssh_command_stdin(
815 mut stdin: tokio::process::ChildStdin,
816 input: &str,
817) -> std::io::Result<()> {
818 use tokio::io::AsyncWriteExt;
819
820 let result = async {
821 stdin.write_all(input.as_bytes()).await?;
822 stdin.shutdown().await
823 }
824 .await;
825 match result {
826 Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),
829 result => result,
830 }
831}
832
833pub(crate) async fn system_ssh_execute(
848 config: &WorkerConfig,
849 command: &str,
850 command_timeout: Duration,
851 server_alive_interval: Duration,
852) -> Result<CommandResult> {
853 use std::process::Stdio;
854 use tokio::io::AsyncReadExt;
855 use tokio::process::Command;
856
857 let prepared = remote_shell_command(config, command);
858 let argv = system_ssh_argv(config, &prepared, server_alive_interval);
859 let input = prepared.stdin_script;
860 let args = argv.into_iter().skip(1);
862
863 let start = std::time::Instant::now();
864 debug!(
865 "system-ssh executing on {}: {}",
866 config.id,
867 crate::util::mask_sensitive_command(command)
868 );
869
870 let mut child = Command::new("ssh")
871 .args(args)
872 .env_clear()
873 .stdin(if input.is_some() {
874 Stdio::piped()
875 } else {
876 Stdio::null()
877 })
878 .stdout(Stdio::piped())
879 .stderr(Stdio::piped())
880 .kill_on_drop(true)
881 .spawn()
882 .with_context(|| {
883 format!(
884 "failed to spawn system ssh for {} (Windows fallback)",
885 config.id
886 )
887 })?;
888
889 let execution = async {
890 let stdout = child.stdout.take();
891 let stderr = child.stderr.take();
892 let stdin = child.stdin.take();
893
894 let stdin_fut = async {
895 if let (Some(stdin), Some(input)) = (stdin, input) {
896 write_ssh_command_stdin(stdin, &input).await?;
897 }
898 Ok::<(), anyhow::Error>(())
899 };
900
901 let stdout_fut = async {
902 let mut buf = String::new();
903 if let Some(out) = stdout {
904 let mut handle = out.take(MAX_OUTPUT_SIZE);
905 handle.read_to_string(&mut buf).await?;
906 }
907 Ok::<String, anyhow::Error>(buf)
908 };
909 let stderr_fut = async {
910 let mut buf = String::new();
911 if let Some(err) = stderr {
912 let mut handle = err.take(MAX_OUTPUT_SIZE);
913 handle.read_to_string(&mut buf).await?;
914 }
915 Ok::<String, anyhow::Error>(buf)
916 };
917
918 let ((), stdout, stderr) = tokio::try_join!(stdin_fut, stdout_fut, stderr_fut)?;
919 let status = child
920 .wait()
921 .await
922 .with_context(|| "system ssh command failed to wait for completion")?;
923 Ok::<_, anyhow::Error>((status, stdout, stderr))
924 };
925
926 match tokio::time::timeout(command_timeout, execution).await {
927 Ok(result) => {
928 let (status, stdout, stderr) = result?;
929 let duration = start.elapsed();
930 let exit_code = status.code().unwrap_or(-1);
931 debug!(
932 "system-ssh command completed on {} (exit={}, duration={}ms)",
933 config.id,
934 exit_code,
935 duration.as_millis()
936 );
937 Ok(CommandResult {
938 exit_code,
939 stdout,
940 stderr,
941 duration_ms: duration.as_millis() as u64,
942 })
943 }
944 Err(_) => {
945 warn!(
950 "system-ssh command timed out on {} after {:?}",
951 config.id, command_timeout
952 );
953 anyhow::bail!("Command timed out after {:?}", command_timeout);
954 }
955 }
956}
957
958pub struct SshPool {
960 connections: Arc<RwLock<HashMap<WorkerId, Arc<RwLock<SshClient>>>>>,
962 options: SshOptions,
964}
965
966impl SshPool {
967 pub fn new(options: SshOptions) -> Self {
969 Self {
970 connections: Arc::new(RwLock::new(HashMap::new())),
971 options,
972 }
973 }
974
975 pub async fn get_or_connect(&self, config: &WorkerConfig) -> Result<Arc<RwLock<SshClient>>> {
987 let shared_client = self.get_or_create_client_entry(config).await;
988
989 let reusable = {
993 let guard = shared_client.read().await;
994 guard.is_connected() && guard.health_check().await.unwrap_or(false)
995 };
996 if reusable {
997 debug!("Reusing live connection to {}", config.id);
998 return Ok(shared_client);
999 }
1000
1001 let mut client_guard = shared_client.write().await;
1006 if !client_guard.is_connected() {
1007 client_guard.connect().await?;
1008 } else if !client_guard.health_check().await.unwrap_or(false) {
1009 warn!(
1010 "Pooled SSH connection to {} failed liveness probe; reconnecting",
1011 config.id
1012 );
1013 client_guard.reconnect().await?;
1014 }
1015 drop(client_guard);
1017
1018 Ok(shared_client)
1019 }
1020
1021 async fn get_or_create_client_entry(&self, config: &WorkerConfig) -> Arc<RwLock<SshClient>> {
1022 let worker_id = config.id.clone();
1023
1024 loop {
1025 let existing_client = {
1026 let connections = self.connections.read().await;
1027 connections.get(&worker_id).cloned()
1028 };
1029
1030 if let Some(client) = existing_client {
1031 let is_configured_for_worker = {
1032 let guard = client.read().await;
1033 guard.is_configured_for(config)
1034 };
1035 if is_configured_for_worker {
1036 return client;
1037 }
1038
1039 let replacement = Arc::new(RwLock::new(SshClient::new(
1040 config.clone(),
1041 pooled_client_options(&self.options, config),
1042 )));
1043 let replaced = {
1044 let mut connections = self.connections.write().await;
1045 if connections
1046 .get(&worker_id)
1047 .is_some_and(|current| Arc::ptr_eq(current, &client))
1048 {
1049 connections.insert(worker_id.clone(), replacement.clone());
1050 true
1051 } else {
1052 false
1053 }
1054 };
1055
1056 if replaced {
1057 debug!(
1058 "Replaced SSH connection entry for {} after endpoint config changed",
1059 worker_id
1060 );
1061 return replacement;
1062 }
1063
1064 continue;
1065 }
1066
1067 let new_client = Arc::new(RwLock::new(SshClient::new(
1068 config.clone(),
1069 pooled_client_options(&self.options, config),
1070 )));
1071 let inserted = {
1072 let mut connections = self.connections.write().await;
1073 if connections.contains_key(&worker_id) {
1074 false
1075 } else {
1076 connections.insert(worker_id.clone(), new_client.clone());
1077 true
1078 }
1079 };
1080
1081 if inserted {
1082 return new_client;
1083 }
1084 }
1085 }
1086
1087 pub async fn close(&self, worker_id: &WorkerId) -> Result<()> {
1089 let client = {
1090 let mut connections = self.connections.write().await;
1091 connections.remove(worker_id)
1092 };
1093
1094 if let Some(client) = client {
1095 let mut client = client.write().await;
1096 client.disconnect().await?;
1097 }
1098
1099 Ok(())
1100 }
1101
1102 pub async fn close_all(&self) -> Result<()> {
1104 let clients: Vec<_> = {
1105 let mut connections = self.connections.write().await;
1106 connections.drain().map(|(_, v)| v).collect()
1107 };
1108
1109 for client in clients {
1110 let mut client = client.write().await;
1111 if let Err(e) = client.disconnect().await {
1112 error!("Error closing connection: {}", e);
1113 }
1114 }
1115
1116 Ok(())
1117 }
1118
1119 pub async fn active_connections(&self) -> usize {
1121 self.connections.read().await.len()
1122 }
1123
1124 pub async fn run_with_timeout(
1144 &self,
1145 config: &WorkerConfig,
1146 command: &str,
1147 command_timeout: Duration,
1148 ) -> Result<CommandResult> {
1149 let client = self.get_or_connect(config).await?;
1150 let guard = client.read().await;
1154 guard.execute_with_timeout(command, command_timeout).await
1155 }
1156}
1157
1158impl Default for SshPool {
1159 fn default() -> Self {
1160 Self::new(SshOptions::default())
1161 }
1162}
1163
1164#[cfg(test)]
1165mod tests {
1166 use super::*;
1167 use crate::test_guard;
1168
1169 #[test]
1170 fn test_command_result_success() {
1171 let _guard = test_guard!();
1172 let result = CommandResult {
1173 exit_code: 0,
1174 stdout: "output".to_string(),
1175 stderr: String::new(),
1176 duration_ms: 100,
1177 };
1178 assert!(result.success());
1179
1180 let failed = CommandResult {
1181 exit_code: 1,
1182 stdout: String::new(),
1183 stderr: "error".to_string(),
1184 duration_ms: 50,
1185 };
1186 assert!(!failed.success());
1187 }
1188
1189 #[test]
1190 fn test_ssh_options_default() {
1191 let _guard = test_guard!();
1192 let options = SshOptions::default();
1193 assert_eq!(options.connect_timeout, Duration::from_secs(10));
1194 assert_eq!(options.command_timeout, Duration::from_secs(300));
1195 assert!(options.server_alive_interval.is_none());
1196 assert!(options.control_persist_idle.is_none());
1197 assert!(!options.control_master);
1198 }
1199
1200 #[test]
1201 fn test_ssh_client_creation() {
1202 let _guard = test_guard!();
1203 let config = WorkerConfig {
1204 id: WorkerId::new("test-worker"),
1205 host: "192.168.1.100".to_string(),
1206 user: "ubuntu".to_string(),
1207 identity_file: "~/.ssh/id_rsa".to_string(),
1208 total_slots: 8,
1209 priority: 100,
1210 tags: vec!["rust".to_string()],
1211 tools: Vec::new(),
1212 };
1213
1214 let client = SshClient::new(config.clone(), SshOptions::default());
1215 assert_eq!(client.worker_id().as_str(), "test-worker");
1216 assert!(!client.is_connected());
1217 }
1218
1219 #[test]
1220 fn test_expected_health_check_output_accepts_sentinel_as_last_line() {
1221 let _guard = test_guard!();
1222
1223 assert!(is_expected_health_check_output("ok\n"));
1224 assert!(is_expected_health_check_output("login banner\nok\n"));
1225 assert!(!is_expected_health_check_output(""));
1226 assert!(!is_expected_health_check_output("not ok\n"));
1227 assert!(!is_expected_health_check_output("ok\npost-command noise\n"));
1228 }
1229
1230 fn worker_config(id: &str, host: &str, user: &str, identity_file: &str) -> WorkerConfig {
1231 WorkerConfig {
1232 id: WorkerId::new(id),
1233 host: host.to_string(),
1234 user: user.to_string(),
1235 identity_file: identity_file.to_string(),
1236 total_slots: 8,
1237 priority: 100,
1238 tags: vec!["rust".to_string()],
1239 tools: Vec::new(),
1240 }
1241 }
1242
1243 #[test]
1244 fn test_ssh_client_configured_for_ignores_scheduling_fields() {
1245 let _guard = test_guard!();
1246 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1247 let client = SshClient::new(config.clone(), SshOptions::default());
1248
1249 let mut scheduling_only_change = config;
1250 scheduling_only_change.total_slots = 16;
1251 scheduling_only_change.priority = 250;
1252 scheduling_only_change.tags = vec!["rust".to_string(), "gpu".to_string()];
1253
1254 assert!(client.is_configured_for(&scheduling_only_change));
1255 }
1256
1257 #[test]
1258 fn test_ssh_client_configured_for_detects_endpoint_changes() {
1259 let _guard = test_guard!();
1260 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1261 let client = SshClient::new(config, SshOptions::default());
1262
1263 assert!(!client.is_configured_for(&worker_config(
1264 "worker-a",
1265 "192.168.1.101",
1266 "ubuntu",
1267 "~/.ssh/id_rsa",
1268 )));
1269 assert!(!client.is_configured_for(&worker_config(
1270 "worker-a",
1271 "192.168.1.100",
1272 "admin",
1273 "~/.ssh/id_rsa",
1274 )));
1275 assert!(!client.is_configured_for(&worker_config(
1276 "worker-a",
1277 "192.168.1.100",
1278 "ubuntu",
1279 "~/.ssh/other_key",
1280 )));
1281 }
1282
1283 #[tokio::test]
1284 async fn test_ssh_pool_reuses_matching_disconnected_entry() {
1285 let _guard = test_guard!();
1286 let pool = SshPool::default();
1287 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1288
1289 let first = pool.get_or_create_client_entry(&config).await;
1290 let second = pool.get_or_create_client_entry(&config).await;
1291
1292 assert!(Arc::ptr_eq(&first, &second));
1293 assert_eq!(pool.active_connections().await, 1);
1294 }
1295
1296 #[tokio::test]
1297 async fn test_ssh_pool_reuses_single_entry_and_close_all_drops_to_zero() {
1298 let _guard = test_guard!();
1304 let options = SshOptions {
1305 control_master: true,
1306 control_persist_idle: Some(Duration::from_secs(60)),
1307 ..Default::default()
1308 };
1309 let pool = SshPool::new(options);
1310 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1311
1312 for _ in 0..5 {
1313 let _entry = pool.get_or_create_client_entry(&config).await;
1314 }
1315 assert_eq!(
1316 pool.active_connections().await,
1317 1,
1318 "repeated borrows of one worker must reuse a single pool entry"
1319 );
1320
1321 pool.close_all().await.expect("close_all should succeed");
1322 assert_eq!(
1323 pool.active_connections().await,
1324 0,
1325 "close_all must drop all pooled connections"
1326 );
1327 }
1328
1329 #[test]
1330 fn test_pool_options_map_to_bounded_idle_persist() {
1331 let _guard = test_guard!();
1335 assert_eq!(
1336 control_persist_mode(true, Some(Duration::from_secs(60))),
1337 ControlPersistMode::IdleFor(NonZeroUsize::new(60).unwrap()),
1338 "pool mux options must keep a warm master for a bounded idle window"
1339 );
1340 }
1341
1342 fn windows_worker_config(id: &str) -> WorkerConfig {
1347 let mut config = worker_config(id, "100.68.2.11", "jeffr", "~/.ssh/surfacebookje_key");
1348 config.tags = vec!["rust".to_string(), crate::types::os_tag("windows")];
1349 config
1350 }
1351
1352 #[test]
1353 fn test_pooled_options_disable_control_master_for_windows_workers() {
1354 let _guard = test_guard!();
1360 let pool_mux = SshOptions {
1361 control_master: true,
1362 control_persist_idle: Some(Duration::from_secs(60)),
1363 ..Default::default()
1364 };
1365 let config = windows_worker_config("wsurf");
1366
1367 let effective = pooled_client_options(&pool_mux, &config);
1368 assert!(
1369 !effective.control_master,
1370 "a Windows worker must never get a pooled ControlMaster"
1371 );
1372 assert_eq!(effective.connect_timeout, pool_mux.connect_timeout);
1376 assert_eq!(effective.command_timeout, pool_mux.command_timeout);
1377 assert_eq!(
1378 effective.server_alive_interval,
1379 pool_mux.server_alive_interval
1380 );
1381 assert_eq!(
1382 effective.control_persist_idle,
1383 pool_mux.control_persist_idle
1384 );
1385 assert_eq!(effective.known_hosts, pool_mux.known_hosts);
1386
1387 let pool_plain = SshOptions {
1389 control_master: false,
1390 ..pool_mux
1391 };
1392 assert!(!pooled_client_options(&pool_plain, &config).control_master);
1393
1394 let mut mixed_case = windows_worker_config("wsurf-2");
1397 mixed_case.tags = vec![crate::types::os_tag("Windows")];
1398 assert!(!pooled_client_options(&pool_mux, &mixed_case).control_master);
1399 }
1400
1401 #[test]
1402 fn test_pooled_options_keep_pool_options_verbatim_for_non_windows_workers() {
1403 let _guard = test_guard!();
1407 let pool_mux = SshOptions {
1408 control_master: true,
1409 control_persist_idle: Some(Duration::from_secs(60)),
1410 ..Default::default()
1411 };
1412 let linux = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1413 let unlabelled = worker_config("contabo-b", "1.2.3.5", "root", "~/.ssh/id_rsa");
1414
1415 for config in [&linux, &unlabelled] {
1416 let effective = pooled_client_options(&pool_mux, config);
1417 assert!(
1418 effective.control_master,
1419 "non-Windows workers keep the pool's mux setting"
1420 );
1421 assert_eq!(
1422 effective.control_persist_idle,
1423 pool_mux.control_persist_idle
1424 );
1425 }
1426
1427 let pool_plain = SshOptions {
1428 control_master: false,
1429 ..pool_mux
1430 };
1431 for config in [&linux, &unlabelled] {
1432 assert!(!pooled_client_options(&pool_plain, config).control_master);
1433 }
1434 }
1435
1436 #[tokio::test]
1437 async fn test_pool_entry_for_windows_worker_is_never_mux() {
1438 let _guard = test_guard!();
1444 let pool = SshPool::new(SshOptions {
1445 control_master: true,
1446 control_persist_idle: Some(Duration::from_secs(60)),
1447 ..Default::default()
1448 });
1449 let config = windows_worker_config("wsurf");
1450
1451 let entry = pool.get_or_create_client_entry(&config).await;
1452 let guard = entry.read().await;
1453 assert!(
1454 !guard.options.control_master,
1455 "pooled entry for a Windows worker must be non-mux"
1456 );
1457 assert_eq!(
1460 guard.options.control_persist_idle,
1461 Some(Duration::from_secs(60))
1462 );
1463 }
1464
1465 #[tokio::test]
1466 async fn test_pool_entry_for_linux_worker_keeps_mux() {
1467 let _guard = test_guard!();
1468 let pool = SshPool::new(SshOptions {
1469 control_master: true,
1470 control_persist_idle: Some(Duration::from_secs(60)),
1471 ..Default::default()
1472 });
1473 let config = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1474
1475 let entry = pool.get_or_create_client_entry(&config).await;
1476 let guard = entry.read().await;
1477 assert!(
1478 guard.options.control_master,
1479 "pooled entry for a Linux worker must keep the pool's mux setting"
1480 );
1481 }
1482
1483 #[tokio::test]
1484 async fn test_ssh_pool_replaces_stale_entry_when_endpoint_changes() {
1485 let _guard = test_guard!();
1486 let pool = SshPool::default();
1487 let old_config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1488 let new_config = worker_config("worker-a", "192.168.1.101", "admin", "~/.ssh/new_key");
1489
1490 let stale = pool.get_or_create_client_entry(&old_config).await;
1491 let replacement = pool.get_or_create_client_entry(&new_config).await;
1492
1493 assert!(!Arc::ptr_eq(&stale, &replacement));
1494 assert_eq!(pool.active_connections().await, 1);
1495
1496 let replacement_guard = replacement.read().await;
1497 assert!(replacement_guard.is_configured_for(&new_config));
1498 }
1499
1500 #[tokio::test]
1501 async fn test_ssh_pool_os_transition_replaces_transport_but_case_and_scheduling_reuse() {
1502 let _guard = test_guard!();
1503 let pool = SshPool::new(SshOptions {
1504 control_master: true,
1505 ..Default::default()
1506 });
1507 let mut config = worker_config("worker-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1508 let unlabelled = pool.get_or_create_client_entry(&config).await;
1509 config.tags.push(crate::types::os_tag("windows"));
1510 let windows = pool.get_or_create_client_entry(&config).await;
1511 assert!(!Arc::ptr_eq(&unlabelled, &windows));
1512 {
1513 let client = windows.read().await;
1514 assert!(prefers_system_ssh(&client.config));
1515 assert!(!client.options.control_master);
1516 }
1517
1518 config.tags = vec!["gpu".into(), "os: WiNdOwS ".into()];
1519 config.total_slots += 1;
1520 config.priority += 1;
1521 let reused = pool.get_or_create_client_entry(&config).await;
1522 assert!(Arc::ptr_eq(&windows, &reused));
1523
1524 config.tags = vec![crate::types::os_tag("linux")];
1525 let linux = pool.get_or_create_client_entry(&config).await;
1526 assert!(!Arc::ptr_eq(&windows, &linux));
1527 {
1528 let client = linux.read().await;
1529 assert!(!prefers_system_ssh(&client.config));
1530 assert!(client.options.control_master);
1531 }
1532 assert!(prefers_system_ssh(&windows.read().await.config));
1535 assert!(unlabelled.read().await.options.control_master);
1536 assert_eq!(pool.active_connections().await, 1);
1537
1538 config.tags.clear();
1539 let removed = pool.get_or_create_client_entry(&config).await;
1540 assert!(!Arc::ptr_eq(&linux, &removed));
1541 assert!(removed.read().await.options.control_master);
1542 assert_eq!(pool.active_connections().await, 1);
1543 }
1544
1545 #[tokio::test]
1546 async fn test_health_check_reports_not_alive_without_session() {
1547 let _guard = test_guard!();
1553 let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1554 let client = SshClient::new(config, SshOptions::default());
1555 assert!(!client.is_connected());
1556 let alive = client
1557 .health_check()
1558 .await
1559 .expect("health_check maps execution errors to Ok(false), never Err");
1560 assert!(!alive, "a client with no session must not report as alive");
1561 }
1562
1563 #[test]
1564 fn test_build_env_prefix_quotes_and_rejects() {
1565 let _guard = test_guard!();
1566 let mut env = HashMap::new();
1567 env.insert("RUSTFLAGS".to_string(), "-C target-cpu=native".to_string());
1568 env.insert("QUOTED".to_string(), "a'b".to_string());
1569 env.insert("BADVAL".to_string(), "line1\nline2".to_string());
1570
1571 let allowlist = vec![
1572 "RUSTFLAGS".to_string(),
1573 "QUOTED".to_string(),
1574 "MISSING".to_string(),
1575 "BADVAL".to_string(),
1576 "BAD=KEY".to_string(),
1577 ];
1578
1579 let prefix = build_env_prefix(&allowlist, |key| env.get(key).cloned());
1580
1581 assert!(prefix.prefix.contains("RUSTFLAGS='-C target-cpu=native'"));
1582 assert!(prefix.prefix.contains("QUOTED='a'\\''b'"));
1584 assert!(!prefix.prefix.contains("MISSING="));
1585 assert!(!prefix.prefix.contains("BADVAL="));
1586 assert!(prefix.rejected.contains(&"BADVAL".to_string()));
1587 assert!(prefix.rejected.contains(&"BAD=KEY".to_string()));
1588 assert_eq!(
1589 prefix.applied,
1590 vec!["RUSTFLAGS".to_string(), "QUOTED".to_string()]
1591 );
1592 }
1593
1594 mod proptest_ssh_escaping {
1599 use super::*;
1600 use proptest::prelude::*;
1601 use std::collections::HashMap;
1602
1603 proptest! {
1604 #![proptest_config(ProptestConfig::with_cases(1000))]
1605
1606 #[test]
1608 fn test_is_valid_env_key_no_panic(s in ".*") {
1609 let _guard = test_guard!();
1610 let _ = is_valid_env_key(&s);
1611 }
1612
1613 #[test]
1615 fn test_is_valid_env_key_accepts_valid(
1616 first in "[a-zA-Z_]",
1617 rest in "[a-zA-Z0-9_]{0,50}"
1618 ) {
1619 let _guard = test_guard!();
1620 let key = format!("{}{}", first, rest);
1621 prop_assert!(is_valid_env_key(&key), "Should accept valid key: {}", key);
1622 }
1623
1624 #[test]
1626 fn test_is_valid_env_key_rejects_digit_start(
1627 digit in "[0-9]",
1628 rest in "[a-zA-Z0-9_]{0,20}"
1629 ) {
1630 let _guard = test_guard!();
1631 let key = format!("{}{}", digit, rest);
1632 prop_assert!(!is_valid_env_key(&key), "Should reject digit-start key: {}", key);
1633 }
1634
1635 #[test]
1637 fn test_shell_escape_value_no_panic(s in ".*") {
1638 let _guard = test_guard!();
1639 let _ = shell_escape_value(&s);
1640 }
1641
1642 #[test]
1644 fn test_shell_escape_value_rejects_unsafe(
1645 prefix in "[a-zA-Z0-9 ]{0,10}",
1646 bad_char in "[\n\r\0]",
1647 suffix in "[a-zA-Z0-9 ]{0,10}"
1648 ) {
1649 let _guard = test_guard!();
1650 let value = format!("{}{}{}", prefix, bad_char, suffix);
1651 prop_assert!(shell_escape_value(&value).is_none(),
1652 "Should reject value with unsafe char: {:?}", value);
1653 }
1654
1655 #[test]
1657 fn test_shell_escape_value_accepts_safe(s in "[a-zA-Z0-9 !@#$%^&*()_+=\\-\\[\\]{}|;:,./<>?]{0,100}") {
1658 let result = shell_escape_value(&s);
1660 prop_assert!(result.is_some(), "Should accept safe value: {:?}", s);
1661
1662 let escaped = match result {
1665 Some(escaped) => escaped,
1666 None => {
1667 prop_assert!(false, "Should accept safe value: {:?}", s);
1668 String::new()
1669 }
1670 };
1671 if s.chars().any(|c| !c.is_ascii_alphanumeric() && c != '_') {
1672 prop_assert!(escaped.starts_with('\'') || escaped.contains('\''),
1674 "Value with special chars should be quoted: {:?} -> {:?}", s, escaped);
1675 }
1676 }
1677
1678 #[test]
1680 fn test_shell_escape_value_escapes_quotes(
1681 prefix in "[a-zA-Z0-9]{0,10}",
1682 suffix in "[a-zA-Z0-9]{0,10}"
1683 ) {
1684 let _guard = test_guard!();
1685 let value = format!("{}'{}", prefix, suffix);
1686 let result = shell_escape_value(&value);
1687 prop_assert!(result.is_some());
1688
1689 let escaped = match result {
1690 Some(escaped) => escaped,
1691 None => {
1692 prop_assert!(false, "Should escape single quote: {}", value);
1693 String::new()
1694 }
1695 };
1696 prop_assert!(escaped.contains("'\\''"),
1698 "Should escape single quote: {} -> {}", value, escaped);
1699 }
1700
1701 #[test]
1703 fn test_build_env_prefix_no_panic(
1704 keys in prop::collection::vec("[a-zA-Z_][a-zA-Z0-9_]{0,10}", 0..10),
1705 values in prop::collection::vec(".*", 0..10)
1706 ) {
1707 let mut env = HashMap::new();
1708 for (i, key) in keys.iter().enumerate() {
1709 if let Some(val) = values.get(i) {
1710 env.insert(key.clone(), val.clone());
1711 }
1712 }
1713
1714 let allowlist: Vec<String> = keys;
1715 let _ = build_env_prefix(&allowlist, |k| env.get(k).cloned());
1716 }
1717
1718 #[test]
1720 fn test_build_env_prefix_rejects_invalid_keys(
1721 invalid_key in "[0-9][a-zA-Z0-9_]{0,10}" ) {
1724 let _guard = test_guard!();
1725 let mut env = HashMap::new();
1726 env.insert(invalid_key.clone(), "value".to_string());
1727
1728 let allowlist = vec![invalid_key.clone()];
1729 let prefix = build_env_prefix(&allowlist, |k| env.get(k).cloned());
1730
1731 prop_assert!(!is_valid_env_key(&invalid_key),
1733 "Key should be invalid: {}", invalid_key);
1734 prop_assert!(prefix.rejected.contains(&invalid_key),
1735 "Should reject invalid key: {}", invalid_key);
1736 prop_assert!(prefix.prefix.is_empty());
1737 }
1738
1739 #[test]
1741 fn test_build_env_prefix_missing_values(
1742 keys in prop::collection::vec("[A-Z_][A-Z0-9_]{0,10}", 1..5)
1743 ) {
1744 let env: HashMap<String, String> = HashMap::new();
1746 let prefix = build_env_prefix(&keys, |k| env.get(k).cloned());
1747
1748 prop_assert!(prefix.prefix.is_empty(), "Should be empty when no values");
1750 prop_assert!(prefix.applied.is_empty());
1751 prop_assert!(prefix.rejected.is_empty());
1753 }
1754 }
1755
1756 #[test]
1758 fn test_shell_escape_edge_cases() {
1759 let _guard = test_guard!();
1760 let result = shell_escape_value("");
1762 assert_eq!(result, Some("''".to_string()));
1763
1764 let result = shell_escape_value("'");
1766 assert_eq!(result, Some("''\\'''".to_string()));
1767
1768 let result = shell_escape_value("'''");
1770 assert_eq!(
1772 result
1773 .as_deref()
1774 .map(|escaped| escaped.matches("'\\''").count()),
1775 Some(3)
1776 );
1777
1778 let result = shell_escape_value("日本語");
1780 assert!(result.is_some());
1781
1782 let result = shell_escape_value("🔥🚀");
1784 assert!(result.is_some());
1785
1786 let result = shell_escape_value("it's a \"test\" with $vars");
1788 assert!(result.is_some());
1789 }
1790
1791 #[test]
1792 fn test_is_valid_env_key_edge_cases() {
1793 let _guard = test_guard!();
1794 assert!(!is_valid_env_key(""));
1796
1797 assert!(is_valid_env_key("_"));
1799
1800 assert!(is_valid_env_key("A"));
1802
1803 assert!(is_valid_env_key("PATH"));
1805 assert!(is_valid_env_key("HOME"));
1806 assert!(is_valid_env_key("RUSTFLAGS"));
1807 assert!(is_valid_env_key("CC"));
1808 assert!(is_valid_env_key("_PRIVATE"));
1809 assert!(is_valid_env_key("MY_VAR_123"));
1810
1811 assert!(!is_valid_env_key("1VAR"));
1813 assert!(!is_valid_env_key("123"));
1814
1815 assert!(!is_valid_env_key("MY-VAR"));
1817 assert!(!is_valid_env_key("MY.VAR"));
1818 assert!(!is_valid_env_key("MY VAR"));
1819 assert!(!is_valid_env_key("MY=VAR"));
1820
1821 assert!(!is_valid_env_key("日本語"));
1823 assert!(!is_valid_env_key("VAR🔥"));
1824 }
1825
1826 #[test]
1827 fn test_build_env_prefix_integration() {
1828 let _guard = test_guard!();
1829 let mut env = HashMap::new();
1831 env.insert("VALID".to_string(), "simple".to_string());
1832 env.insert("WITH_QUOTE".to_string(), "it's here".to_string());
1833 env.insert("NEWLINE".to_string(), "line1\nline2".to_string());
1834 env.insert("UNICODE".to_string(), "日本語".to_string());
1835 env.insert("EMPTY".to_string(), String::new());
1836 env.insert("123INVALID".to_string(), "value".to_string());
1837
1838 let allowlist = vec![
1839 "VALID".to_string(),
1840 "WITH_QUOTE".to_string(),
1841 "NEWLINE".to_string(),
1842 "UNICODE".to_string(),
1843 "EMPTY".to_string(),
1844 "123INVALID".to_string(),
1845 "MISSING".to_string(),
1846 ];
1847
1848 let prefix = build_env_prefix(&allowlist, |k| env.get(k).cloned());
1849
1850 assert!(prefix.applied.contains(&"VALID".to_string()));
1852 assert!(prefix.prefix.contains("VALID=simple"));
1854
1855 assert!(prefix.applied.contains(&"WITH_QUOTE".to_string()));
1857
1858 assert!(prefix.rejected.contains(&"NEWLINE".to_string()));
1860
1861 assert!(prefix.applied.contains(&"UNICODE".to_string()));
1863
1864 assert!(prefix.applied.contains(&"EMPTY".to_string()));
1866
1867 assert!(prefix.rejected.contains(&"123INVALID".to_string()));
1869
1870 assert!(!prefix.applied.contains(&"MISSING".to_string()));
1872 assert!(!prefix.rejected.contains(&"MISSING".to_string()));
1873 }
1874
1875 #[test]
1876 fn test_shell_escape_roundtrip_safety() {
1877 let _guard = test_guard!();
1878 let test_values = [
1880 "simple",
1881 "with spaces",
1882 "with\ttab",
1883 "special!@#$%^&*()",
1884 "quoted\"value",
1885 "path/to/file",
1886 "-flag",
1887 "--long-flag=value",
1888 "",
1889 ];
1890
1891 for value in &test_values {
1892 let escaped = shell_escape_value(value);
1893 assert!(escaped.is_some(), "Should escape: {:?}", value);
1894 }
1895 }
1896 }
1897
1898 fn windows_worker(id: &str) -> WorkerConfig {
1903 let mut config = worker_config(id, "100.68.2.11", "jeffr", "~/.ssh/surfacebookje_key");
1904 config.tags = vec!["rust".to_string(), crate::types::os_tag("windows")];
1905 config
1906 }
1907
1908 #[test]
1909 fn test_prefers_system_ssh_for_windows_tag() {
1910 let _guard = test_guard!();
1913 let cfg = windows_worker("wsurf");
1914 assert!(prefers_system_ssh(&cfg));
1915 }
1916
1917 #[test]
1918 fn test_prefers_system_ssh_case_normalized() {
1919 let _guard = test_guard!();
1922 let mut cfg = windows_worker("wsurf-mixed");
1923 cfg.tags = vec![crate::types::os_tag("Windows")];
1924 assert!(prefers_system_ssh(&cfg));
1925
1926 let mut upper = windows_worker("wsurf-upper");
1927 upper.tags = vec![crate::types::os_tag("WINDOWS")];
1928 assert!(prefers_system_ssh(&upper));
1929 }
1930
1931 #[test]
1932 fn test_prefers_system_ssh_false_for_linux() {
1933 let _guard = test_guard!();
1936 let mut linux = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1937 linux.tags = vec!["rust".to_string(), crate::types::os_tag("linux")];
1938 assert!(!prefers_system_ssh(&linux));
1939
1940 let unlabelled = worker_config("contabo-b", "1.2.3.5", "root", "~/.ssh/id_rsa");
1941 assert!(!prefers_system_ssh(&unlabelled));
1942 }
1943
1944 #[test]
1945 fn test_system_ssh_argv_basic_windows() {
1946 let _guard = test_guard!();
1951 let cfg = windows_worker("wsurf");
1952 let command = remote_shell_command(&cfg, "uname -a");
1953 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(2));
1954 let s: Vec<String> = argv
1955 .iter()
1956 .map(|o| o.to_string_lossy().into_owned())
1957 .collect();
1958
1959 assert_eq!(s[0], "ssh", "argv[0] is the program name");
1960 assert_eq!(s[1], "-i");
1961 assert!(
1964 s[2].ends_with(".ssh/surfacebookje_key"),
1965 "identity file is the tilde-expanded path: {}",
1966 s[2]
1967 );
1968 assert!(s.contains(&"-o".to_string()));
1971 assert!(s.contains(&"BatchMode=yes".to_string()));
1972 assert!(s.contains(&"ConnectTimeout=8".to_string()));
1973 assert!(s.contains(&"StrictHostKeyChecking=accept-new".to_string()));
1974 assert!(s.contains(&"ServerAliveInterval=2".to_string()));
1975 assert!(!s.iter().any(|arg| arg == "-p"));
1977 assert!(s.contains(&"jeffr@100.68.2.11".to_string()));
1979 assert_eq!(s[s.len() - 1], "sh -s", "Windows gets only a fixed reader");
1980 }
1981
1982 #[test]
1983 fn test_system_ssh_script_transport_keeps_posix_argv_and_exact_stdin_reader() {
1984 let _guard = test_guard!();
1985 let linux = worker_config("linux", "1.2.3.4", "root", "~/.ssh/id_rsa");
1986 let windows = windows_worker("windows");
1987 let script = "printf '%s\\n' '%PATH% !VAR! & | < > ^ $()'\nprintf done";
1988 let prepared = remote_shell_command(&linux, script);
1989 let argv = system_ssh_argv(&linux, &prepared, Duration::ZERO);
1990 assert_eq!(argv.last().unwrap(), script);
1991 assert!(prepared.stdin_script.is_none());
1992 for command in [script, "sh -s", "", "sh -s\n"] {
1993 let prepared = remote_shell_command(&windows, command);
1994 let argv = system_ssh_argv(&windows, &prepared, Duration::ZERO);
1995 assert_eq!(argv.last().unwrap(), "sh -s");
1996 assert_eq!(prepared.stdin_script.is_none(), command == "sh -s");
1997 }
1998 }
1999
2000 #[tokio::test]
2001 async fn test_system_ssh_stdin_roundtrip_matches_null_stdin_shell() {
2002 use std::process::Stdio;
2003
2004 let _guard = test_guard!();
2005 let config = windows_worker("windows");
2006 for script in [
2009 "",
2010 "printf '%s\\n' \"apostrophe's\" 'double\"quote' '%PATH% !VAR! ^ & | < >' '$() `x` \\'",
2011 "cat\nprintf 'AFTER\\n'\nread value || printf 'EMPTY\\n'",
2012 "cat <<'END'\nheredoc ' \" % ! $() \\\nEND\nprintf 'done\\n'",
2013 "printf '%s\\n' 'first\nsecond' # trailing comment",
2014 "printf 'before\\n'; printf 'diagnostic\\n' >&2; exit 37",
2015 "# comment ending with a backslash \\",
2016 ] {
2017 let expected = tokio::process::Command::new("sh")
2018 .args(["-c", script])
2019 .stdin(Stdio::null())
2020 .output()
2021 .await
2022 .unwrap();
2023 let mut child = tokio::process::Command::new("sh")
2024 .arg("-s")
2025 .stdin(Stdio::piped())
2026 .stdout(Stdio::piped())
2027 .stderr(Stdio::piped())
2028 .kill_on_drop(true)
2029 .spawn()
2030 .unwrap();
2031 let stdin = child.stdin.take().unwrap();
2032 let payload = remote_shell_command(&config, script).stdin_script.unwrap();
2033 let (written, observed) = tokio::time::timeout(Duration::from_secs(5), async {
2034 tokio::join!(
2035 write_ssh_command_stdin(stdin, &payload),
2036 child.wait_with_output()
2037 )
2038 })
2039 .await
2040 .unwrap();
2041 written.unwrap();
2042 let observed = observed.unwrap();
2043 assert_eq!(observed.status.code(), expected.status.code(), "{script:?}");
2044 assert_eq!(observed.stdout, expected.stdout, "{script:?}");
2045 assert_eq!(observed.stderr, expected.stderr, "{script:?}");
2046 }
2047 }
2048
2049 #[tokio::test]
2050 async fn test_system_ssh_split_stdin_never_becomes_command_input() {
2051 use std::process::Stdio;
2052 use tokio::io::AsyncWriteExt;
2053
2054 let _guard = test_guard!();
2055 let script = "cat\nprintf 'AFTER\\n'\nread value || printf 'EMPTY\\n'";
2056 let payload = remote_shell_command(&windows_worker("windows"), script)
2057 .stdin_script
2058 .unwrap();
2059 let split = payload.find("cat\n").unwrap() + "cat\n".len();
2060 let mut child = tokio::process::Command::new("sh")
2061 .arg("-s")
2062 .stdin(Stdio::piped())
2063 .stdout(Stdio::piped())
2064 .stderr(Stdio::piped())
2065 .kill_on_drop(true)
2066 .spawn()
2067 .unwrap();
2068 let mut stdin = child.stdin.take().unwrap();
2069 let writer = async {
2070 stdin.write_all(&payload.as_bytes()[..split]).await.unwrap();
2071 tokio::time::sleep(Duration::from_millis(100)).await;
2072 stdin.write_all(&payload.as_bytes()[split..]).await.unwrap();
2073 stdin.shutdown().await.unwrap();
2074 };
2075 let ((), output) = tokio::time::timeout(Duration::from_secs(5), async {
2076 tokio::join!(writer, child.wait_with_output())
2077 })
2078 .await
2079 .unwrap();
2080 let output = output.unwrap();
2081 assert!(output.status.success(), "{:?}", output.stderr);
2082 assert_eq!(output.stdout, b"AFTER\nEMPTY\n");
2083 assert!(output.stderr.is_empty());
2084 }
2085
2086 #[tokio::test]
2087 async fn test_system_ssh_closed_stdin_preserves_early_exit_diagnostic() {
2088 use std::process::Stdio;
2089
2090 let _guard = test_guard!();
2091 let mut child = tokio::process::Command::new("sh")
2092 .args(["-c", "printf 'remote failure\\n' >&2; exit 37"])
2093 .stdin(Stdio::piped())
2094 .stdout(Stdio::piped())
2095 .stderr(Stdio::piped())
2096 .kill_on_drop(true)
2097 .spawn()
2098 .unwrap();
2099 let stdin = child.stdin.take().unwrap();
2100 child.wait().await.unwrap();
2103 write_ssh_command_stdin(stdin, &"x".repeat(128 * 1024))
2104 .await
2105 .unwrap();
2106 let output = child.wait_with_output().await.unwrap();
2107 assert_eq!(output.status.code(), Some(37));
2108 assert_eq!(output.stderr, b"remote failure\n");
2109 assert!(output.stdout.is_empty());
2110 }
2111
2112 #[test]
2113 fn test_system_ssh_argv_zero_server_alive_omits_keepalive() {
2114 let _guard = test_guard!();
2118 let cfg = windows_worker("wsurf");
2119 let command = remote_shell_command(&cfg, "echo hi");
2120 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(0));
2121 let s: Vec<String> = argv
2122 .iter()
2123 .map(|o| o.to_string_lossy().into_owned())
2124 .collect();
2125 assert!(
2126 !s.iter().any(|a| a.starts_with("ServerAliveInterval=")),
2127 "no ServerAliveInterval flag when interval is 0; got {:?}",
2128 s
2129 );
2130 }
2131
2132 #[test]
2133 fn test_system_ssh_argv_long_keepalive_preserved() {
2134 let _guard = test_guard!();
2138 let cfg = windows_worker("wsurf");
2139 let command = remote_shell_command(&cfg, "true");
2140 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(60));
2141 let s: Vec<String> = argv
2142 .iter()
2143 .map(|o| o.to_string_lossy().into_owned())
2144 .collect();
2145 assert!(s.contains(&"ServerAliveInterval=60".to_string()));
2146 }
2147
2148 #[test]
2149 fn test_system_ssh_argv_uses_tilde_expanded_identity() {
2150 let _guard = test_guard!();
2155 let mut cfg = windows_worker("wsurf");
2156 cfg.identity_file = "~/.ssh/operator_key".to_string();
2157 let command = remote_shell_command(&cfg, "true");
2158 let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(2));
2159 let s: Vec<String> = argv
2160 .iter()
2161 .map(|o| o.to_string_lossy().into_owned())
2162 .collect();
2163 assert!(
2164 !s.contains(&"~/.ssh/operator_key".to_string()),
2165 "identity file must be tilde-expanded, got {:?}",
2166 s
2167 );
2168 assert!(
2169 s.contains(&"~/.ssh/operator_key".to_string())
2170 || s.iter().any(|a| a.ends_with(".ssh/operator_key")),
2171 "tilde-expanded path appears in argv: {:?}",
2172 s
2173 );
2174 }
2175
2176 #[test]
2177 fn test_dispatch_helper_matches_declared_os_key() {
2178 let _guard = test_guard!();
2183 let cfg = windows_worker("wsurf");
2184 assert_eq!(declared_os(&cfg.tags).as_deref(), Some("windows"));
2185 assert!(prefers_system_ssh(&cfg));
2186 }
2187}
2188
2189#[derive(Debug, Clone, PartialEq, Eq)]
2191enum ControlPersistMode {
2192 Closed,
2194 IdleFor(NonZeroUsize),
2196 TooLarge(u64),
2198}
2199
2200fn control_persist_mode(control_master: bool, idle: Option<Duration>) -> ControlPersistMode {
2205 match idle {
2206 Some(idle) if control_master && !idle.is_zero() => match usize::try_from(idle.as_secs()) {
2207 Ok(secs) => match NonZeroUsize::new(secs) {
2208 Some(nonzero) => ControlPersistMode::IdleFor(nonzero),
2209 None => ControlPersistMode::Closed,
2210 },
2211 Err(_) => ControlPersistMode::TooLarge(idle.as_secs()),
2212 },
2213 _ => ControlPersistMode::Closed,
2214 }
2215}
2216
2217#[cfg(test)]
2218mod control_persist_tests {
2219 use super::*;
2220 use std::time::Duration;
2221
2222 #[test]
2223 fn non_mux_sessions_never_persist_forever() {
2224 assert_eq!(
2226 control_persist_mode(false, None),
2227 ControlPersistMode::Closed
2228 );
2229 assert_eq!(
2230 control_persist_mode(false, Some(Duration::from_secs(60))),
2231 ControlPersistMode::Closed
2232 );
2233 }
2234
2235 #[test]
2236 fn mux_without_idle_closes() {
2237 assert_eq!(control_persist_mode(true, None), ControlPersistMode::Closed);
2238 assert_eq!(
2239 control_persist_mode(true, Some(Duration::from_secs(0))),
2240 ControlPersistMode::Closed
2241 );
2242 }
2243
2244 #[test]
2245 fn mux_with_idle_keeps_warm() {
2246 assert_eq!(
2247 control_persist_mode(true, Some(Duration::from_secs(60))),
2248 ControlPersistMode::IdleFor(NonZeroUsize::new(60).unwrap())
2249 );
2250 }
2251}