Skip to main content

rch_common/
ssh.rs

1//! SSH client utilities for remote command execution.
2//!
3//! Provides connection management, command execution, and pooling support
4//! for the remote compilation pipeline.
5//!
6//! This module is only available on Unix platforms (requires openssh crate).
7
8use 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
21// Re-export platform-independent utilities for backwards compatibility
22pub 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
27/// Default SSH connection timeout.
28const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
29
30/// Default command execution timeout.
31const DEFAULT_COMMAND_TIMEOUT: Duration = Duration::from_secs(300);
32
33/// Maximum size for command output (stdout/stderr) to prevent OOM (10MB).
34const MAX_OUTPUT_SIZE: u64 = 10 * 1024 * 1024;
35
36const HEALTH_CHECK_COMMAND: &str = "echo ok";
37
38/// Bound on the pool's reuse liveness probe (`echo ok` over an existing
39/// master). A live master answers in milliseconds; a hung one must not hold a
40/// caller for the client's multi-minute command timeout.
41const 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/// SSH connection options.
56#[derive(Debug, Clone)]
57pub struct SshOptions {
58    /// Connection timeout.
59    pub connect_timeout: Duration,
60    /// Command execution timeout.
61    pub command_timeout: Duration,
62    /// Server keepalive interval (`ssh -o ServerAliveInterval`).
63    ///
64    /// Defaults to `None` (OpenSSH default; keepalive disabled).
65    pub server_alive_interval: Option<Duration>,
66    /// How long the SSH ControlMaster should remain alive while idle.
67    ///
68    /// Only applies when `control_master` is true (connection reuse). `Some(n)`
69    /// with n > 0 keeps the master warm for n idle seconds (`ControlPersist=Ns`).
70    /// `Some(0s)`/`None`, or ANY non-mux per-call session, closes the master
71    /// after the initial connection (`ControlPersist=no`). The OpenSSH crate
72    /// `ControlPersist=yes` (forever) default is never used — it leaked a
73    /// master process per per-call SSH. See `control_persist_mode`.
74    pub control_persist_idle: Option<Duration>,
75    /// SSH control master mode for connection reuse.
76    pub control_master: bool,
77    /// Known hosts policy.
78    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            // Default to a plain SSH session. ControlMaster is an optimization
89            // and stale local control sockets can poison otherwise healthy
90            // connections. Callers that explicitly want mux reuse can opt in.
91            control_master: false,
92            known_hosts: KnownHostsPolicy::Add,
93        }
94    }
95}
96
97/// Known hosts policy for SSH connections.
98#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99pub enum KnownHostsPolicy {
100    /// Strictly verify known hosts (recommended for production).
101    Strict,
102    /// Add unknown hosts automatically (for development).
103    Add,
104    /// Accept all hosts without verification (INSECURE - testing only).
105    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
144/// SSH client for a single worker connection.
145pub struct SshClient {
146    /// Worker configuration.
147    config: WorkerConfig,
148    /// SSH options.
149    options: SshOptions,
150    /// Active SSH session (if connected).
151    session: Option<Session>,
152}
153
154impl SshClient {
155    /// Create a new SSH client for a worker.
156    pub fn new(config: WorkerConfig, options: SshOptions) -> Self {
157        Self {
158            config,
159            options,
160            session: None,
161        }
162    }
163
164    /// Get the worker ID.
165    pub fn worker_id(&self) -> &WorkerId {
166        &self.config.id
167    }
168
169    /// Check if connected to the worker.
170    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    /// Connect to the remote worker.
183    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, not info: the telemetry pool connects to every worker each poll
218        // cycle, so at info this floods the daemon log (8M+ lines / multi-GB).
219        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        // Add identity file if specified
259        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        // Always pin ControlPersist explicitly. The openssh crate ALWAYS
265        // spawns a control-master; leaving it unset defaults to
266        // `ControlPersist=yes` (Forever), which leaks a master process for
267        // every short-lived per-call SSH (telemetry/health/capabilities run
268        // with control_master=false). ClosedAfterInitialConnection lets
269        // non-mux masters self-terminate; only the explicit mux-reuse path
270        // with a configured idle keeps a warm master.
271        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        // Control-master socket directory only matters when reusing connections.
285        if control_master {
286            // Use a short control directory path to stay within the Unix domain
287            // socket path limit (104 bytes on macOS, 108 on Linux).  The openssh
288            // crate appends a `%C` hash (~32 chars) to form the socket filename,
289            // so the directory path itself must be short.
290            //
291            // On macOS `std::env::temp_dir()` returns a long path under
292            // /var/folders/…/T/ which, combined with the hash, exceeds 104 bytes.
293            // We therefore prefer `~/.ssh/rch` (short, stable, correct perms).
294            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                // Set restrictive permissions (0700) to prevent symlink attacks
314                // and unauthorized access to SSH control sockets
315                #[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    /// Disconnect from the worker.
334    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, not info: paired with the connect log above; floods at info.
339            debug!("Disconnected from {}", self.config.id);
340        }
341        Ok(())
342    }
343
344    /// Force a fresh connection, dropping any existing (possibly dead) session.
345    ///
346    /// [`connect`](Self::connect) is a no-op when a `Session` object already
347    /// exists, so a pooled client whose underlying SSH ControlMaster has died
348    /// cannot recover through it. `reconnect` tears the stale session down first
349    /// (closing it best-effort — a dead master may itself error on close, but
350    /// [`disconnect`](Self::disconnect) `take()`s the session before closing so
351    /// it is gone regardless) and then establishes a new one.
352    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    /// Execute a command on the remote worker using this client's configured
363    /// command timeout.
364    pub async fn execute(&self, command: &str) -> Result<CommandResult> {
365        self.execute_with_timeout(command, self.options.command_timeout)
366            .await
367    }
368
369    /// Execute a command on the remote worker with an explicit per-call timeout.
370    ///
371    /// Used by the connection pool ([`SshPool::run_with_timeout`]) so a single
372    /// pooled client can serve calls with different timeouts (a fast health
373    /// probe vs. a slower cleanup) without mutating the shared client options.
374    pub async fn execute_with_timeout(
375        &self,
376        command: &str,
377        command_timeout: Duration,
378    ) -> Result<CommandResult> {
379        // Windows fallback (bd-kzy2x): the openssh crate cannot execute
380        // commands on Windows OpenSSH (the slave never completes the
381        // command channel), so we dispatch through the system `ssh` binary
382        // for declared-OS Windows workers. The connect path above is
383        // untouched and still uses the openssh crate (it succeeds); only
384        // the execute step swaps. Linux / unlabelled workers keep the
385        // openssh-crate path verbatim.
386        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            // Read stdout and stderr concurrently to avoid deadlock if one pipe fills.
419            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                    // Drain the rest to prevent SIGPIPE or blocking
429                    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                    // Drain the rest to prevent SIGPIPE or blocking
448                    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                // Timeout occurred. The inner future (which owns `child`) is dropped,
492                // but dropping an openssh RemoteChild does NOT kill the remote process
493                // if a ControlMaster is active. The remote process will only terminate
494                // when the caller disconnects the SshClient (closing the Session/
495                // ControlMaster). Callers should ensure disconnect() is called after
496                // a timeout to avoid leaked remote processes.
497                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    /// Execute a command and stream output in real-time.
507    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        // Use a channel to aggregate stream events from reader tasks.
540        // This avoids cancellation safety issues with select! over AsyncBufReadExt::read_line.
541        let (tx, mut rx) = mpsc::channel(100);
542
543        // Spawn stdout reader
544        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, // EOF
553                        Ok(_) => {
554                            if tx.send(StreamEvent::Stdout(line.clone())).await.is_err() {
555                                break; // Receiver dropped
556                            }
557                        }
558                        Err(_) => break, // Read error
559                    }
560                }
561            });
562        }
563
564        // Spawn stderr reader
565        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, // EOF
574                        Ok(_) => {
575                            if tx.send(StreamEvent::Stderr(line.clone())).await.is_err() {
576                                break; // Receiver dropped
577                            }
578                        }
579                        Err(_) => break, // Read error
580                    }
581                }
582            });
583        }
584
585        // Drop original tx so rx closes when tasks finish
586        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            // Process events until channel closes (EOF from both streams)
598            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                // Timeout occurred - the spawned reader tasks will terminate when they
640                // try to send on rx (which is dropped when this scope exits).
641                // The child process is also dropped here, but openssh may not kill
642                // the remote process immediately. Log the situation for visibility.
643                //
644                // Note: The reader tasks are detached (tokio::spawn) so they continue
645                // briefly until they hit EOF or the send fails. This is acceptable
646                // because they're lightweight and will terminate quickly once the
647                // channel closes.
648                warn!(
649                    "Command (streaming) timed out on {} after {:?}, cleaning up",
650                    self.config.id, self.options.command_timeout
651                );
652                // rx is dropped here, which will cause senders to fail on next send
653                // child is dropped here, which signals termination to openssh
654                anyhow::bail!("Command timed out after {:?}", self.options.command_timeout);
655            }
656        }
657    }
658
659    /// Check if the worker is reachable via SSH.
660    pub async fn health_check(&self) -> Result<bool> {
661        self.health_check_with_timeout(self.options.command_timeout)
662            .await
663    }
664
665    /// [`health_check`](Self::health_check) bounded by `timeout` instead of
666    /// the client's command timeout (300s for pooled clients).
667    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
681/// Derive the per-worker [`SshOptions`] a pool hands to each pooled
682/// [`SshClient`].
683///
684/// POLICY — Windows workers never get ControlMaster multiplexing. Windows
685/// OpenSSH has no ControlMaster support: a pooled (multiplexed) session
686/// appears to connect, then every command over it hangs and dies at the
687/// command stage ("Failed to wait for command completion"), so each pooled
688/// probe fails deterministically and the daemon marks the worker permanently
689/// unreachable even though fresh one-shot SSH to the same host succeeds
690/// (2026-08-30 incident: worker `wsurf`, `os = "windows"` — healthy via
691/// `rch workers probe` fresh connections, "unreachable" via every pooled
692/// health check). Applying the override HERE — the single place pooled
693/// clients are constructed from the pool's global options — fixes every pool
694/// consumer (daemon shared build/telemetry pool, dedicated health pool) at
695/// once. [`SshClient::connect`]'s ControlMaster-failure retry cannot catch
696/// this class: the mux connect does not fail, the commands do.
697///
698/// `control_persist_idle` needs no separate neutralization: it is consulted
699/// only when `control_master` is true (see [`control_persist_mode`]), so it
700/// is inert once mux is disabled and is left verbatim. Non-Windows workers
701/// get the pool options verbatim in full.
702fn 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
710// System-ssh command-execution fallback (bd-kzy2x).
711//
712// The `openssh` crate 0.11.6 cannot execute commands on Windows OpenSSH
713// regardless of `control_master` / `control_persist` / keepalive settings:
714// the slave never completes the command channel and every command hangs at
715// the execute stage even though the connect succeeds. The CLI `ssh` binary
716// speaks the OpenSSH wire protocol directly and works fine. For workers that
717// declare `os = "windows"` we therefore spawn the system `ssh` binary at
718// command-execution time, mirroring the proven system-ssh pattern already
719// used by the fleet preflight path (`rch/src/fleet/ssh.rs::SshExecutor`) and
720// the CLI worker init / probe paths. The connect path stays the openssh
721// crate (it succeeds) — the fallback is purely an execute-time dispatch
722// inside `SshClient::execute_with_timeout`.
723//
724// Policy key (kept in lockstep with the existing pool-layer override
725// `pooled_client_options`): `declared_os(&config.tags) == Some("windows")`.
726// `os = "Windows"` is case-normalized by `declared_os` so both spellings
727// route through the same fallback.
728
729/// Should `SshClient::execute_with_timeout` dispatch through the system `ssh`
730/// binary for this worker instead of the `openssh` crate?
731pub(crate) fn prefers_system_ssh(config: &WorkerConfig) -> bool {
732    declared_os(&config.tags).as_deref() == Some("windows")
733}
734
735/// Remote command and optional stdin payload for a system SSH invocation.
736#[derive(Debug, Clone, PartialEq, Eq)]
737pub struct RemoteShellCommand {
738    /// Command argument passed to SSH after the destination.
739    pub command: String,
740    /// Script to write and close concurrently with draining stdout/stderr.
741    /// When absent, the caller should provide null stdin.
742    pub stdin_script: Option<String>,
743}
744
745/// Prepare a POSIX script without exposing its contents to Windows cmd.exe.
746/// Windows uses a fixed `sh -s` reader, which replaces itself with a command
747/// shell whose stdin is empty. An exact `sh -s` keeps its existing EOF behavior.
748/// POSIX workers receive the original command argument and null stdin.
749#[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
769/// Build the argv for the system-ssh fallback, mirroring the proven CLI
770/// system-ssh pattern (see `rch/src/fleet/ssh.rs::SshExecutor::build_ssh_args`
771/// and `rch/src/commands/workers_init.rs`). Pure / testable: no process is
772/// spawned. The first element is the program name `"ssh"` so the vector can
773/// be passed straight to `Command::new` callers that prefer the explicit
774/// first arg, but the standard `Command::new("ssh").args(...)` callers can
775/// skip it — see `system_ssh_execute` which drops it.
776///
777/// Argv layout:
778/// `ssh -i <identity_file> -o BatchMode=yes -o ConnectTimeout=8 -o
779///   StrictHostKeyChecking=accept-new -o ServerAliveInterval=<secs>
780///   <user>@<host> <command>`
781///
782/// `WorkerConfig` has no port field today; the SSH port is carried as part
783/// of the host (`host:port`) and SSH's own config handles non-default
784/// ports, so the argv has no `-p` flag. The 8s `ConnectTimeout` matches the
785/// fleet SshExecutor default (`DEFAULT_CONNECT_TIMEOUT_SECS = 10` would be
786/// the natural match too — kept at 8 to match the FINDINGS spec exactly).
787/// `server_alive_interval` is rounded up to 1s minimum when non-zero to
788/// avoid `ServerAliveInterval=0` (which OpenSSH treats as "disable
789/// keepalives" — fine — but matches the spec "0 -> omit" intent).
790/// Windows receives only the fixed POSIX reader `sh -s`; the original script
791/// travels through stdin so cmd.exe never interprets its contents.
792pub(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        // `0` means "no keepalive" in OpenSSH; otherwise emit the requested
812        // interval in whole seconds (sub-second values are rounded up to 1s
813        // because OpenSSH only accepts integer seconds for this option).
814        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
823/// Write and close an SSH script pipe within the caller's execution deadline.
824/// Run concurrently with output draining. A closed pipe is allowed so the
825/// caller can collect an early SSH failure's actual exit status and stderr.
826///
827/// # Errors
828/// Returns write or shutdown errors other than an expected broken pipe.
829pub 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        // An authentication failure or early remote exit can close stdin
842        // before the payload is written. Preserve its actual status/stderr.
843        Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),
844        result => result,
845    }
846}
847
848/// Execute a command on a Windows worker via the system `ssh` binary.
849///
850/// Mirrors the openssh-crate path's size cap (`MAX_OUTPUT_SIZE`), timeout
851/// semantics (tokio timeout that bails on expiry), and `CommandResult`
852/// shape. The local ssh process is set `kill_on_drop(true)` so a timeout
853/// cannot leak an `ssh` child — the openssh-crate timeout path likewise
854/// relies on the caller to drop the child (and warns about the leak risk);
855/// here we get a hard SIGKILL on drop instead. Output reads are
856/// concurrent over `tokio::io::BufReader` to avoid a single-stream back-
857/// pressure deadlock (same pattern as the openssh path). `env_clear()` is
858/// used so a hostile / noisy environment from the calling daemon cannot
859/// perturb the local `ssh` invocation; SSH agent forwarding and standard
860/// paths still resolve via `~/.ssh/config` because ssh reads them from the
861/// filesystem, not the env.
862pub(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    // Drop the program name — `Command::new("ssh")` is the standard form.
876    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            // `kill_on_drop(true)` ensures the local ssh process is
961            // SIGKILLed when `child` drops here, so no ssh child can leak
962            // even if the remote end is wedged. This only owns the local
963            // SSH child; remote process termination is not guaranteed.
964            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
973/// Connection pool for managing multiple SSH connections.
974pub struct SshPool {
975    /// Pool of active connections.
976    connections: Arc<RwLock<HashMap<WorkerId, Arc<RwLock<SshClient>>>>>,
977    /// Default SSH options.
978    options: SshOptions,
979}
980
981impl SshPool {
982    /// Create a new connection pool.
983    pub fn new(options: SshOptions) -> Self {
984        Self {
985            connections: Arc::new(RwLock::new(HashMap::new())),
986            options,
987        }
988    }
989
990    /// Get or create a connection to a worker.
991    ///
992    /// Validates liveness on borrow: a pooled `SshClient` reports
993    /// [`is_connected`](SshClient::is_connected) purely from the presence of a
994    /// `Session` object, which cannot detect that the underlying SSH
995    /// ControlMaster has died (idle `ControlPersist` expiry, a network blip, or
996    /// the remote `sshd` being restarted). Returning such a client hands the
997    /// caller a dead master, and its next `execute()` fails with
998    /// broken-pipe/connection-closed and no chance to recover. So before reusing
999    /// a connected client this probes it with a lightweight health check, and on
1000    /// failure reconnects once under the per-worker write lock.
1001    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    /// [`get_or_connect`](Self::get_or_connect) with the liveness probe bounded
1007    /// by `probe_timeout`.
1008    ///
1009    /// The probe used to run under the client's command timeout (300s for the
1010    /// daemon's pools), and a failed probe is followed by a reconnect and a
1011    /// second probe. A worker whose master was up but whose commands hung (fork
1012    /// exhaustion, D-state) therefore held a caller with a 10-15s budget (health
1013    /// probe, toolchain preflight during selection) for ~10 minutes, and the
1014    /// sequential health loop stalled the whole fleet's state behind it.
1015    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        // Fast path: reuse only a session that is both present AND verified live.
1023        // The probe needs `&self` only (health_check → execute), so it runs under
1024        // a shared read lock and never blocks other borrowers of this worker.
1025        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        // Slow path: (re)connect under the per-worker write lock. Re-evaluate the
1039        // state here because another task may have connected/reconnected while we
1040        // waited for the lock, and to distinguish "never connected" (connect)
1041        // from "session exists but the master is dead" (reconnect).
1042        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 write lock before returning
1057        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    /// Close a specific connection.
1129    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    /// Close all connections.
1144    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    /// Get the number of active connections.
1161    pub async fn active_connections(&self) -> usize {
1162        self.connections.read().await.len()
1163    }
1164
1165    /// Run a single command on a worker over a POOLED (warm, reused) connection,
1166    /// bounded by `command_timeout`.
1167    ///
1168    /// This is the entry point daemon subsystems use instead of the throwaway
1169    /// `SshClient::new().connect()...execute()...disconnect()` dance. It
1170    /// [`get_or_connect`](Self::get_or_connect)s a live master (validating
1171    /// liveness on borrow and reconnecting a dead one), runs exactly one command,
1172    /// and — crucially — does NOT disconnect afterwards, keeping the
1173    /// ControlMaster warm for the next call. That is the whole point: reuse one
1174    /// master per worker instead of spawning (and, under the old code, leaking) a
1175    /// fresh master for every telemetry/health/cleanup poll.
1176    ///
1177    /// The per-command timeout is applied via a temporary [`SshOptions`] override
1178    /// on the pooled client so it does not disturb the pool's shared default
1179    /// (e.g. a long build vs. a short health probe). The connect timeout and
1180    /// control-master/persist settings come from the pool's options — except
1181    /// that Windows workers are ALWAYS non-mux (see [`pooled_client_options`]):
1182    /// Windows OpenSSH cannot multiplex, so a pooled mux session hangs at the
1183    /// command stage.
1184    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        // Execute under a shared read lock: execute() needs only `&self`, so
1194        // concurrent callers for the same worker can multiplex over the one
1195        // master without serializing on a write lock.
1196        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        // Leak guard: repeated borrows of the SAME worker must reuse ONE pool
1342        // entry (one warm master, not one-per-call), and close_all() must empty
1343        // the pool. Uses the entry-creation path (get_or_create_client_entry) so
1344        // the test needs no live SSH host; get_or_connect layers only a liveness
1345        // probe + reconnect on top of this same entry map.
1346        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        // Leak guard: the pool's mux options must map to a BOUNDED
1375        // ControlPersist=IdleFor(60), never Closed/Forever. control_persist_mode
1376        // is the single decision point configure_builder uses.
1377        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    // ======================================================================
1386    // Windows workers never get pooled ControlMaster (bd-wgbx9)
1387    // ======================================================================
1388
1389    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        // Regression (bd-wgbx9): Windows OpenSSH has no ControlMaster — a
1398        // pooled mux session hangs at the command stage ("Failed to wait for
1399        // command completion") and the worker is falsely marked unreachable
1400        // while fresh one-shot SSH succeeds. Whatever the pool's global
1401        // options say, a Windows worker's pooled client must be non-mux.
1402        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        // The override must be surgical: everything else stays verbatim.
1416        // control_persist_idle is left as-is because it is inert once mux is
1417        // off (control_persist_mode only reads it when control_master is true).
1418        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        // A pool that already disabled mux stays disabled.
1431        let pool_plain = SshOptions {
1432            control_master: false,
1433            ..pool_mux
1434        };
1435        assert!(!pooled_client_options(&pool_plain, &config).control_master);
1436
1437        // declared_os normalizes case: `os = "Windows"` in workers.toml is the
1438        // same reserved tag and must disable mux too.
1439        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        // Linux/unlabelled workers keep the pool's options VERBATIM — the
1447        // override must not degrade warm-master reuse for the rest of the
1448        // fleet, under either a mux or a hypothetical non-mux pool.
1449        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        // End-to-end at the pool layer: the daemon's health pool and shared
1482        // build/telemetry pool both construct per-worker clients via
1483        // get_or_create_client_entry with pool-global mux options — the
1484        // Windows entry must come out non-mux. Uses the entry-creation path
1485        // (like the other pool tests) so no live SSH host is needed.
1486        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        // The inert companion setting is untouched, proving the override was
1501        // surgical rather than a blanket options reset.
1502        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        // Existing borrowers retain their original transport while the pool
1576        // exposes only the replacement to subsequent requests.
1577        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        // get_or_connect's liveness probe relies on health_check() returning a
1591        // non-true value (without erroring) for a client that has no live
1592        // session, so a dead/empty pooled entry is reconnected rather than
1593        // falsely handed back as "reused". execute() errors "Not connected",
1594        // which health_check() maps to Ok(false).
1595        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        // shell_escape uses '\'' style (end string, escaped quote, start string)
1626        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    // ==========================================================================
1638    // Proptest: SSH command escaping with special chars (bd-2elj)
1639    // ==========================================================================
1640
1641    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 1: is_valid_env_key never panics on arbitrary strings
1650            #[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 2: Valid env keys start with letter/_ and contain only alphanum/_
1657            #[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 3: Env keys starting with digit are invalid
1668            #[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 4: shell_escape_value never panics on arbitrary strings
1679            #[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 5: shell_escape_value rejects newlines/carriage returns/NUL
1686            #[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 6: shell_escape_value handles safe values
1699            #[test]
1700            fn test_shell_escape_value_accepts_safe(s in "[a-zA-Z0-9 !@#$%^&*()_+=\\-\\[\\]{}|;:,./<>?]{0,100}") {
1701                // These don't contain \n, \r, or \0
1702                let result = shell_escape_value(&s);
1703                prop_assert!(result.is_some(), "Should accept safe value: {:?}", s);
1704
1705                // shell_escape only quotes values that need it (contain special chars)
1706                // Simple alphanumeric strings may be returned unquoted
1707                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                    // Values with special chars should be quoted
1716                    prop_assert!(escaped.starts_with('\'') || escaped.contains('\''),
1717                        "Value with special chars should be quoted: {:?} -> {:?}", s, escaped);
1718                }
1719            }
1720
1721            // Test 7: shell_escape_value properly escapes single quotes
1722            #[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                // shell_escape uses '\'' style (end string, escaped quote, start string)
1740                prop_assert!(escaped.contains("'\\''"),
1741                    "Should escape single quote: {} -> {}", value, escaped);
1742            }
1743
1744            // Test 8: build_env_prefix never panics
1745            #[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 9: build_env_prefix rejects invalid keys (non-empty after trim)
1762            #[test]
1763            fn test_build_env_prefix_rejects_invalid_keys(
1764                // Generate keys that are invalid even after trimming
1765                invalid_key in "[0-9][a-zA-Z0-9_]{0,10}"  // Starts with digit
1766            ) {
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                // Key should be rejected since it starts with a digit
1775                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 10: build_env_prefix handles missing values gracefully
1783            #[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                // Empty env - all keys missing
1788                let env: HashMap<String, String> = HashMap::new();
1789                let prefix = build_env_prefix(&keys, |k| env.get(k).cloned());
1790
1791                // Should be empty prefix since no values found
1792                prop_assert!(prefix.prefix.is_empty(), "Should be empty when no values");
1793                prop_assert!(prefix.applied.is_empty());
1794                // Missing values don't count as rejected
1795                prop_assert!(prefix.rejected.is_empty());
1796            }
1797        }
1798
1799        // Targeted edge case tests
1800        #[test]
1801        fn test_shell_escape_edge_cases() {
1802            let _guard = test_guard!();
1803            // Empty string
1804            let result = shell_escape_value("");
1805            assert_eq!(result, Some("''".to_string()));
1806
1807            // Just single quote - shell_escape uses '\'' style (end string, escaped quote, start string)
1808            let result = shell_escape_value("'");
1809            assert_eq!(result, Some("''\\'''".to_string()));
1810
1811            // Multiple single quotes
1812            let result = shell_escape_value("'''");
1813            // shell_escape uses '\'' style for each single quote
1814            assert_eq!(
1815                result
1816                    .as_deref()
1817                    .map(|escaped| escaped.matches("'\\''").count()),
1818                Some(3)
1819            );
1820
1821            // Unicode
1822            let result = shell_escape_value("日本語");
1823            assert!(result.is_some());
1824
1825            // Emoji
1826            let result = shell_escape_value("🔥🚀");
1827            assert!(result.is_some());
1828
1829            // Mixed quotes and special chars
1830            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            // Empty
1838            assert!(!is_valid_env_key(""));
1839
1840            // Single underscore
1841            assert!(is_valid_env_key("_"));
1842
1843            // Single letter
1844            assert!(is_valid_env_key("A"));
1845
1846            // Typical env vars
1847            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            // Invalid: starts with number
1855            assert!(!is_valid_env_key("1VAR"));
1856            assert!(!is_valid_env_key("123"));
1857
1858            // Invalid: contains special chars
1859            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            // Invalid: Unicode
1865            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            // Complex scenario with mixed valid/invalid
1873            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            // VALID should be applied
1894            assert!(prefix.applied.contains(&"VALID".to_string()));
1895            // shell_escape doesn't quote simple alphanumeric strings
1896            assert!(prefix.prefix.contains("VALID=simple"));
1897
1898            // WITH_QUOTE should be applied with escaped quote
1899            assert!(prefix.applied.contains(&"WITH_QUOTE".to_string()));
1900
1901            // NEWLINE should be rejected (unsafe value)
1902            assert!(prefix.rejected.contains(&"NEWLINE".to_string()));
1903
1904            // UNICODE should be applied (safe unicode)
1905            assert!(prefix.applied.contains(&"UNICODE".to_string()));
1906
1907            // EMPTY should be applied
1908            assert!(prefix.applied.contains(&"EMPTY".to_string()));
1909
1910            // 123INVALID should be rejected (invalid key)
1911            assert!(prefix.rejected.contains(&"123INVALID".to_string()));
1912
1913            // MISSING should not appear in either list (not found = silently ignored)
1914            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            // Values that when escaped and passed through shell should reconstruct original
1922            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    // ======================================================================
1942    // System-ssh fallback for Windows workers (bd-kzy2x)
1943    // ======================================================================
1944
1945    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        // The dispatch key in lockstep with `pooled_client_options`:
1954        // declared_os == "windows" -> system-ssh fallback.
1955        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        // declared_os lower-cases the tag's OS, so `os = "Windows"` (mixed
1963        // case) and `os = "WINDOWS"` (all caps) both route to the fallback.
1964        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        // `os:linux` and unlabelled workers do NOT trigger the fallback —
1977        // they keep the openssh-crate path verbatim.
1978        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        // Baseline Windows worker (no port today — WorkerConfig has no
1990        // `port` field; non-default ports travel as `host:port`). The argv
1991        // ends with the destination then a fixed reader, never has a `-p`
1992        // flag, and carries the keepalive when non-zero.
1993        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        // shellexpand::tilde expands `~`; on this machine it goes to a
2005        // $HOME-anchored path, so the suffix is the right invariant.
2006        assert!(
2007            s[2].ends_with(".ssh/surfacebookje_key"),
2008            "identity file is the tilde-expanded path: {}",
2009            s[2]
2010        );
2011        // Pair -o/value: BatchMode=yes, ConnectTimeout=8,
2012        // StrictHostKeyChecking=accept-new, ServerAliveInterval=2.
2013        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        // No -p flag — WorkerConfig has no port.
2019        assert!(!s.iter().any(|arg| arg == "-p"));
2020        // The script never reaches the account's cmd.exe command line.
2021        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        // These are real POSIX shell oracles for the stdin payload; native
2050        // Windows SSH dispatch is qualified separately on a Windows worker.
2051        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        // Force the pipe closed before writing; a tiny concurrent payload
2144        // could fit in the buffer and accidentally miss BrokenPipe coverage.
2145        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        // `server_alive_interval = 0` must NOT emit `-o ServerAliveInterval=0`
2158        // (which OpenSSH treats as "disable keepalives" — fine, but the spec
2159        // calls for omitting the flag entirely at 0).
2160        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        // Non-default keepalive values must be preserved verbatim (no
2178        // truncation, no off-by-one). Use 60s — the same value the openssh
2179        // pool applies for warm masters.
2180        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        // `shellexpand::tilde` is invoked on the identity_file path so an
2194        // operator can write `~/.ssh/foo` in workers.toml and have ssh
2195        // see the expanded path. Compare the raw `~`-prefixed form is
2196        // NOT in the argv.
2197        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        // The dispatch key MUST be the same one `pooled_client_options` uses
2222        // so the two Windows fallbacks (no ControlMaster on the pool side,
2223        // system-ssh on the execute side) cannot diverge on which workers
2224        // they apply to.
2225        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/// How an SSH session's `ControlPersist` should be configured.
2233#[derive(Debug, Clone, PartialEq, Eq)]
2234enum ControlPersistMode {
2235    /// Close the control-master once the initial connection ends (ControlPersist=no).
2236    Closed,
2237    /// Keep a warm control-master for the given idle seconds (ControlPersist=Ns).
2238    IdleFor(NonZeroUsize),
2239    /// Requested idle exceeded usize; caller falls back to Closed.
2240    TooLarge(u64),
2241}
2242
2243/// Decide `ControlPersist` for an SSH session. A warm master is kept ONLY for the
2244/// explicit connection-reuse path (`control_master` + a configured non-zero idle);
2245/// every other case closes after the initial connection so per-call SSH sessions
2246/// (telemetry/health/capabilities) cannot leak `ControlPersist=yes` masters.
2247fn 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        // The per-call paths (control_master=false) must close after use.
2268        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}