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::{Session, Stdio};
11use std::collections::HashMap;
12use std::ffi::OsString;
13use std::num::NonZeroUsize;
14use std::path::{Path, PathBuf};
15use std::process::ExitStatus;
16use std::sync::Arc;
17use std::time::Duration;
18use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
19use tokio::sync::{RwLock, mpsc, watch};
20use tokio::time::Instant;
21use tracing::{debug, error, warn};
22
23// Re-export platform-independent utilities for backwards compatibility
24pub use crate::ssh_utils::{
25    CommandResult, EnvPrefix, build_env_prefix, is_retryable_transport_error,
26    is_retryable_transport_error_text, is_valid_env_key, shell_escape_value,
27};
28
29/// Default SSH connection timeout.
30const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
31
32/// Default command execution timeout.
33const DEFAULT_COMMAND_TIMEOUT: Duration = Duration::from_secs(300);
34
35/// Maximum size for command output (stdout/stderr) to prevent OOM (10MB).
36const MAX_OUTPUT_SIZE: u64 = 10 * 1024 * 1024;
37
38const HEALTH_CHECK_COMMAND: &str = "echo ok";
39
40/// Bound on the pool's reuse liveness probe (`echo ok` over an existing
41/// master). A live master answers in milliseconds; a hung one must not hold a
42/// caller for the client's multi-minute command timeout.
43const POOL_LIVENESS_PROBE_TIMEOUT: Duration = Duration::from_secs(10);
44
45/// Reaping a killed local SSH process must not wedge disconnect/reload either.
46const SSH_CLEANUP_TIMEOUT: Duration = Duration::from_secs(1);
47const SSH_READY_POLL_INTERVAL: Duration = Duration::from_millis(20);
48
49#[derive(Clone)]
50struct SshProcessState {
51    active_commands: usize,
52    idle_since: Instant,
53    draining: bool,
54    stopping: bool,
55    outcome: Option<std::result::Result<ExitStatus, String>>,
56}
57
58/// Owns a foreground SSH process from spawn through wait/reap. Dropping a
59/// connect future or a client requests termination; the supervisor retains the
60/// child and socket directory until it has reaped the child.
61struct OwnedSshProcess {
62    state: watch::Sender<SshProcessState>,
63    supervisor: tokio::task::JoinHandle<()>,
64}
65
66impl OwnedSshProcess {
67    fn spawn(
68        command: &mut tokio::process::Command,
69        control_dir: Option<Arc<tempfile::TempDir>>,
70        idle_timeout: Option<Duration>,
71    ) -> Result<Self> {
72        let mut child = command.kill_on_drop(true).spawn()?;
73        let (state, mut changes) = watch::channel(SshProcessState {
74            active_commands: 0,
75            idle_since: Instant::now(),
76            draining: false,
77            stopping: false,
78            outcome: None,
79        });
80        let updates = state.clone();
81        let supervisor = tokio::spawn(async move {
82            // In particular, do not unlink a control socket while its master
83            // is still alive following cancellation of the caller.
84            let _control_dir = control_dir;
85            let outcome = loop {
86                let current = changes.borrow_and_update().clone();
87                if current.stopping {
88                    if let Err(error) = child.start_kill() {
89                        debug!("Failed to signal owned SSH child: {error}");
90                    }
91                    break child.wait().await;
92                }
93                let idle_deadline = idle_timeout
94                    .filter(|_| current.active_commands == 0)
95                    .and_then(|idle| current.idle_since.checked_add(idle));
96                tokio::select! {
97                    outcome = child.wait() => break outcome,
98                    _ = changes.changed() => {},
99                    _ = async {
100                        if let Some(deadline) = idle_deadline {
101                            tokio::time::sleep_until(deadline).await;
102                        } else {
103                            std::future::pending::<()>().await;
104                        }
105                    } => {
106                        // Serialize expiry with command admission: a command
107                        // starting as the timer fires either owns a use guard
108                        // or observes a closing session before it is spawned.
109                        updates.send_if_modified(|state| {
110                            if !state.stopping && state.active_commands == 0
111                                && idle_timeout.is_some_and(|idle| {
112                                    state.idle_since.elapsed() >= idle
113                                })
114                            {
115                                state.stopping = true;
116                                true
117                            } else {
118                                false
119                            }
120                        });
121                    }
122                }
123            };
124            updates.send_modify(|state| {
125                state.stopping = true;
126                state.outcome = Some(outcome.map_err(|error| error.to_string()));
127            });
128        });
129        Ok(Self { state, supervisor })
130    }
131
132    fn is_running(&self) -> bool {
133        let state = self.state.borrow();
134        !state.draining
135            && !state.stopping
136            && state.outcome.is_none()
137            && !self.supervisor.is_finished()
138    }
139
140    fn begin_use(&self) -> Result<SshProcessUse> {
141        let mut admitted = false;
142        self.state.send_if_modified(|state| {
143            if !state.draining
144                && !state.stopping
145                && state.outcome.is_none()
146                && let Some(active) = state.active_commands.checked_add(1)
147            {
148                state.active_commands = active;
149                admitted = true;
150            }
151            admitted
152        });
153        anyhow::ensure!(admitted, "SSH connection is closing or disconnected");
154        Ok(SshProcessUse {
155            state: self.state.clone(),
156            completed: false,
157        })
158    }
159
160    fn request_stop(&self) {
161        self.state.send_modify(|state| state.stopping = true);
162    }
163
164    async fn wait_until(&self, deadline: Instant) -> Result<ExitStatus> {
165        let mut changes = self.state.subscribe();
166        loop {
167            if let Some(outcome) = changes.borrow_and_update().outcome.clone() {
168                return outcome.map_err(anyhow::Error::msg);
169            }
170            tokio::time::timeout_at(deadline, changes.changed())
171                .await
172                .context(
173                    "SSH connection timed out during authentication or control socket readiness",
174                )?
175                .context("SSH child supervisor stopped unexpectedly")?;
176        }
177    }
178
179    async fn stop_before(&mut self, deadline: Instant) -> Result<()> {
180        self.request_stop();
181        tokio::time::timeout_at(deadline, &mut self.supervisor)
182            .await
183            .context("Timed out reaping owned SSH child")?
184            .context("SSH child supervisor failed")
185    }
186}
187
188impl Drop for OwnedSshProcess {
189    fn drop(&mut self) {
190        // Do not abort the supervisor: it must wait/reap after signalling the
191        // child, even when our caller was itself cancelled at its deadline.
192        self.request_stop();
193    }
194}
195
196struct SshProcessUse {
197    state: watch::Sender<SshProcessState>,
198    completed: bool,
199}
200
201impl SshProcessUse {
202    fn complete(mut self) {
203        self.completed = true;
204    }
205}
206
207impl Drop for SshProcessUse {
208    fn drop(&mut self) {
209        self.state.send_modify(|state| {
210            // A cancelled/failed command may leave remote work attached to
211            // the mux. Retire the transport after its already-running users
212            // finish, without disconnecting those unrelated commands.
213            state.draining |= !self.completed;
214            state.active_commands = state.active_commands.saturating_sub(1);
215            if state.active_commands == 0 {
216                state.idle_since = Instant::now();
217                state.stopping |= state.draining;
218            }
219        });
220    }
221}
222
223struct OwnedSshSession {
224    // Session::resume does not run openssh's unbounded synchronous `ssh -O
225    // exit` in Drop. Our foreground process supervisor performs that cleanup.
226    session: Session,
227    process: OwnedSshProcess,
228    _control_dir: Arc<tempfile::TempDir>,
229}
230
231fn is_expected_health_check_output(stdout: &str) -> bool {
232    stdout
233        .trim()
234        .lines()
235        .last()
236        .is_some_and(is_health_check_sentinel)
237}
238
239fn is_health_check_sentinel(line: &str) -> bool {
240    matches!(line.trim(), "ok")
241}
242
243/// SSH connection options.
244#[derive(Debug, Clone)]
245pub struct SshOptions {
246    /// Wall-clock deadline for connection setup, including authentication and
247    /// ControlMaster readiness. A mux fallback shares the same deadline.
248    pub connect_timeout: Duration,
249    /// Command execution timeout.
250    pub command_timeout: Duration,
251    /// Server keepalive interval (`ssh -o ServerAliveInterval`).
252    ///
253    /// Defaults to `None` (OpenSSH default; keepalive disabled).
254    pub server_alive_interval: Option<Duration>,
255    /// How long the SSH ControlMaster should remain alive while idle.
256    ///
257    /// Only applies when `control_master` is true (connection reuse). `Some(n)`
258    /// with n > 0 keeps the master warm for n idle seconds. Active commands
259    /// prevent expiry. `Some(0s)`/`None`, or a non-mux per-call session, keeps
260    /// the master only for this client's lifetime. RCH owns this idle timer:
261    /// OpenSSH's `ControlPersist` forks an unowned process even without `-f`,
262    /// so the SSH child always uses `ControlPersist=no`.
263    pub control_persist_idle: Option<Duration>,
264    /// SSH control master mode for connection reuse.
265    pub control_master: bool,
266    /// Known hosts policy.
267    pub known_hosts: KnownHostsPolicy,
268}
269
270impl Default for SshOptions {
271    fn default() -> Self {
272        Self {
273            connect_timeout: DEFAULT_CONNECT_TIMEOUT,
274            command_timeout: DEFAULT_COMMAND_TIMEOUT,
275            server_alive_interval: None,
276            control_persist_idle: None,
277            // Default to a plain SSH session. ControlMaster is an optimization
278            // and stale local control sockets can poison otherwise healthy
279            // connections. Callers that explicitly want mux reuse can opt in.
280            control_master: false,
281            known_hosts: KnownHostsPolicy::Add,
282        }
283    }
284}
285
286/// Known hosts policy for SSH connections.
287#[derive(Debug, Clone, Copy, PartialEq, Eq)]
288pub enum KnownHostsPolicy {
289    /// Strictly verify known hosts (recommended for production).
290    Strict,
291    /// Add unknown hosts automatically (for development).
292    Add,
293    /// Accept all hosts without verification (INSECURE - testing only).
294    AcceptAll,
295}
296
297#[cfg(test)]
298mod retry_tests {
299    use super::*;
300    use crate::test_guard;
301
302    #[test]
303    fn test_retryable_transport_error_text() {
304        let _guard = test_guard!();
305        assert!(is_retryable_transport_error_text(
306            "ssh: connect to host 1.2.3.4 port 22: Connection timed out"
307        ));
308        assert!(is_retryable_transport_error_text(
309            "kex_exchange_identification: Connection reset by peer"
310        ));
311        assert!(is_retryable_transport_error_text("Broken pipe"));
312        assert!(is_retryable_transport_error_text("Network is unreachable"));
313    }
314
315    #[test]
316    fn test_non_retryable_transport_error_text() {
317        let _guard = test_guard!();
318        assert!(!is_retryable_transport_error_text(
319            "Permission denied (publickey)."
320        ));
321        assert!(!is_retryable_transport_error_text(
322            "Host key verification failed."
323        ));
324        assert!(!is_retryable_transport_error_text(
325            "Could not resolve hostname worker.example.com: Name or service not known"
326        ));
327        assert!(!is_retryable_transport_error_text(
328            "Identity file /nope/id_rsa not accessible: No such file or directory"
329        ));
330    }
331}
332
333/// SSH client for a single worker connection.
334pub struct SshClient {
335    /// Worker configuration.
336    config: WorkerConfig,
337    /// SSH options.
338    options: SshOptions,
339    /// Active SSH session (if connected).
340    session: Option<OwnedSshSession>,
341}
342
343impl SshClient {
344    /// Create a new SSH client for a worker.
345    pub fn new(config: WorkerConfig, options: SshOptions) -> Self {
346        Self {
347            config,
348            options,
349            session: None,
350        }
351    }
352
353    /// Get the worker ID.
354    pub fn worker_id(&self) -> &WorkerId {
355        &self.config.id
356    }
357
358    /// Check if connected to the worker.
359    pub fn is_connected(&self) -> bool {
360        self.session
361            .as_ref()
362            .is_some_and(|session| session.process.is_running())
363    }
364
365    #[cfg(test)]
366    fn is_configured_for(&self, config: &WorkerConfig) -> bool {
367        same_ssh_endpoint(&self.config, config)
368    }
369
370    /// Connect to the remote worker.
371    pub async fn connect(&mut self) -> Result<()> {
372        self.connect_using(Path::new("ssh"), None).await
373    }
374
375    async fn connect_using(
376        &mut self,
377        ssh_program: &Path,
378        control_directory: Option<&Path>,
379    ) -> Result<()> {
380        if self.is_connected() {
381            debug!("Already connected to {}", self.config.id);
382            return Ok(());
383        }
384        // A master that exited or expired is no longer a usable session.
385        self.session.take();
386        let deadline = Instant::now()
387            .checked_add(self.options.connect_timeout)
388            .context("SSH connect_timeout exceeds the supported clock range")?;
389
390        let destination = format!("{}@{}", self.config.user, self.config.host);
391        debug!("Connecting to {} via SSH...", destination);
392
393        let session = match self
394            .connect_with_mode(
395                &destination,
396                self.options.control_master,
397                ssh_program,
398                control_directory,
399                deadline,
400            )
401            .await
402        {
403            Ok(session) => session,
404            Err(primary_error) if self.options.control_master && Instant::now() < deadline => {
405                warn!(
406                    "SSH ControlMaster connection to {} failed ({}). Retrying without ControlMaster.",
407                    destination, primary_error
408                );
409                self.connect_with_mode(
410                    &destination,
411                    false,
412                    ssh_program,
413                    control_directory,
414                    deadline,
415                )
416                .await
417                .with_context(|| {
418                    format!(
419                        "Failed to connect to {} after retrying without ControlMaster",
420                        destination
421                    )
422                })?
423            }
424            Err(primary_error) => {
425                return Err(primary_error)
426                    .with_context(|| format!("Failed to connect to {}", destination));
427            }
428        };
429
430        // debug, not info: the telemetry pool connects to every worker each poll
431        // cycle, so at info this floods the daemon log (8M+ lines / multi-GB).
432        debug!("Connected to {} ({})", self.config.id, self.config.host);
433        self.session = Some(session);
434        Ok(())
435    }
436
437    async fn connect_with_mode(
438        &self,
439        destination: &str,
440        control_master: bool,
441        ssh_program: &Path,
442        control_directory: Option<&Path>,
443        deadline: Instant,
444    ) -> Result<OwnedSshSession> {
445        anyhow::ensure!(
446            Instant::now() < deadline,
447            "SSH connection timed out before spawn"
448        );
449        let directory = control_directory
450            .map(Path::to_path_buf)
451            .unwrap_or_else(ssh_control_directory);
452        std::fs::create_dir_all(&directory)
453            .with_context(|| format!("Failed to create SSH control directory {directory:?}"))?;
454        {
455            use std::os::unix::fs::PermissionsExt;
456            std::fs::set_permissions(&directory, std::fs::Permissions::from_mode(0o700))
457                .context("Failed to protect SSH control directory")?;
458        }
459        let control_dir = Arc::new(
460            tempfile::Builder::new()
461                .prefix("rch-")
462                .tempdir_in(directory)
463                .context("Failed to create private SSH control directory")?,
464        );
465        let socket = control_dir.path().join("master");
466        let log = control_dir.path().join("log");
467        let idle_timeout =
468            match control_persist_mode(control_master, self.options.control_persist_idle) {
469                ControlPersistMode::IdleFor(seconds) => {
470                    Some(Duration::from_secs(seconds.get() as u64))
471                }
472                ControlPersistMode::TooLarge(seconds) => {
473                    warn!("control_persist_idle too large ({seconds}s); closing with client");
474                    None
475                }
476                ControlPersistMode::Closed => None,
477            };
478        if let Some(idle) = idle_timeout {
479            anyhow::ensure!(
480                Instant::now().checked_add(idle).is_some(),
481                "SSH idle timeout exceeds the supported clock range"
482            );
483        }
484        // -F /dev/null probes the client itself without running user Match
485        // directives. Older clients do not understand ForkAfterAuthentication;
486        // a global IgnoreUnknown override would break the user's own list.
487        let can_disable_background =
488            ssh_can_disable_background(ssh_program, control_dir.clone(), deadline).await?;
489        let mut command = self.master_command(
490            ssh_program,
491            destination,
492            &socket,
493            &log,
494            can_disable_background,
495        );
496        let mut process =
497            OwnedSshProcess::spawn(&mut command, Some(control_dir.clone()), idle_timeout)
498                .with_context(|| format!("Failed to spawn SSH connection to {destination}"))?;
499        // Authentication is activity too; a short idle policy must not cut it
500        // off before the connection deadline.
501        let connecting = process.begin_use()?;
502        let ready = wait_for_ssh_master(&process, ssh_program, &socket, &log, deadline).await;
503        if let Err(error) = ready {
504            // Cleanup consumes only the remaining connection budget. If the
505            // caller's deadline already fired, the supervisor still owns and
506            // reaps the child after this future returns or is cancelled.
507            let _ = process.stop_before(deadline).await;
508            return Err(error).with_context(|| format!("Failed to connect to {destination}"));
509        }
510        connecting.complete();
511        Ok(OwnedSshSession {
512            session: Session::resume(socket.into_boxed_path(), Some(log.into_boxed_path())),
513            process,
514            _control_dir: control_dir,
515        })
516    }
517
518    fn master_command(
519        &self,
520        ssh_program: &Path,
521        destination: &str,
522        socket: &Path,
523        log: &Path,
524        can_disable_background: bool,
525    ) -> tokio::process::Command {
526        let mut command = tokio::process::Command::new(ssh_program);
527        command
528            .stdin(std::process::Stdio::null())
529            .stdout(std::process::Stdio::null())
530            .stderr(std::process::Stdio::null())
531            .arg("-E")
532            .arg(log)
533            .arg("-S")
534            .arg(socket)
535            .args(["-M", "-N", "-T", "-o", "BatchMode=yes"])
536            // ControlPersist forks after authentication even without `-f`.
537            // Keep the master in this owned foreground child instead.
538            .args(["-o", "ControlPersist=no"])
539            .arg("-o")
540            .arg(match self.options.known_hosts {
541                KnownHostsPolicy::Strict => "StrictHostKeyChecking=yes",
542                KnownHostsPolicy::Add => "StrictHostKeyChecking=accept-new",
543                KnownHostsPolicy::AcceptAll => "StrictHostKeyChecking=no",
544            })
545            .arg("-o")
546            .arg(format!(
547                "ConnectTimeout={}",
548                self.options.connect_timeout.as_secs()
549            ));
550        if can_disable_background {
551            command.args(["-o", "ForkAfterAuthentication=no"]);
552        }
553        if let Some(interval) = self.options.server_alive_interval {
554            command
555                .arg("-o")
556                .arg(format!("ServerAliveInterval={}", interval.as_secs()));
557        }
558        let identity = shellexpand::tilde(&self.config.identity_file);
559        if Path::new(identity.as_ref()).exists() {
560            command
561                .arg("-i")
562                .arg(identity.as_ref())
563                .args(["-o", "IdentitiesOnly=yes"]);
564        }
565        command.arg(destination);
566        command
567    }
568
569    /// Disconnect from the worker.
570    pub async fn disconnect(&mut self) -> Result<()> {
571        if let Some(mut session) = self.session.take() {
572            debug!("Disconnecting from {}", self.config.id);
573            session
574                .process
575                .stop_before(Instant::now() + SSH_CLEANUP_TIMEOUT)
576                .await?;
577            // debug, not info: paired with the connect log above; floods at info.
578            debug!("Disconnected from {}", self.config.id);
579        }
580        Ok(())
581    }
582
583    /// Force a fresh connection, dropping any existing (possibly dead) session.
584    ///
585    /// `reconnect` tears the old master down even if its local process is
586    /// running but its transport no longer answers, then establishes a new one.
587    pub async fn reconnect(&mut self) -> Result<()> {
588        if let Err(e) = self.disconnect().await {
589            debug!(
590                "Ignoring error closing stale session to {} before reconnect: {}",
591                self.config.id, e
592            );
593        }
594        self.connect().await
595    }
596
597    /// Execute a command on the remote worker using this client's configured
598    /// command timeout.
599    pub async fn execute(&self, command: &str) -> Result<CommandResult> {
600        self.execute_with_timeout(command, self.options.command_timeout)
601            .await
602    }
603
604    /// Execute a command on the remote worker with an explicit per-call timeout.
605    ///
606    /// Used by the connection pool ([`SshPool::run_with_timeout`]) so a single
607    /// pooled client can serve calls with different timeouts (a fast health
608    /// probe vs. a slower cleanup) without mutating the shared client options.
609    pub async fn execute_with_timeout(
610        &self,
611        command: &str,
612        command_timeout: Duration,
613    ) -> Result<CommandResult> {
614        // Windows fallback (bd-kzy2x): the openssh crate cannot execute
615        // commands on Windows OpenSSH (the slave never completes the
616        // command channel), so we dispatch through the system `ssh` binary
617        // for declared-OS Windows workers. Connection setup still uses our
618        // supervised SSH master; Linux / unlabelled workers execute through
619        // the openssh session attached to that master.
620        if prefers_system_ssh(&self.config) {
621            return system_ssh_execute(
622                &self.config,
623                command,
624                command_timeout,
625                self.options
626                    .server_alive_interval
627                    .unwrap_or(Duration::from_secs(2)),
628            )
629            .await;
630        }
631
632        let session = self.session.as_ref().context("Not connected to worker")?;
633        let activity = session.process.begin_use()?;
634
635        let start = std::time::Instant::now();
636        debug!(
637            "Executing on {}: {}",
638            self.config.id,
639            crate::util::mask_sensitive_command(command)
640        );
641
642        let mut child = session
643            .session
644            .command("sh")
645            .arg("-c")
646            .arg(command)
647            .stdout(Stdio::piped())
648            .stderr(Stdio::piped())
649            .spawn()
650            .await
651            .with_context(|| format!("Failed to spawn command on {}", self.config.id))?;
652
653        let execution_future = async {
654            // Read stdout and stderr concurrently to avoid deadlock if one pipe fills.
655            let stdout_handle = child.stdout().take();
656            let stderr_handle = child.stderr().take();
657
658            let stdout_fut = async {
659                if let Some(out) = stdout_handle {
660                    let reader = BufReader::new(out);
661                    let mut take = reader.take(MAX_OUTPUT_SIZE);
662                    let mut buf = String::new();
663                    take.read_to_string(&mut buf).await?;
664                    // Drain the rest to prevent SIGPIPE or blocking
665                    let mut reader = take.into_inner();
666                    let mut sink = tokio::io::sink();
667                    tokio::io::copy(&mut reader, &mut sink).await?;
668                    if buf.len() >= MAX_OUTPUT_SIZE as usize {
669                        buf.push_str("\n...[output truncated]...\n");
670                    }
671                    Ok::<String, anyhow::Error>(buf)
672                } else {
673                    Ok(String::new())
674                }
675            };
676
677            let stderr_fut = async {
678                if let Some(err) = stderr_handle {
679                    let reader = BufReader::new(err);
680                    let mut take = reader.take(MAX_OUTPUT_SIZE);
681                    let mut buf = String::new();
682                    take.read_to_string(&mut buf).await?;
683                    // Drain the rest to prevent SIGPIPE or blocking
684                    let mut reader = take.into_inner();
685                    let mut sink = tokio::io::sink();
686                    tokio::io::copy(&mut reader, &mut sink).await?;
687                    if buf.len() >= MAX_OUTPUT_SIZE as usize {
688                        buf.push_str("\n...[output truncated]...\n");
689                    }
690                    Ok::<String, anyhow::Error>(buf)
691                } else {
692                    Ok(String::new())
693                }
694            };
695
696            let (stdout, stderr) = tokio::try_join!(stdout_fut, stderr_fut)?;
697
698            let status = child
699                .wait()
700                .await
701                .with_context(|| "Failed to wait for command completion")?;
702
703            Ok::<_, anyhow::Error>((status, stdout, stderr))
704        };
705
706        match tokio::time::timeout(command_timeout, execution_future).await {
707            Ok(result) => {
708                let (status, stdout, stderr) = result?;
709                let duration = start.elapsed();
710                let exit_code = status.code().unwrap_or(-1);
711
712                debug!(
713                    "Command completed on {} (exit={}, duration={}ms)",
714                    self.config.id,
715                    exit_code,
716                    duration.as_millis()
717                );
718
719                activity.complete();
720                Ok(CommandResult {
721                    exit_code,
722                    stdout,
723                    stderr,
724                    duration_ms: duration.as_millis() as u64,
725                })
726            }
727            Err(_) => {
728                // openssh kills the local slave when its RemoteChild drops.
729                // The incomplete activity guard also retires the master once
730                // other active commands finish, so cancelled remote work is
731                // not left attached to an indefinitely warm connection.
732                warn!(
733                    "Command timed out on {} after {:?}",
734                    self.config.id, command_timeout
735                );
736                anyhow::bail!("Command timed out after {:?}", command_timeout);
737            }
738        }
739    }
740
741    /// Execute a command and stream output in real-time.
742    pub async fn execute_streaming<F, G>(
743        &self,
744        command: &str,
745        mut on_stdout: F,
746        mut on_stderr: G,
747    ) -> Result<CommandResult>
748    where
749        F: FnMut(&str),
750        G: FnMut(&str),
751    {
752        let session = self.session.as_ref().context("Not connected to worker")?;
753        let activity = session.process.begin_use()?;
754
755        let start = std::time::Instant::now();
756        debug!(
757            "Executing (streaming) on {}: {}",
758            self.config.id,
759            crate::util::mask_sensitive_command(command)
760        );
761
762        let mut child = session
763            .session
764            .command("sh")
765            .arg("-c")
766            .arg(command)
767            .stdout(Stdio::piped())
768            .stderr(Stdio::piped())
769            .spawn()
770            .await
771            .with_context(|| format!("Failed to spawn command on {}", self.config.id))?;
772
773        let stdout = child.stdout().take();
774        let stderr = child.stderr().take();
775
776        // Use a channel to aggregate stream events from reader tasks.
777        // This avoids cancellation safety issues with select! over AsyncBufReadExt::read_line.
778        let (tx, mut rx) = mpsc::channel(100);
779
780        // Spawn stdout reader
781        if let Some(out) = stdout {
782            let tx = tx.clone();
783            tokio::spawn(async move {
784                let mut reader = BufReader::new(out);
785                let mut line = String::new();
786                loop {
787                    line.clear();
788                    match reader.read_line(&mut line).await {
789                        Ok(0) => break, // EOF
790                        Ok(_) => {
791                            if tx.send(StreamEvent::Stdout(line.clone())).await.is_err() {
792                                break; // Receiver dropped
793                            }
794                        }
795                        Err(_) => break, // Read error
796                    }
797                }
798            });
799        }
800
801        // Spawn stderr reader
802        if let Some(err) = stderr {
803            let tx = tx.clone();
804            tokio::spawn(async move {
805                let mut reader = BufReader::new(err);
806                let mut line = String::new();
807                loop {
808                    line.clear();
809                    match reader.read_line(&mut line).await {
810                        Ok(0) => break, // EOF
811                        Ok(_) => {
812                            if tx.send(StreamEvent::Stderr(line.clone())).await.is_err() {
813                                break; // Receiver dropped
814                            }
815                        }
816                        Err(_) => break, // Read error
817                    }
818                }
819            });
820        }
821
822        // Drop original tx so rx closes when tasks finish
823        drop(tx);
824
825        let mut stdout_acc = String::new();
826        let mut stderr_acc = String::new();
827
828        enum StreamEvent {
829            Stdout(String),
830            Stderr(String),
831        }
832
833        let streaming_future = async {
834            // Process events until channel closes (EOF from both streams)
835            while let Some(event) = rx.recv().await {
836                match event {
837                    StreamEvent::Stdout(line) => {
838                        on_stdout(&line);
839                        if stdout_acc.len() < MAX_OUTPUT_SIZE as usize {
840                            stdout_acc.push_str(&line);
841                            if stdout_acc.len() >= MAX_OUTPUT_SIZE as usize {
842                                stdout_acc.push_str("\n...[output truncated]...\n");
843                            }
844                        }
845                    }
846                    StreamEvent::Stderr(line) => {
847                        on_stderr(&line);
848                        if stderr_acc.len() < MAX_OUTPUT_SIZE as usize {
849                            stderr_acc.push_str(&line);
850                            if stderr_acc.len() >= MAX_OUTPUT_SIZE as usize {
851                                stderr_acc.push_str("\n...[output truncated]...\n");
852                            }
853                        }
854                    }
855                }
856            }
857
858            let status = child.wait().await?;
859            Ok::<_, anyhow::Error>(status)
860        };
861
862        match tokio::time::timeout(self.options.command_timeout, streaming_future).await {
863            Ok(result) => {
864                let status = result?;
865                let duration = start.elapsed();
866                let exit_code = status.code().unwrap_or(-1);
867
868                activity.complete();
869                Ok(CommandResult {
870                    exit_code,
871                    stdout: stdout_acc,
872                    stderr: stderr_acc,
873                    duration_ms: duration.as_millis() as u64,
874                })
875            }
876            Err(_) => {
877                // Dropping RemoteChild kills the local slave and closes its
878                // pipes, ending the reader tasks. The incomplete activity
879                // guard retires the master after other active commands finish.
880                warn!(
881                    "Command (streaming) timed out on {} after {:?}, cleaning up",
882                    self.config.id, self.options.command_timeout
883                );
884                anyhow::bail!("Command timed out after {:?}", self.options.command_timeout);
885            }
886        }
887    }
888
889    /// Check if the worker is reachable via SSH.
890    pub async fn health_check(&self) -> Result<bool> {
891        self.health_check_with_timeout(self.options.command_timeout)
892            .await
893    }
894
895    /// [`health_check`](Self::health_check) bounded by `timeout` instead of
896    /// the client's command timeout (300s for pooled clients).
897    pub async fn health_check_with_timeout(&self, timeout: Duration) -> Result<bool> {
898        match self
899            .execute_with_timeout(HEALTH_CHECK_COMMAND, timeout)
900            .await
901        {
902            Ok(result) => Ok(result.success() && is_expected_health_check_output(&result.stdout)),
903            Err(e) => {
904                warn!("Health check failed for {}: {}", self.config.id, e);
905                Ok(false)
906            }
907        }
908    }
909}
910
911fn ssh_control_directory() -> PathBuf {
912    // Keep socket names short on macOS, where temp_dir commonly lives under
913    // a long /var/folders path. Each connection gets a fresh private subdir.
914    if let Some(home) = dirs::home_dir() {
915        home.join(".ssh").join("rch")
916    } else if let Some(runtime) = dirs::runtime_dir() {
917        runtime.join("rch-ssh")
918    } else {
919        let username = whoami::username().unwrap_or_else(|_| "unknown".to_owned());
920        std::env::temp_dir().join(format!("rch-ssh-{username}"))
921    }
922}
923
924async fn ssh_can_disable_background(
925    ssh_program: &Path,
926    control_dir: Arc<tempfile::TempDir>,
927    deadline: Instant,
928) -> Result<bool> {
929    use std::io::Read;
930
931    anyhow::ensure!(
932        Instant::now() < deadline,
933        "SSH connection timed out before capability probe"
934    );
935    let configuration = control_dir.path().join("client-config");
936    let mut command = tokio::process::Command::new(ssh_program);
937    command
938        .args(["-G", "-F", "/dev/null", "-T", "none"])
939        .stdin(std::process::Stdio::null())
940        .stdout(std::fs::File::create(&configuration)?)
941        .stderr(std::process::Stdio::null());
942    let mut probe = OwnedSshProcess::spawn(&mut command, Some(control_dir.clone()), None)?;
943    let checked = probe.wait_until(deadline).await;
944    let _ = probe.stop_before(deadline).await;
945    anyhow::ensure!(checked?.success(), "SSH client capability probe failed");
946    let mut output = String::new();
947    std::fs::File::open(configuration)?
948        .take(65_537)
949        .read_to_string(&mut output)?;
950    anyhow::ensure!(
951        output.len() <= 65_536,
952        "SSH client capability response is too large"
953    );
954    anyhow::ensure!(
955        output.lines().any(|line| {
956            let mut words = line.split_whitespace();
957            words.next() == Some("hostname") && words.next() == Some("none")
958        }),
959        "SSH client returned an invalid capability response"
960    );
961    Ok(output.lines().any(|line| {
962        line.split_whitespace()
963            .next()
964            .is_some_and(|key| key.eq_ignore_ascii_case("forkafterauthentication"))
965    }))
966}
967
968async fn wait_for_ssh_master(
969    process: &OwnedSshProcess,
970    ssh_program: &Path,
971    socket: &Path,
972    log: &Path,
973    deadline: Instant,
974) -> Result<()> {
975    loop {
976        let outcome = process.state.borrow().outcome.clone();
977        if let Some(outcome) = outcome {
978            let status = outcome.map_err(anyhow::Error::msg)?;
979            let diagnostic = ssh_connect_diagnostic(log);
980            anyhow::bail!("SSH master exited before becoming ready ({status}): {diagnostic}");
981        }
982        anyhow::ensure!(
983            Instant::now() < deadline,
984            "SSH connection timed out during authentication or control socket readiness"
985        );
986        if socket.exists() {
987            // Socket existence alone does not prove the mux accepts commands.
988            // This probe is also owned and shares the authentication deadline;
989            // openssh::Session::check itself does not kill its child on drop.
990            let mut command = tokio::process::Command::new(ssh_program);
991            command
992                .stdin(std::process::Stdio::null())
993                .stdout(std::process::Stdio::null())
994                .stderr(std::process::Stdio::null())
995                .arg("-S")
996                .arg(socket)
997                .args(["-o", "BatchMode=yes", "-O", "check", "none"]);
998            let mut probe = OwnedSshProcess::spawn(&mut command, None, None)?;
999            let checked = probe.wait_until(deadline).await;
1000            let _ = probe.stop_before(deadline).await;
1001            if checked?.success() && process.is_running() {
1002                return Ok(());
1003            }
1004        }
1005        tokio::time::sleep_until(deadline.min(Instant::now() + SSH_READY_POLL_INTERVAL)).await;
1006    }
1007}
1008
1009fn ssh_connect_diagnostic(log: &Path) -> String {
1010    use std::io::Read;
1011
1012    let mut diagnostic = Vec::new();
1013    if let Ok(file) = std::fs::File::open(log) {
1014        let _ = file.take(16 * 1024).read_to_end(&mut diagnostic);
1015    }
1016    String::from_utf8_lossy(&diagnostic).trim().to_owned()
1017}
1018
1019/// Derive the per-worker [`SshOptions`] a pool hands to each pooled
1020/// [`SshClient`].
1021///
1022/// POLICY — Windows workers never get ControlMaster multiplexing. Windows
1023/// OpenSSH has no ControlMaster support: a pooled (multiplexed) session
1024/// appears to connect, then every command over it hangs and dies at the
1025/// command stage ("Failed to wait for command completion"), so each pooled
1026/// probe fails deterministically and the daemon marks the worker permanently
1027/// unreachable even though fresh one-shot SSH to the same host succeeds
1028/// (2026-08-30 incident: worker `wsurf`, `os = "windows"` — healthy via
1029/// `rch workers probe` fresh connections, "unreachable" via every pooled
1030/// health check). Applying the override HERE — the single place pooled
1031/// clients are constructed from the pool's global options — fixes every pool
1032/// consumer (daemon shared build/telemetry pool, dedicated health pool) at
1033/// once. [`SshClient::connect`]'s ControlMaster-failure retry cannot catch
1034/// this class: the mux connect does not fail, the commands do.
1035///
1036/// `control_persist_idle` needs no separate neutralization: it is consulted
1037/// only when `control_master` is true (see [`control_persist_mode`]), so it
1038/// is inert once mux is disabled and is left verbatim. Non-Windows workers
1039/// get the pool options verbatim in full.
1040fn pooled_client_options(pool_options: &SshOptions, config: &WorkerConfig) -> SshOptions {
1041    let mut options = pool_options.clone();
1042    if declared_os(&config.tags).as_deref() == Some("windows") {
1043        options.control_master = false;
1044    }
1045    options
1046}
1047
1048// System-ssh command-execution fallback (bd-kzy2x).
1049//
1050// The `openssh` crate 0.11.6 cannot execute commands on Windows OpenSSH
1051// regardless of `control_master` / `control_persist` / keepalive settings:
1052// the slave never completes the command channel and every command hangs at
1053// the execute stage even though the connect succeeds. The CLI `ssh` binary
1054// speaks the OpenSSH wire protocol directly and works fine. For workers that
1055// declare `os = "windows"` we therefore spawn the system `ssh` binary at
1056// command-execution time, mirroring the proven system-ssh pattern already
1057// used by the fleet preflight path (`rch/src/fleet/ssh.rs::SshExecutor`) and
1058// the CLI worker init / probe paths. Connection setup uses the supervised
1059// master above; the fallback is an execute-time dispatch inside
1060// `SshClient::execute_with_timeout`.
1061//
1062// Policy key (kept in lockstep with the existing pool-layer override
1063// `pooled_client_options`): `declared_os(&config.tags) == Some("windows")`.
1064// `os = "Windows"` is case-normalized by `declared_os` so both spellings
1065// route through the same fallback.
1066
1067/// Should `SshClient::execute_with_timeout` dispatch through the system `ssh`
1068/// binary for this worker instead of the `openssh` crate?
1069pub(crate) fn prefers_system_ssh(config: &WorkerConfig) -> bool {
1070    declared_os(&config.tags).as_deref() == Some("windows")
1071}
1072
1073/// Remote command and optional stdin payload for a system SSH invocation.
1074#[derive(Debug, Clone, PartialEq, Eq)]
1075pub struct RemoteShellCommand {
1076    /// Command argument passed to SSH after the destination.
1077    pub command: String,
1078    /// Script to write and close concurrently with draining stdout/stderr.
1079    /// When absent, the caller should provide null stdin.
1080    pub stdin_script: Option<String>,
1081}
1082
1083/// Prepare a POSIX script without exposing its contents to Windows cmd.exe.
1084/// Windows uses a fixed `sh -s` reader, which replaces itself with a command
1085/// shell whose stdin is empty. An exact `sh -s` keeps its existing EOF behavior.
1086/// POSIX workers receive the original command argument and null stdin.
1087#[must_use]
1088pub fn remote_shell_command(config: &WorkerConfig, command: &str) -> RemoteShellCommand {
1089    if prefers_system_ssh(config) {
1090        RemoteShellCommand {
1091            command: "sh -s".into(),
1092            stdin_script: (command != "sh -s").then(|| {
1093                format!(
1094                    "exec sh -c {} </dev/null\n",
1095                    shell_escape::escape(command.into())
1096                )
1097            }),
1098        }
1099    } else {
1100        RemoteShellCommand {
1101            command: command.into(),
1102            stdin_script: None,
1103        }
1104    }
1105}
1106
1107pub use crate::ssh_utils::identities_only_args;
1108
1109/// Build the argv for the system-ssh fallback, mirroring the proven CLI
1110/// system-ssh pattern (see `rch/src/fleet/ssh.rs::SshExecutor::build_ssh_args`
1111/// and `rch/src/commands/workers_init.rs`). Pure / testable: no process is
1112/// spawned. The first element is the program name `"ssh"` so the vector can
1113/// be passed straight to `Command::new` callers that prefer the explicit
1114/// first arg, but the standard `Command::new("ssh").args(...)` callers can
1115/// skip it — see `system_ssh_execute` which drops it.
1116///
1117/// Argv layout:
1118/// `ssh -i <identity_file> -o BatchMode=yes -o ConnectTimeout=8 -o
1119///   StrictHostKeyChecking=accept-new -o ServerAliveInterval=<secs>
1120///   <user>@<host> <command>`
1121///
1122/// `WorkerConfig` has no port field today; the SSH port is carried as part
1123/// of the host (`host:port`) and SSH's own config handles non-default
1124/// ports, so the argv has no `-p` flag. The 8s `ConnectTimeout` matches the
1125/// fleet SshExecutor default (`DEFAULT_CONNECT_TIMEOUT_SECS = 10` would be
1126/// the natural match too — kept at 8 to match the FINDINGS spec exactly).
1127/// `server_alive_interval` is rounded up to 1s minimum when non-zero to
1128/// avoid `ServerAliveInterval=0` (which OpenSSH treats as "disable
1129/// keepalives" — fine — but matches the spec "0 -> omit" intent).
1130/// Windows receives only the fixed POSIX reader `sh -s`; the original script
1131/// travels through stdin so cmd.exe never interprets its contents.
1132pub(crate) fn system_ssh_argv(
1133    config: &WorkerConfig,
1134    command: &RemoteShellCommand,
1135    server_alive_interval: Duration,
1136) -> Vec<OsString> {
1137    let identity_path = shellexpand::tilde(&config.identity_file);
1138    let destination = format!("{}@{}", config.user, config.host);
1139
1140    let mut argv: Vec<OsString> = Vec::with_capacity(12);
1141    argv.push(OsString::from("ssh"));
1142    argv.push(OsString::from("-i"));
1143    argv.push(OsString::from(identity_path.as_ref()));
1144    if let Some(opts) = identities_only_args(&config.identity_file) {
1145        argv.extend(opts.map(OsString::from));
1146    }
1147    argv.push(OsString::from("-o"));
1148    argv.push(OsString::from("BatchMode=yes"));
1149    argv.push(OsString::from("-o"));
1150    argv.push(OsString::from("ConnectTimeout=8"));
1151    argv.push(OsString::from("-o"));
1152    argv.push(OsString::from("StrictHostKeyChecking=accept-new"));
1153    if !server_alive_interval.is_zero() {
1154        // `0` means "no keepalive" in OpenSSH; otherwise emit the requested
1155        // interval in whole seconds (sub-second values are rounded up to 1s
1156        // because OpenSSH only accepts integer seconds for this option).
1157        let secs = server_alive_interval.as_secs().max(1);
1158        argv.push(OsString::from("-o"));
1159        argv.push(OsString::from(format!("ServerAliveInterval={secs}")));
1160    }
1161    argv.push(OsString::from(destination));
1162    argv.push(OsString::from(&command.command));
1163    argv
1164}
1165
1166/// Write and close an SSH script pipe within the caller's execution deadline.
1167/// Run concurrently with output draining. A closed pipe is allowed so the
1168/// caller can collect an early SSH failure's actual exit status and stderr.
1169///
1170/// # Errors
1171/// Returns write or shutdown errors other than an expected broken pipe.
1172pub async fn write_ssh_command_stdin(
1173    mut stdin: tokio::process::ChildStdin,
1174    input: &str,
1175) -> std::io::Result<()> {
1176    use tokio::io::AsyncWriteExt;
1177
1178    let result = async {
1179        stdin.write_all(input.as_bytes()).await?;
1180        stdin.shutdown().await
1181    }
1182    .await;
1183    match result {
1184        // An authentication failure or early remote exit can close stdin
1185        // before the payload is written. Preserve its actual status/stderr.
1186        Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),
1187        result => result,
1188    }
1189}
1190
1191/// Execute a command on a Windows worker via the system `ssh` binary.
1192///
1193/// Mirrors the openssh-crate path's size cap (`MAX_OUTPUT_SIZE`), timeout
1194/// semantics (tokio timeout that bails on expiry), and `CommandResult`
1195/// shape. The local ssh process is set `kill_on_drop(true)` so a timeout
1196/// cannot leak an `ssh` child — the openssh-crate timeout path likewise
1197/// relies on the caller to drop the child (and warns about the leak risk);
1198/// here we get a hard SIGKILL on drop instead. Output reads are
1199/// concurrent over `tokio::io::BufReader` to avoid a single-stream back-
1200/// pressure deadlock (same pattern as the openssh path). `env_clear()` is
1201/// used so a hostile / noisy environment from the calling daemon cannot
1202/// perturb the local `ssh` invocation; SSH agent forwarding and standard
1203/// paths still resolve via `~/.ssh/config` because ssh reads them from the
1204/// filesystem, not the env.
1205pub(crate) async fn system_ssh_execute(
1206    config: &WorkerConfig,
1207    command: &str,
1208    command_timeout: Duration,
1209    server_alive_interval: Duration,
1210) -> Result<CommandResult> {
1211    use std::process::Stdio;
1212    use tokio::io::AsyncReadExt;
1213    use tokio::process::Command;
1214
1215    let prepared = remote_shell_command(config, command);
1216    let argv = system_ssh_argv(config, &prepared, server_alive_interval);
1217    let input = prepared.stdin_script;
1218    // Drop the program name — `Command::new("ssh")` is the standard form.
1219    let args = argv.into_iter().skip(1);
1220
1221    let start = std::time::Instant::now();
1222    debug!(
1223        "system-ssh executing on {}: {}",
1224        config.id,
1225        crate::util::mask_sensitive_command(command)
1226    );
1227
1228    let mut child = Command::new("ssh")
1229        .args(args)
1230        .env_clear()
1231        .stdin(if input.is_some() {
1232            Stdio::piped()
1233        } else {
1234            Stdio::null()
1235        })
1236        .stdout(Stdio::piped())
1237        .stderr(Stdio::piped())
1238        .kill_on_drop(true)
1239        .spawn()
1240        .with_context(|| {
1241            format!(
1242                "failed to spawn system ssh for {} (Windows fallback)",
1243                config.id
1244            )
1245        })?;
1246
1247    let execution = async {
1248        let stdout = child.stdout.take();
1249        let stderr = child.stderr.take();
1250        let stdin = child.stdin.take();
1251
1252        let stdin_fut = async {
1253            if let (Some(stdin), Some(input)) = (stdin, input) {
1254                write_ssh_command_stdin(stdin, &input).await?;
1255            }
1256            Ok::<(), anyhow::Error>(())
1257        };
1258
1259        let stdout_fut = async {
1260            let mut buf = String::new();
1261            if let Some(out) = stdout {
1262                let mut handle = out.take(MAX_OUTPUT_SIZE);
1263                handle.read_to_string(&mut buf).await?;
1264            }
1265            Ok::<String, anyhow::Error>(buf)
1266        };
1267        let stderr_fut = async {
1268            let mut buf = String::new();
1269            if let Some(err) = stderr {
1270                let mut handle = err.take(MAX_OUTPUT_SIZE);
1271                handle.read_to_string(&mut buf).await?;
1272            }
1273            Ok::<String, anyhow::Error>(buf)
1274        };
1275
1276        let ((), stdout, stderr) = tokio::try_join!(stdin_fut, stdout_fut, stderr_fut)?;
1277        let status = child
1278            .wait()
1279            .await
1280            .with_context(|| "system ssh command failed to wait for completion")?;
1281        Ok::<_, anyhow::Error>((status, stdout, stderr))
1282    };
1283
1284    match tokio::time::timeout(command_timeout, execution).await {
1285        Ok(result) => {
1286            let (status, stdout, stderr) = result?;
1287            let duration = start.elapsed();
1288            let exit_code = status.code().unwrap_or(-1);
1289            debug!(
1290                "system-ssh command completed on {} (exit={}, duration={}ms)",
1291                config.id,
1292                exit_code,
1293                duration.as_millis()
1294            );
1295            Ok(CommandResult {
1296                exit_code,
1297                stdout,
1298                stderr,
1299                duration_ms: duration.as_millis() as u64,
1300            })
1301        }
1302        Err(_) => {
1303            // `kill_on_drop(true)` ensures the local ssh process is
1304            // SIGKILLed when `child` drops here, so no ssh child can leak
1305            // even if the remote end is wedged. This only owns the local
1306            // SSH child; remote process termination is not guaranteed.
1307            warn!(
1308                "system-ssh command timed out on {} after {:?}",
1309                config.id, command_timeout
1310            );
1311            anyhow::bail!("Command timed out after {:?}", command_timeout);
1312        }
1313    }
1314}
1315
1316fn same_ssh_endpoint(left: &WorkerConfig, right: &WorkerConfig) -> bool {
1317    left.id == right.id
1318        && left.host == right.host
1319        && left.user == right.user
1320        && left.identity_file == right.identity_file
1321        && declared_os(&left.tags) == declared_os(&right.tags)
1322}
1323
1324#[derive(Clone)]
1325struct PooledSshClient {
1326    // Endpoint identity is immutable and does not require the client's lock.
1327    // An old caller can hold that lock across authentication while a new
1328    // endpoint is admitted independently into the pool.
1329    config: WorkerConfig,
1330    client: Arc<RwLock<SshClient>>,
1331}
1332
1333/// Connection pool for managing multiple SSH connections.
1334pub struct SshPool {
1335    /// Pool of active connections.
1336    connections: Arc<RwLock<HashMap<WorkerId, PooledSshClient>>>,
1337    /// Default SSH options.
1338    options: SshOptions,
1339}
1340
1341impl SshPool {
1342    /// Create a new connection pool.
1343    pub fn new(options: SshOptions) -> Self {
1344        Self {
1345            connections: Arc::new(RwLock::new(HashMap::new())),
1346            options,
1347        }
1348    }
1349
1350    /// Get or create a connection to a worker.
1351    ///
1352    /// Validates transport liveness on borrow. The client detects exited or
1353    /// expired local masters, but a running master can still have a broken or
1354    /// stalled remote transport. Probe before reuse and reconnect once under
1355    /// the per-worker write lock if the transport no longer answers.
1356    pub async fn get_or_connect(&self, config: &WorkerConfig) -> Result<Arc<RwLock<SshClient>>> {
1357        self.get_or_connect_probing_within(config, POOL_LIVENESS_PROBE_TIMEOUT)
1358            .await
1359    }
1360
1361    /// [`get_or_connect`](Self::get_or_connect) with the liveness probe bounded
1362    /// by `probe_timeout`.
1363    ///
1364    /// The probe used to run under the client's command timeout (300s for the
1365    /// daemon's pools), and a failed probe is followed by a reconnect and a
1366    /// second probe. A worker whose master was up but whose commands hung (fork
1367    /// exhaustion, D-state) therefore held a caller with a 10-15s budget (health
1368    /// probe, toolchain preflight during selection) for ~10 minutes, and the
1369    /// sequential health loop stalled the whole fleet's state behind it.
1370    async fn get_or_connect_probing_within(
1371        &self,
1372        config: &WorkerConfig,
1373        probe_timeout: Duration,
1374    ) -> Result<Arc<RwLock<SshClient>>> {
1375        let shared_client = self.get_or_create_client_entry(config).await;
1376
1377        // Fast path: reuse only a session that is both present AND verified live.
1378        // The probe needs `&self` only (health_check → execute), so it runs under
1379        // a shared read lock and never blocks other borrowers of this worker.
1380        let reusable = {
1381            let guard = shared_client.read().await;
1382            guard.is_connected()
1383                && guard
1384                    .health_check_with_timeout(probe_timeout)
1385                    .await
1386                    .unwrap_or(false)
1387        };
1388        if reusable {
1389            debug!("Reusing live connection to {}", config.id);
1390            return Ok(shared_client);
1391        }
1392
1393        // Slow path: (re)connect under the per-worker write lock. Re-evaluate the
1394        // state here because another task may have connected/reconnected while we
1395        // waited for the lock, and to distinguish "never connected" (connect)
1396        // from "session exists but the master is dead" (reconnect).
1397        let mut client_guard = shared_client.write().await;
1398        if !client_guard.is_connected() {
1399            client_guard.connect().await?;
1400        } else if !client_guard
1401            .health_check_with_timeout(probe_timeout)
1402            .await
1403            .unwrap_or(false)
1404        {
1405            warn!(
1406                "Pooled SSH connection to {} failed liveness probe; reconnecting",
1407                config.id
1408            );
1409            client_guard.reconnect().await?;
1410        }
1411        // Drop write lock before returning
1412        drop(client_guard);
1413
1414        Ok(shared_client)
1415    }
1416
1417    async fn get_or_create_client_entry(&self, config: &WorkerConfig) -> Arc<RwLock<SshClient>> {
1418        let worker_id = config.id.clone();
1419
1420        loop {
1421            let existing_entry = {
1422                let connections = self.connections.read().await;
1423                connections.get(&worker_id).cloned()
1424            };
1425
1426            if let Some(entry) = existing_entry {
1427                if same_ssh_endpoint(&entry.config, config) {
1428                    return entry.client;
1429                }
1430
1431                let replacement = Arc::new(RwLock::new(SshClient::new(
1432                    config.clone(),
1433                    pooled_client_options(&self.options, config),
1434                )));
1435                let replaced = {
1436                    let mut connections = self.connections.write().await;
1437                    if connections
1438                        .get(&worker_id)
1439                        .is_some_and(|current| Arc::ptr_eq(&current.client, &entry.client))
1440                    {
1441                        connections.insert(
1442                            worker_id.clone(),
1443                            PooledSshClient {
1444                                config: config.clone(),
1445                                client: replacement.clone(),
1446                            },
1447                        );
1448                        true
1449                    } else {
1450                        false
1451                    }
1452                };
1453
1454                if replaced {
1455                    debug!(
1456                        "Replaced SSH connection entry for {} after endpoint config changed",
1457                        worker_id
1458                    );
1459                    return replacement;
1460                }
1461
1462                continue;
1463            }
1464
1465            let new_client = Arc::new(RwLock::new(SshClient::new(
1466                config.clone(),
1467                pooled_client_options(&self.options, config),
1468            )));
1469            let inserted = {
1470                let mut connections = self.connections.write().await;
1471                if connections.contains_key(&worker_id) {
1472                    false
1473                } else {
1474                    connections.insert(
1475                        worker_id.clone(),
1476                        PooledSshClient {
1477                            config: config.clone(),
1478                            client: new_client.clone(),
1479                        },
1480                    );
1481                    true
1482                }
1483            };
1484
1485            if inserted {
1486                return new_client;
1487            }
1488        }
1489    }
1490
1491    /// Close a specific connection.
1492    pub async fn close(&self, worker_id: &WorkerId) -> Result<()> {
1493        let client = {
1494            let mut connections = self.connections.write().await;
1495            connections.remove(worker_id).map(|entry| entry.client)
1496        };
1497
1498        if let Some(client) = client {
1499            let mut client = client.write().await;
1500            client.disconnect().await?;
1501        }
1502
1503        Ok(())
1504    }
1505
1506    /// Close all connections.
1507    pub async fn close_all(&self) -> Result<()> {
1508        let clients: Vec<_> = {
1509            let mut connections = self.connections.write().await;
1510            connections.drain().map(|(_, entry)| entry.client).collect()
1511        };
1512
1513        for client in clients {
1514            let mut client = client.write().await;
1515            if let Err(e) = client.disconnect().await {
1516                error!("Error closing connection: {}", e);
1517            }
1518        }
1519
1520        Ok(())
1521    }
1522
1523    /// Get the number of active connections.
1524    pub async fn active_connections(&self) -> usize {
1525        self.connections.read().await.len()
1526    }
1527
1528    /// Run a single command on a worker over a POOLED (warm, reused) connection,
1529    /// bounded by `command_timeout`.
1530    ///
1531    /// This is the entry point daemon subsystems use instead of the throwaway
1532    /// `SshClient::new().connect()...execute()...disconnect()` dance. It
1533    /// [`get_or_connect`](Self::get_or_connect)s a live master (validating
1534    /// liveness on borrow and reconnecting a dead one), runs exactly one command,
1535    /// and — crucially — does NOT disconnect afterwards, keeping the
1536    /// ControlMaster warm for the next call. That is the whole point: reuse one
1537    /// master per worker instead of spawning (and, under the old code, leaking) a
1538    /// fresh master for every telemetry/health/cleanup poll.
1539    ///
1540    /// The per-command timeout is applied via a temporary [`SshOptions`] override
1541    /// on the pooled client so it does not disturb the pool's shared default
1542    /// (e.g. a long build vs. a short health probe). The connect timeout and
1543    /// control-master/persist settings come from the pool's options — except
1544    /// that Windows workers are ALWAYS non-mux (see [`pooled_client_options`]):
1545    /// Windows OpenSSH cannot multiplex, so a pooled mux session hangs at the
1546    /// command stage.
1547    pub async fn run_with_timeout(
1548        &self,
1549        config: &WorkerConfig,
1550        command: &str,
1551        command_timeout: Duration,
1552    ) -> Result<CommandResult> {
1553        let client = self
1554            .get_or_connect_probing_within(config, command_timeout.min(POOL_LIVENESS_PROBE_TIMEOUT))
1555            .await?;
1556        // Execute under a shared read lock: execute() needs only `&self`, so
1557        // concurrent callers for the same worker can multiplex over the one
1558        // master without serializing on a write lock.
1559        let guard = client.read().await;
1560        guard.execute_with_timeout(command, command_timeout).await
1561    }
1562}
1563
1564impl Default for SshPool {
1565    fn default() -> Self {
1566        Self::new(SshOptions::default())
1567    }
1568}
1569
1570#[cfg(test)]
1571mod tests {
1572    use super::*;
1573    use crate::test_guard;
1574
1575    #[test]
1576    fn test_command_result_success() {
1577        let _guard = test_guard!();
1578        let result = CommandResult {
1579            exit_code: 0,
1580            stdout: "output".to_string(),
1581            stderr: String::new(),
1582            duration_ms: 100,
1583        };
1584        assert!(result.success());
1585
1586        let failed = CommandResult {
1587            exit_code: 1,
1588            stdout: String::new(),
1589            stderr: "error".to_string(),
1590            duration_ms: 50,
1591        };
1592        assert!(!failed.success());
1593    }
1594
1595    #[test]
1596    fn test_ssh_options_default() {
1597        let _guard = test_guard!();
1598        let options = SshOptions::default();
1599        assert_eq!(options.connect_timeout, Duration::from_secs(10));
1600        assert_eq!(options.command_timeout, Duration::from_secs(300));
1601        assert!(options.server_alive_interval.is_none());
1602        assert!(options.control_persist_idle.is_none());
1603        assert!(!options.control_master);
1604    }
1605
1606    #[test]
1607    fn test_ssh_client_creation() {
1608        let _guard = test_guard!();
1609        let config = WorkerConfig {
1610            id: WorkerId::new("test-worker"),
1611            host: "192.168.1.100".to_string(),
1612            user: "ubuntu".to_string(),
1613            identity_file: "~/.ssh/id_rsa".to_string(),
1614            total_slots: 8,
1615            priority: 100,
1616            tags: vec!["rust".to_string()],
1617            tools: Vec::new(),
1618        };
1619
1620        let client = SshClient::new(config.clone(), SshOptions::default());
1621        assert_eq!(client.worker_id().as_str(), "test-worker");
1622        assert!(!client.is_connected());
1623    }
1624
1625    #[test]
1626    fn test_expected_health_check_output_accepts_sentinel_as_last_line() {
1627        let _guard = test_guard!();
1628
1629        assert!(is_expected_health_check_output("ok\n"));
1630        assert!(is_expected_health_check_output("login banner\nok\n"));
1631        assert!(!is_expected_health_check_output(""));
1632        assert!(!is_expected_health_check_output("not ok\n"));
1633        assert!(!is_expected_health_check_output("ok\npost-command noise\n"));
1634    }
1635
1636    fn worker_config(id: &str, host: &str, user: &str, identity_file: &str) -> WorkerConfig {
1637        WorkerConfig {
1638            id: WorkerId::new(id),
1639            host: host.to_string(),
1640            user: user.to_string(),
1641            identity_file: identity_file.to_string(),
1642            total_slots: 8,
1643            priority: 100,
1644            tags: vec!["rust".to_string()],
1645            tools: Vec::new(),
1646        }
1647    }
1648
1649    #[test]
1650    fn test_ssh_client_configured_for_ignores_scheduling_fields() {
1651        let _guard = test_guard!();
1652        let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1653        let client = SshClient::new(config.clone(), SshOptions::default());
1654
1655        let mut scheduling_only_change = config;
1656        scheduling_only_change.total_slots = 16;
1657        scheduling_only_change.priority = 250;
1658        scheduling_only_change.tags = vec!["rust".to_string(), "gpu".to_string()];
1659
1660        assert!(client.is_configured_for(&scheduling_only_change));
1661    }
1662
1663    #[test]
1664    fn test_ssh_client_configured_for_detects_endpoint_changes() {
1665        let _guard = test_guard!();
1666        let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1667        let client = SshClient::new(config, SshOptions::default());
1668
1669        assert!(!client.is_configured_for(&worker_config(
1670            "worker-a",
1671            "192.168.1.101",
1672            "ubuntu",
1673            "~/.ssh/id_rsa",
1674        )));
1675        assert!(!client.is_configured_for(&worker_config(
1676            "worker-a",
1677            "192.168.1.100",
1678            "admin",
1679            "~/.ssh/id_rsa",
1680        )));
1681        assert!(!client.is_configured_for(&worker_config(
1682            "worker-a",
1683            "192.168.1.100",
1684            "ubuntu",
1685            "~/.ssh/other_key",
1686        )));
1687    }
1688
1689    #[tokio::test]
1690    async fn test_ssh_pool_reuses_matching_disconnected_entry() {
1691        let _guard = test_guard!();
1692        let pool = SshPool::default();
1693        let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1694
1695        let first = pool.get_or_create_client_entry(&config).await;
1696        let second = pool.get_or_create_client_entry(&config).await;
1697
1698        assert!(Arc::ptr_eq(&first, &second));
1699        assert_eq!(pool.active_connections().await, 1);
1700    }
1701
1702    #[tokio::test]
1703    async fn test_ssh_pool_reuses_single_entry_and_close_all_drops_to_zero() {
1704        // Leak guard: repeated borrows of the SAME worker must reuse ONE pool
1705        // entry (one warm master, not one-per-call), and close_all() must empty
1706        // the pool. Uses the entry-creation path (get_or_create_client_entry) so
1707        // the test needs no live SSH host; get_or_connect layers only a liveness
1708        // probe + reconnect on top of this same entry map.
1709        let _guard = test_guard!();
1710        let options = SshOptions {
1711            control_master: true,
1712            control_persist_idle: Some(Duration::from_secs(60)),
1713            ..Default::default()
1714        };
1715        let pool = SshPool::new(options);
1716        let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1717
1718        for _ in 0..5 {
1719            let _entry = pool.get_or_create_client_entry(&config).await;
1720        }
1721        assert_eq!(
1722            pool.active_connections().await,
1723            1,
1724            "repeated borrows of one worker must reuse a single pool entry"
1725        );
1726
1727        pool.close_all().await.expect("close_all should succeed");
1728        assert_eq!(
1729            pool.active_connections().await,
1730            0,
1731            "close_all must drop all pooled connections"
1732        );
1733    }
1734
1735    #[test]
1736    fn test_pool_options_map_to_bounded_idle_persist() {
1737        // Leak guard: the pool's mux options must map to a BOUNDED
1738        // ControlPersist=IdleFor(60), never Closed/Forever. control_persist_mode
1739        // is the single decision point configure_builder uses.
1740        let _guard = test_guard!();
1741        assert_eq!(
1742            control_persist_mode(true, Some(Duration::from_secs(60))),
1743            ControlPersistMode::IdleFor(NonZeroUsize::new(60).unwrap()),
1744            "pool mux options must keep a warm master for a bounded idle window"
1745        );
1746    }
1747
1748    // ======================================================================
1749    // Windows workers never get pooled ControlMaster (bd-wgbx9)
1750    // ======================================================================
1751
1752    fn windows_worker_config(id: &str) -> WorkerConfig {
1753        let mut config = worker_config(id, "100.68.2.11", "jeffr", "~/.ssh/surfacebookje_key");
1754        config.tags = vec!["rust".to_string(), crate::types::os_tag("windows")];
1755        config
1756    }
1757
1758    #[test]
1759    fn test_pooled_options_disable_control_master_for_windows_workers() {
1760        // Regression (bd-wgbx9): Windows OpenSSH has no ControlMaster — a
1761        // pooled mux session hangs at the command stage ("Failed to wait for
1762        // command completion") and the worker is falsely marked unreachable
1763        // while fresh one-shot SSH succeeds. Whatever the pool's global
1764        // options say, a Windows worker's pooled client must be non-mux.
1765        let _guard = test_guard!();
1766        let pool_mux = SshOptions {
1767            control_master: true,
1768            control_persist_idle: Some(Duration::from_secs(60)),
1769            ..Default::default()
1770        };
1771        let config = windows_worker_config("wsurf");
1772
1773        let effective = pooled_client_options(&pool_mux, &config);
1774        assert!(
1775            !effective.control_master,
1776            "a Windows worker must never get a pooled ControlMaster"
1777        );
1778        // The override must be surgical: everything else stays verbatim.
1779        // control_persist_idle is left as-is because it is inert once mux is
1780        // off (control_persist_mode only reads it when control_master is true).
1781        assert_eq!(effective.connect_timeout, pool_mux.connect_timeout);
1782        assert_eq!(effective.command_timeout, pool_mux.command_timeout);
1783        assert_eq!(
1784            effective.server_alive_interval,
1785            pool_mux.server_alive_interval
1786        );
1787        assert_eq!(
1788            effective.control_persist_idle,
1789            pool_mux.control_persist_idle
1790        );
1791        assert_eq!(effective.known_hosts, pool_mux.known_hosts);
1792
1793        // A pool that already disabled mux stays disabled.
1794        let pool_plain = SshOptions {
1795            control_master: false,
1796            ..pool_mux
1797        };
1798        assert!(!pooled_client_options(&pool_plain, &config).control_master);
1799
1800        // declared_os normalizes case: `os = "Windows"` in workers.toml is the
1801        // same reserved tag and must disable mux too.
1802        let mut mixed_case = windows_worker_config("wsurf-2");
1803        mixed_case.tags = vec![crate::types::os_tag("Windows")];
1804        assert!(!pooled_client_options(&pool_mux, &mixed_case).control_master);
1805    }
1806
1807    #[test]
1808    fn test_pooled_options_keep_pool_options_verbatim_for_non_windows_workers() {
1809        // Linux/unlabelled workers keep the pool's options VERBATIM — the
1810        // override must not degrade warm-master reuse for the rest of the
1811        // fleet, under either a mux or a hypothetical non-mux pool.
1812        let _guard = test_guard!();
1813        let pool_mux = SshOptions {
1814            control_master: true,
1815            control_persist_idle: Some(Duration::from_secs(60)),
1816            ..Default::default()
1817        };
1818        let linux = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1819        let unlabelled = worker_config("contabo-b", "1.2.3.5", "root", "~/.ssh/id_rsa");
1820
1821        for config in [&linux, &unlabelled] {
1822            let effective = pooled_client_options(&pool_mux, config);
1823            assert!(
1824                effective.control_master,
1825                "non-Windows workers keep the pool's mux setting"
1826            );
1827            assert_eq!(
1828                effective.control_persist_idle,
1829                pool_mux.control_persist_idle
1830            );
1831        }
1832
1833        let pool_plain = SshOptions {
1834            control_master: false,
1835            ..pool_mux
1836        };
1837        for config in [&linux, &unlabelled] {
1838            assert!(!pooled_client_options(&pool_plain, config).control_master);
1839        }
1840    }
1841
1842    #[tokio::test]
1843    async fn test_pool_entry_for_windows_worker_is_never_mux() {
1844        // End-to-end at the pool layer: the daemon's health pool and shared
1845        // build/telemetry pool both construct per-worker clients via
1846        // get_or_create_client_entry with pool-global mux options — the
1847        // Windows entry must come out non-mux. Uses the entry-creation path
1848        // (like the other pool tests) so no live SSH host is needed.
1849        let _guard = test_guard!();
1850        let pool = SshPool::new(SshOptions {
1851            control_master: true,
1852            control_persist_idle: Some(Duration::from_secs(60)),
1853            ..Default::default()
1854        });
1855        let config = windows_worker_config("wsurf");
1856
1857        let entry = pool.get_or_create_client_entry(&config).await;
1858        let guard = entry.read().await;
1859        assert!(
1860            !guard.options.control_master,
1861            "pooled entry for a Windows worker must be non-mux"
1862        );
1863        // The inert companion setting is untouched, proving the override was
1864        // surgical rather than a blanket options reset.
1865        assert_eq!(
1866            guard.options.control_persist_idle,
1867            Some(Duration::from_secs(60))
1868        );
1869    }
1870
1871    #[tokio::test]
1872    async fn test_pool_entry_for_linux_worker_keeps_mux() {
1873        let _guard = test_guard!();
1874        let pool = SshPool::new(SshOptions {
1875            control_master: true,
1876            control_persist_idle: Some(Duration::from_secs(60)),
1877            ..Default::default()
1878        });
1879        let config = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1880
1881        let entry = pool.get_or_create_client_entry(&config).await;
1882        let guard = entry.read().await;
1883        assert!(
1884            guard.options.control_master,
1885            "pooled entry for a Linux worker must keep the pool's mux setting"
1886        );
1887    }
1888
1889    #[tokio::test]
1890    async fn test_ssh_pool_replaces_stale_entry_when_endpoint_changes() {
1891        let _guard = test_guard!();
1892        let pool = SshPool::default();
1893        let old_config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1894        let new_config = worker_config("worker-a", "192.168.1.101", "admin", "~/.ssh/new_key");
1895
1896        let stale = pool.get_or_create_client_entry(&old_config).await;
1897        let replacement = pool.get_or_create_client_entry(&new_config).await;
1898
1899        assert!(!Arc::ptr_eq(&stale, &replacement));
1900        assert_eq!(pool.active_connections().await, 1);
1901
1902        let replacement_guard = replacement.read().await;
1903        assert!(replacement_guard.is_configured_for(&new_config));
1904    }
1905
1906    #[tokio::test]
1907    async fn test_ssh_pool_retarget_does_not_wait_for_old_authentication_lock() {
1908        let _guard = test_guard!();
1909        let pool = SshPool::default();
1910        let old_config = worker_config("worker-a", "old-endpoint", "builder", "~/.ssh/id_rsa");
1911        let new_config = worker_config("worker-a", "new-endpoint", "builder", "~/.ssh/id_rsa");
1912        let old = pool.get_or_create_client_entry(&old_config).await;
1913        let authenticating = old.write().await;
1914        let unchanged = tokio::time::timeout(
1915            Duration::from_millis(100),
1916            pool.get_or_create_client_entry(&old_config),
1917        )
1918        .await
1919        .expect("immutable identity lookup must not wait on authentication");
1920        assert!(Arc::ptr_eq(&old, &unchanged));
1921        let replacement = tokio::time::timeout(
1922            Duration::from_millis(100),
1923            pool.get_or_create_client_entry(&new_config),
1924        )
1925        .await
1926        .expect("new endpoint must not wait on the old endpoint's authentication");
1927        assert!(!Arc::ptr_eq(&old, &replacement));
1928        assert_eq!(pool.active_connections().await, 1);
1929        assert!(replacement.read().await.is_configured_for(&new_config));
1930        drop(authenticating);
1931    }
1932
1933    #[tokio::test]
1934    async fn test_ssh_pool_os_transition_replaces_transport_but_case_and_scheduling_reuse() {
1935        let _guard = test_guard!();
1936        let pool = SshPool::new(SshOptions {
1937            control_master: true,
1938            ..Default::default()
1939        });
1940        let mut config = worker_config("worker-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
1941        let unlabelled = pool.get_or_create_client_entry(&config).await;
1942        config.tags.push(crate::types::os_tag("windows"));
1943        let windows = pool.get_or_create_client_entry(&config).await;
1944        assert!(!Arc::ptr_eq(&unlabelled, &windows));
1945        {
1946            let client = windows.read().await;
1947            assert!(prefers_system_ssh(&client.config));
1948            assert!(!client.options.control_master);
1949        }
1950
1951        config.tags = vec!["gpu".into(), "os: WiNdOwS ".into()];
1952        config.total_slots += 1;
1953        config.priority += 1;
1954        let reused = pool.get_or_create_client_entry(&config).await;
1955        assert!(Arc::ptr_eq(&windows, &reused));
1956
1957        config.tags = vec![crate::types::os_tag("linux")];
1958        let linux = pool.get_or_create_client_entry(&config).await;
1959        assert!(!Arc::ptr_eq(&windows, &linux));
1960        {
1961            let client = linux.read().await;
1962            assert!(!prefers_system_ssh(&client.config));
1963            assert!(client.options.control_master);
1964        }
1965        // Existing borrowers retain their original transport while the pool
1966        // exposes only the replacement to subsequent requests.
1967        assert!(prefers_system_ssh(&windows.read().await.config));
1968        assert!(unlabelled.read().await.options.control_master);
1969        assert_eq!(pool.active_connections().await, 1);
1970
1971        config.tags.clear();
1972        let removed = pool.get_or_create_client_entry(&config).await;
1973        assert!(!Arc::ptr_eq(&linux, &removed));
1974        assert!(removed.read().await.options.control_master);
1975        assert_eq!(pool.active_connections().await, 1);
1976    }
1977
1978    #[tokio::test]
1979    async fn test_health_check_reports_not_alive_without_session() {
1980        // get_or_connect's liveness probe relies on health_check() returning a
1981        // non-true value (without erroring) for a client that has no live
1982        // session, so a dead/empty pooled entry is reconnected rather than
1983        // falsely handed back as "reused". execute() errors "Not connected",
1984        // which health_check() maps to Ok(false).
1985        let _guard = test_guard!();
1986        let config = worker_config("worker-a", "192.168.1.100", "ubuntu", "~/.ssh/id_rsa");
1987        let client = SshClient::new(config, SshOptions::default());
1988        assert!(!client.is_connected());
1989        let alive = client
1990            .health_check()
1991            .await
1992            .expect("health_check maps execution errors to Ok(false), never Err");
1993        assert!(!alive, "a client with no session must not report as alive");
1994    }
1995
1996    #[test]
1997    fn test_build_env_prefix_quotes_and_rejects() {
1998        let _guard = test_guard!();
1999        let mut env = HashMap::new();
2000        env.insert("RUSTFLAGS".to_string(), "-C target-cpu=native".to_string());
2001        env.insert("QUOTED".to_string(), "a'b".to_string());
2002        env.insert("BADVAL".to_string(), "line1\nline2".to_string());
2003
2004        let allowlist = vec![
2005            "RUSTFLAGS".to_string(),
2006            "QUOTED".to_string(),
2007            "MISSING".to_string(),
2008            "BADVAL".to_string(),
2009            "BAD=KEY".to_string(),
2010        ];
2011
2012        let prefix = build_env_prefix(&allowlist, |key| env.get(key).cloned());
2013
2014        assert!(prefix.prefix.contains("RUSTFLAGS='-C target-cpu=native'"));
2015        // shell_escape uses '\'' style (end string, escaped quote, start string)
2016        assert!(prefix.prefix.contains("QUOTED='a'\\''b'"));
2017        assert!(!prefix.prefix.contains("MISSING="));
2018        assert!(!prefix.prefix.contains("BADVAL="));
2019        assert!(prefix.rejected.contains(&"BADVAL".to_string()));
2020        assert!(prefix.rejected.contains(&"BAD=KEY".to_string()));
2021        assert_eq!(
2022            prefix.applied,
2023            vec!["RUSTFLAGS".to_string(), "QUOTED".to_string()]
2024        );
2025    }
2026
2027    // ==========================================================================
2028    // Proptest: SSH command escaping with special chars (bd-2elj)
2029    // ==========================================================================
2030
2031    mod proptest_ssh_escaping {
2032        use super::*;
2033        use proptest::prelude::*;
2034        use std::collections::HashMap;
2035
2036        proptest! {
2037            #![proptest_config(ProptestConfig::with_cases(1000))]
2038
2039            // Test 1: is_valid_env_key never panics on arbitrary strings
2040            #[test]
2041            fn test_is_valid_env_key_no_panic(s in ".*") {
2042        let _guard = test_guard!();
2043                let _ = is_valid_env_key(&s);
2044            }
2045
2046            // Test 2: Valid env keys start with letter/_ and contain only alphanum/_
2047            #[test]
2048            fn test_is_valid_env_key_accepts_valid(
2049                first in "[a-zA-Z_]",
2050                rest in "[a-zA-Z0-9_]{0,50}"
2051            ) {
2052        let _guard = test_guard!();
2053                let key = format!("{}{}", first, rest);
2054                prop_assert!(is_valid_env_key(&key), "Should accept valid key: {}", key);
2055            }
2056
2057            // Test 3: Env keys starting with digit are invalid
2058            #[test]
2059            fn test_is_valid_env_key_rejects_digit_start(
2060                digit in "[0-9]",
2061                rest in "[a-zA-Z0-9_]{0,20}"
2062            ) {
2063        let _guard = test_guard!();
2064                let key = format!("{}{}", digit, rest);
2065                prop_assert!(!is_valid_env_key(&key), "Should reject digit-start key: {}", key);
2066            }
2067
2068            // Test 4: shell_escape_value never panics on arbitrary strings
2069            #[test]
2070            fn test_shell_escape_value_no_panic(s in ".*") {
2071        let _guard = test_guard!();
2072                let _ = shell_escape_value(&s);
2073            }
2074
2075            // Test 5: shell_escape_value rejects newlines/carriage returns/NUL
2076            #[test]
2077            fn test_shell_escape_value_rejects_unsafe(
2078                prefix in "[a-zA-Z0-9 ]{0,10}",
2079                bad_char in "[\n\r\0]",
2080                suffix in "[a-zA-Z0-9 ]{0,10}"
2081            ) {
2082        let _guard = test_guard!();
2083                let value = format!("{}{}{}", prefix, bad_char, suffix);
2084                prop_assert!(shell_escape_value(&value).is_none(),
2085                    "Should reject value with unsafe char: {:?}", value);
2086            }
2087
2088            // Test 6: shell_escape_value handles safe values
2089            #[test]
2090            fn test_shell_escape_value_accepts_safe(s in "[a-zA-Z0-9 !@#$%^&*()_+=\\-\\[\\]{}|;:,./<>?]{0,100}") {
2091                // These don't contain \n, \r, or \0
2092                let result = shell_escape_value(&s);
2093                prop_assert!(result.is_some(), "Should accept safe value: {:?}", s);
2094
2095                // shell_escape only quotes values that need it (contain special chars)
2096                // Simple alphanumeric strings may be returned unquoted
2097                let escaped = match result {
2098                    Some(escaped) => escaped,
2099                    None => {
2100                        prop_assert!(false, "Should accept safe value: {:?}", s);
2101                        String::new()
2102                    }
2103                };
2104                if s.chars().any(|c| !c.is_ascii_alphanumeric() && c != '_') {
2105                    // Values with special chars should be quoted
2106                    prop_assert!(escaped.starts_with('\'') || escaped.contains('\''),
2107                        "Value with special chars should be quoted: {:?} -> {:?}", s, escaped);
2108                }
2109            }
2110
2111            // Test 7: shell_escape_value properly escapes single quotes
2112            #[test]
2113            fn test_shell_escape_value_escapes_quotes(
2114                prefix in "[a-zA-Z0-9]{0,10}",
2115                suffix in "[a-zA-Z0-9]{0,10}"
2116            ) {
2117        let _guard = test_guard!();
2118                let value = format!("{}'{}", prefix, suffix);
2119                let result = shell_escape_value(&value);
2120                prop_assert!(result.is_some());
2121
2122                let escaped = match result {
2123                    Some(escaped) => escaped,
2124                    None => {
2125                        prop_assert!(false, "Should escape single quote: {}", value);
2126                        String::new()
2127                    }
2128                };
2129                // shell_escape uses '\'' style (end string, escaped quote, start string)
2130                prop_assert!(escaped.contains("'\\''"),
2131                    "Should escape single quote: {} -> {}", value, escaped);
2132            }
2133
2134            // Test 8: build_env_prefix never panics
2135            #[test]
2136            fn test_build_env_prefix_no_panic(
2137                keys in prop::collection::vec("[a-zA-Z_][a-zA-Z0-9_]{0,10}", 0..10),
2138                values in prop::collection::vec(".*", 0..10)
2139            ) {
2140                let mut env = HashMap::new();
2141                for (i, key) in keys.iter().enumerate() {
2142                    if let Some(val) = values.get(i) {
2143                        env.insert(key.clone(), val.clone());
2144                    }
2145                }
2146
2147                let allowlist: Vec<String> = keys;
2148                let _ = build_env_prefix(&allowlist, |k| env.get(k).cloned());
2149            }
2150
2151            // Test 9: build_env_prefix rejects invalid keys (non-empty after trim)
2152            #[test]
2153            fn test_build_env_prefix_rejects_invalid_keys(
2154                // Generate keys that are invalid even after trimming
2155                invalid_key in "[0-9][a-zA-Z0-9_]{0,10}"  // Starts with digit
2156            ) {
2157        let _guard = test_guard!();
2158                let mut env = HashMap::new();
2159                env.insert(invalid_key.clone(), "value".to_string());
2160
2161                let allowlist = vec![invalid_key.clone()];
2162                let prefix = build_env_prefix(&allowlist, |k| env.get(k).cloned());
2163
2164                // Key should be rejected since it starts with a digit
2165                prop_assert!(!is_valid_env_key(&invalid_key),
2166                    "Key should be invalid: {}", invalid_key);
2167                prop_assert!(prefix.rejected.contains(&invalid_key),
2168                    "Should reject invalid key: {}", invalid_key);
2169                prop_assert!(prefix.prefix.is_empty());
2170            }
2171
2172            // Test 10: build_env_prefix handles missing values gracefully
2173            #[test]
2174            fn test_build_env_prefix_missing_values(
2175                keys in prop::collection::vec("[A-Z_][A-Z0-9_]{0,10}", 1..5)
2176            ) {
2177                // Empty env - all keys missing
2178                let env: HashMap<String, String> = HashMap::new();
2179                let prefix = build_env_prefix(&keys, |k| env.get(k).cloned());
2180
2181                // Should be empty prefix since no values found
2182                prop_assert!(prefix.prefix.is_empty(), "Should be empty when no values");
2183                prop_assert!(prefix.applied.is_empty());
2184                // Missing values don't count as rejected
2185                prop_assert!(prefix.rejected.is_empty());
2186            }
2187        }
2188
2189        // Targeted edge case tests
2190        #[test]
2191        fn test_shell_escape_edge_cases() {
2192            let _guard = test_guard!();
2193            // Empty string
2194            let result = shell_escape_value("");
2195            assert_eq!(result, Some("''".to_string()));
2196
2197            // Just single quote - shell_escape uses '\'' style (end string, escaped quote, start string)
2198            let result = shell_escape_value("'");
2199            assert_eq!(result, Some("''\\'''".to_string()));
2200
2201            // Multiple single quotes
2202            let result = shell_escape_value("'''");
2203            // shell_escape uses '\'' style for each single quote
2204            assert_eq!(
2205                result
2206                    .as_deref()
2207                    .map(|escaped| escaped.matches("'\\''").count()),
2208                Some(3)
2209            );
2210
2211            // Unicode
2212            let result = shell_escape_value("日本語");
2213            assert!(result.is_some());
2214
2215            // Emoji
2216            let result = shell_escape_value("🔥🚀");
2217            assert!(result.is_some());
2218
2219            // Mixed quotes and special chars
2220            let result = shell_escape_value("it's a \"test\" with $vars");
2221            assert!(result.is_some());
2222        }
2223
2224        #[test]
2225        fn test_is_valid_env_key_edge_cases() {
2226            let _guard = test_guard!();
2227            // Empty
2228            assert!(!is_valid_env_key(""));
2229
2230            // Single underscore
2231            assert!(is_valid_env_key("_"));
2232
2233            // Single letter
2234            assert!(is_valid_env_key("A"));
2235
2236            // Typical env vars
2237            assert!(is_valid_env_key("PATH"));
2238            assert!(is_valid_env_key("HOME"));
2239            assert!(is_valid_env_key("RUSTFLAGS"));
2240            assert!(is_valid_env_key("CC"));
2241            assert!(is_valid_env_key("_PRIVATE"));
2242            assert!(is_valid_env_key("MY_VAR_123"));
2243
2244            // Invalid: starts with number
2245            assert!(!is_valid_env_key("1VAR"));
2246            assert!(!is_valid_env_key("123"));
2247
2248            // Invalid: contains special chars
2249            assert!(!is_valid_env_key("MY-VAR"));
2250            assert!(!is_valid_env_key("MY.VAR"));
2251            assert!(!is_valid_env_key("MY VAR"));
2252            assert!(!is_valid_env_key("MY=VAR"));
2253
2254            // Invalid: Unicode
2255            assert!(!is_valid_env_key("日本語"));
2256            assert!(!is_valid_env_key("VAR🔥"));
2257        }
2258
2259        #[test]
2260        fn test_build_env_prefix_integration() {
2261            let _guard = test_guard!();
2262            // Complex scenario with mixed valid/invalid
2263            let mut env = HashMap::new();
2264            env.insert("VALID".to_string(), "simple".to_string());
2265            env.insert("WITH_QUOTE".to_string(), "it's here".to_string());
2266            env.insert("NEWLINE".to_string(), "line1\nline2".to_string());
2267            env.insert("UNICODE".to_string(), "日本語".to_string());
2268            env.insert("EMPTY".to_string(), String::new());
2269            env.insert("123INVALID".to_string(), "value".to_string());
2270
2271            let allowlist = vec![
2272                "VALID".to_string(),
2273                "WITH_QUOTE".to_string(),
2274                "NEWLINE".to_string(),
2275                "UNICODE".to_string(),
2276                "EMPTY".to_string(),
2277                "123INVALID".to_string(),
2278                "MISSING".to_string(),
2279            ];
2280
2281            let prefix = build_env_prefix(&allowlist, |k| env.get(k).cloned());
2282
2283            // VALID should be applied
2284            assert!(prefix.applied.contains(&"VALID".to_string()));
2285            // shell_escape doesn't quote simple alphanumeric strings
2286            assert!(prefix.prefix.contains("VALID=simple"));
2287
2288            // WITH_QUOTE should be applied with escaped quote
2289            assert!(prefix.applied.contains(&"WITH_QUOTE".to_string()));
2290
2291            // NEWLINE should be rejected (unsafe value)
2292            assert!(prefix.rejected.contains(&"NEWLINE".to_string()));
2293
2294            // UNICODE should be applied (safe unicode)
2295            assert!(prefix.applied.contains(&"UNICODE".to_string()));
2296
2297            // EMPTY should be applied
2298            assert!(prefix.applied.contains(&"EMPTY".to_string()));
2299
2300            // 123INVALID should be rejected (invalid key)
2301            assert!(prefix.rejected.contains(&"123INVALID".to_string()));
2302
2303            // MISSING should not appear in either list (not found = silently ignored)
2304            assert!(!prefix.applied.contains(&"MISSING".to_string()));
2305            assert!(!prefix.rejected.contains(&"MISSING".to_string()));
2306        }
2307
2308        #[test]
2309        fn test_shell_escape_roundtrip_safety() {
2310            let _guard = test_guard!();
2311            // Values that when escaped and passed through shell should reconstruct original
2312            let test_values = [
2313                "simple",
2314                "with spaces",
2315                "with\ttab",
2316                "special!@#$%^&*()",
2317                "quoted\"value",
2318                "path/to/file",
2319                "-flag",
2320                "--long-flag=value",
2321                "",
2322            ];
2323
2324            for value in &test_values {
2325                let escaped = shell_escape_value(value);
2326                assert!(escaped.is_some(), "Should escape: {:?}", value);
2327            }
2328        }
2329    }
2330
2331    // ======================================================================
2332    // System-ssh fallback for Windows workers (bd-kzy2x)
2333    // ======================================================================
2334
2335    fn windows_worker(id: &str) -> WorkerConfig {
2336        let mut config = worker_config(id, "100.68.2.11", "jeffr", "~/.ssh/surfacebookje_key");
2337        config.tags = vec!["rust".to_string(), crate::types::os_tag("windows")];
2338        config
2339    }
2340
2341    #[test]
2342    fn test_prefers_system_ssh_for_windows_tag() {
2343        // The dispatch key in lockstep with `pooled_client_options`:
2344        // declared_os == "windows" -> system-ssh fallback.
2345        let _guard = test_guard!();
2346        let cfg = windows_worker("wsurf");
2347        assert!(prefers_system_ssh(&cfg));
2348    }
2349
2350    #[test]
2351    fn test_prefers_system_ssh_case_normalized() {
2352        // declared_os lower-cases the tag's OS, so `os = "Windows"` (mixed
2353        // case) and `os = "WINDOWS"` (all caps) both route to the fallback.
2354        let _guard = test_guard!();
2355        let mut cfg = windows_worker("wsurf-mixed");
2356        cfg.tags = vec![crate::types::os_tag("Windows")];
2357        assert!(prefers_system_ssh(&cfg));
2358
2359        let mut upper = windows_worker("wsurf-upper");
2360        upper.tags = vec![crate::types::os_tag("WINDOWS")];
2361        assert!(prefers_system_ssh(&upper));
2362    }
2363
2364    #[test]
2365    fn test_prefers_system_ssh_false_for_linux() {
2366        // `os:linux` and unlabelled workers do NOT trigger the fallback —
2367        // they keep the openssh-crate path verbatim.
2368        let _guard = test_guard!();
2369        let mut linux = worker_config("contabo-a", "1.2.3.4", "root", "~/.ssh/id_rsa");
2370        linux.tags = vec!["rust".to_string(), crate::types::os_tag("linux")];
2371        assert!(!prefers_system_ssh(&linux));
2372
2373        let unlabelled = worker_config("contabo-b", "1.2.3.5", "root", "~/.ssh/id_rsa");
2374        assert!(!prefers_system_ssh(&unlabelled));
2375    }
2376
2377    #[test]
2378    fn test_system_ssh_argv_basic_windows() {
2379        // Baseline Windows worker (no port today — WorkerConfig has no
2380        // `port` field; non-default ports travel as `host:port`). The argv
2381        // ends with the destination then a fixed reader, never has a `-p`
2382        // flag, and carries the keepalive when non-zero.
2383        let _guard = test_guard!();
2384        let cfg = windows_worker("wsurf");
2385        let command = remote_shell_command(&cfg, "uname -a");
2386        let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(2));
2387        let s: Vec<String> = argv
2388            .iter()
2389            .map(|o| o.to_string_lossy().into_owned())
2390            .collect();
2391
2392        assert_eq!(s[0], "ssh", "argv[0] is the program name");
2393        assert_eq!(s[1], "-i");
2394        // shellexpand::tilde expands `~`; on this machine it goes to a
2395        // $HOME-anchored path, so the suffix is the right invariant.
2396        assert!(
2397            s[2].ends_with(".ssh/surfacebookje_key"),
2398            "identity file is the tilde-expanded path: {}",
2399            s[2]
2400        );
2401        // Pair -o/value: BatchMode=yes, ConnectTimeout=8,
2402        // StrictHostKeyChecking=accept-new, ServerAliveInterval=2.
2403        assert!(s.contains(&"-o".to_string()));
2404        assert!(s.contains(&"BatchMode=yes".to_string()));
2405        assert!(s.contains(&"ConnectTimeout=8".to_string()));
2406        assert!(s.contains(&"StrictHostKeyChecking=accept-new".to_string()));
2407        assert!(s.contains(&"ServerAliveInterval=2".to_string()));
2408        // No -p flag — WorkerConfig has no port.
2409        assert!(!s.iter().any(|arg| arg == "-p"));
2410        // The script never reaches the account's cmd.exe command line.
2411        assert!(s.contains(&"jeffr@100.68.2.11".to_string()));
2412        assert_eq!(s[s.len() - 1], "sh -s", "Windows gets only a fixed reader");
2413    }
2414
2415    #[test]
2416    fn test_system_ssh_script_transport_keeps_posix_argv_and_exact_stdin_reader() {
2417        let _guard = test_guard!();
2418        let linux = worker_config("linux", "1.2.3.4", "root", "~/.ssh/id_rsa");
2419        let windows = windows_worker("windows");
2420        let script = "printf '%s\\n' '%PATH% !VAR! & | < > ^ $()'\nprintf done";
2421        let prepared = remote_shell_command(&linux, script);
2422        let argv = system_ssh_argv(&linux, &prepared, Duration::ZERO);
2423        assert_eq!(argv.last().unwrap(), script);
2424        assert!(prepared.stdin_script.is_none());
2425        for command in [script, "sh -s", "", "sh -s\n"] {
2426            let prepared = remote_shell_command(&windows, command);
2427            let argv = system_ssh_argv(&windows, &prepared, Duration::ZERO);
2428            assert_eq!(argv.last().unwrap(), "sh -s");
2429            assert_eq!(prepared.stdin_script.is_none(), command == "sh -s");
2430        }
2431    }
2432
2433    #[tokio::test]
2434    async fn test_system_ssh_stdin_roundtrip_matches_null_stdin_shell() {
2435        use std::process::Stdio;
2436
2437        let _guard = test_guard!();
2438        let config = windows_worker("windows");
2439        // These are real POSIX shell oracles for the stdin payload; native
2440        // Windows SSH dispatch is qualified separately on a Windows worker.
2441        for script in [
2442            "",
2443            "printf '%s\\n' \"apostrophe's\" 'double\"quote' '%PATH% !VAR! ^ & | < >' '$() `x` \\'",
2444            "cat\nprintf 'AFTER\\n'\nread value || printf 'EMPTY\\n'",
2445            "cat <<'END'\nheredoc ' \" % ! $() \\\nEND\nprintf 'done\\n'",
2446            "printf '%s\\n' 'first\nsecond' # trailing comment",
2447            "printf 'before\\n'; printf 'diagnostic\\n' >&2; exit 37",
2448            "# comment ending with a backslash \\",
2449        ] {
2450            let expected = tokio::process::Command::new("sh")
2451                .args(["-c", script])
2452                .stdin(Stdio::null())
2453                .output()
2454                .await
2455                .unwrap();
2456            let mut child = tokio::process::Command::new("sh")
2457                .arg("-s")
2458                .stdin(Stdio::piped())
2459                .stdout(Stdio::piped())
2460                .stderr(Stdio::piped())
2461                .kill_on_drop(true)
2462                .spawn()
2463                .unwrap();
2464            let stdin = child.stdin.take().unwrap();
2465            let payload = remote_shell_command(&config, script).stdin_script.unwrap();
2466            let (written, observed) = tokio::time::timeout(Duration::from_secs(5), async {
2467                tokio::join!(
2468                    write_ssh_command_stdin(stdin, &payload),
2469                    child.wait_with_output()
2470                )
2471            })
2472            .await
2473            .unwrap();
2474            written.unwrap();
2475            let observed = observed.unwrap();
2476            assert_eq!(observed.status.code(), expected.status.code(), "{script:?}");
2477            assert_eq!(observed.stdout, expected.stdout, "{script:?}");
2478            assert_eq!(observed.stderr, expected.stderr, "{script:?}");
2479        }
2480    }
2481
2482    #[tokio::test]
2483    async fn test_system_ssh_split_stdin_never_becomes_command_input() {
2484        use std::process::Stdio;
2485        use tokio::io::AsyncWriteExt;
2486
2487        let _guard = test_guard!();
2488        let script = "cat\nprintf 'AFTER\\n'\nread value || printf 'EMPTY\\n'";
2489        let payload = remote_shell_command(&windows_worker("windows"), script)
2490            .stdin_script
2491            .unwrap();
2492        let split = payload.find("cat\n").unwrap() + "cat\n".len();
2493        let mut child = tokio::process::Command::new("sh")
2494            .arg("-s")
2495            .stdin(Stdio::piped())
2496            .stdout(Stdio::piped())
2497            .stderr(Stdio::piped())
2498            .kill_on_drop(true)
2499            .spawn()
2500            .unwrap();
2501        let mut stdin = child.stdin.take().unwrap();
2502        let writer = async {
2503            stdin.write_all(&payload.as_bytes()[..split]).await.unwrap();
2504            tokio::time::sleep(Duration::from_millis(100)).await;
2505            stdin.write_all(&payload.as_bytes()[split..]).await.unwrap();
2506            stdin.shutdown().await.unwrap();
2507        };
2508        let ((), output) = tokio::time::timeout(Duration::from_secs(5), async {
2509            tokio::join!(writer, child.wait_with_output())
2510        })
2511        .await
2512        .unwrap();
2513        let output = output.unwrap();
2514        assert!(output.status.success(), "{:?}", output.stderr);
2515        assert_eq!(output.stdout, b"AFTER\nEMPTY\n");
2516        assert!(output.stderr.is_empty());
2517    }
2518
2519    #[tokio::test]
2520    async fn test_system_ssh_closed_stdin_preserves_early_exit_diagnostic() {
2521        use std::process::Stdio;
2522
2523        let _guard = test_guard!();
2524        let mut child = tokio::process::Command::new("sh")
2525            .args(["-c", "printf 'remote failure\\n' >&2; exit 37"])
2526            .stdin(Stdio::piped())
2527            .stdout(Stdio::piped())
2528            .stderr(Stdio::piped())
2529            .kill_on_drop(true)
2530            .spawn()
2531            .unwrap();
2532        let stdin = child.stdin.take().unwrap();
2533        // Force the pipe closed before writing; a tiny concurrent payload
2534        // could fit in the buffer and accidentally miss BrokenPipe coverage.
2535        child.wait().await.unwrap();
2536        write_ssh_command_stdin(stdin, &"x".repeat(128 * 1024))
2537            .await
2538            .unwrap();
2539        let output = child.wait_with_output().await.unwrap();
2540        assert_eq!(output.status.code(), Some(37));
2541        assert_eq!(output.stderr, b"remote failure\n");
2542        assert!(output.stdout.is_empty());
2543    }
2544
2545    #[test]
2546    fn test_system_ssh_argv_zero_server_alive_omits_keepalive() {
2547        // `server_alive_interval = 0` must NOT emit `-o ServerAliveInterval=0`
2548        // (which OpenSSH treats as "disable keepalives" — fine, but the spec
2549        // calls for omitting the flag entirely at 0).
2550        let _guard = test_guard!();
2551        let cfg = windows_worker("wsurf");
2552        let command = remote_shell_command(&cfg, "echo hi");
2553        let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(0));
2554        let s: Vec<String> = argv
2555            .iter()
2556            .map(|o| o.to_string_lossy().into_owned())
2557            .collect();
2558        assert!(
2559            !s.iter().any(|a| a.starts_with("ServerAliveInterval=")),
2560            "no ServerAliveInterval flag when interval is 0; got {:?}",
2561            s
2562        );
2563    }
2564
2565    #[test]
2566    fn test_system_ssh_argv_long_keepalive_preserved() {
2567        // Non-default keepalive values must be preserved verbatim (no
2568        // truncation, no off-by-one). Use 60s — the same value the openssh
2569        // pool applies for warm masters.
2570        let _guard = test_guard!();
2571        let cfg = windows_worker("wsurf");
2572        let command = remote_shell_command(&cfg, "true");
2573        let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(60));
2574        let s: Vec<String> = argv
2575            .iter()
2576            .map(|o| o.to_string_lossy().into_owned())
2577            .collect();
2578        assert!(s.contains(&"ServerAliveInterval=60".to_string()));
2579    }
2580
2581    #[test]
2582    fn test_system_ssh_argv_uses_tilde_expanded_identity() {
2583        // `shellexpand::tilde` is invoked on the identity_file path so an
2584        // operator can write `~/.ssh/foo` in workers.toml and have ssh
2585        // see the expanded path. Compare the raw `~`-prefixed form is
2586        // NOT in the argv.
2587        let _guard = test_guard!();
2588        let mut cfg = windows_worker("wsurf");
2589        cfg.identity_file = "~/.ssh/operator_key".to_string();
2590        let command = remote_shell_command(&cfg, "true");
2591        let argv = system_ssh_argv(&cfg, &command, Duration::from_secs(2));
2592        let s: Vec<String> = argv
2593            .iter()
2594            .map(|o| o.to_string_lossy().into_owned())
2595            .collect();
2596        assert!(
2597            !s.contains(&"~/.ssh/operator_key".to_string()),
2598            "identity file must be tilde-expanded, got {:?}",
2599            s
2600        );
2601        assert!(
2602            s.contains(&"~/.ssh/operator_key".to_string())
2603                || s.iter().any(|a| a.ends_with(".ssh/operator_key")),
2604            "tilde-expanded path appears in argv: {:?}",
2605            s
2606        );
2607    }
2608
2609    #[test]
2610    fn test_identities_only_only_when_identity_file_exists() {
2611        // bd-ebszo: an existing `-i` key must be the ONLY key offered (no agent
2612        // keys first); a missing key file keeps agent fallback working.
2613        let _guard = test_guard!();
2614        let dir = tempfile::tempdir().expect("tempdir");
2615        let key = dir.path().join("worker_key");
2616        std::fs::write(&key, b"not a real key").expect("write key");
2617        let key = key.to_string_lossy().into_owned();
2618        assert_eq!(
2619            identities_only_args(&key),
2620            Some(["-o", "IdentitiesOnly=yes"])
2621        );
2622        assert_eq!(identities_only_args("/nonexistent/rch/worker_key"), None);
2623
2624        let mut cfg = windows_worker("wsurf");
2625        cfg.identity_file = key;
2626        let command = remote_shell_command(&cfg, "true");
2627        let argv = system_ssh_argv(&cfg, &command, Duration::ZERO);
2628        assert!(argv.contains(&OsString::from("IdentitiesOnly=yes")));
2629
2630        cfg.identity_file = "/nonexistent/rch/worker_key".to_string();
2631        let argv = system_ssh_argv(&cfg, &command, Duration::ZERO);
2632        assert!(!argv.contains(&OsString::from("IdentitiesOnly=yes")));
2633    }
2634
2635    #[test]
2636    fn test_dispatch_helper_matches_declared_os_key() {
2637        // The dispatch key MUST be the same one `pooled_client_options` uses
2638        // so the two Windows fallbacks (no ControlMaster on the pool side,
2639        // system-ssh on the execute side) cannot diverge on which workers
2640        // they apply to.
2641        let _guard = test_guard!();
2642        let cfg = windows_worker("wsurf");
2643        assert_eq!(declared_os(&cfg.tags).as_deref(), Some("windows"));
2644        assert!(prefers_system_ssh(&cfg));
2645    }
2646}
2647
2648#[cfg(test)]
2649mod ssh_connection_tests {
2650    use super::*;
2651    use std::os::unix::fs::PermissionsExt;
2652
2653    struct FakeSsh {
2654        directory: tempfile::TempDir,
2655        program: PathBuf,
2656    }
2657
2658    impl FakeSsh {
2659        fn new(mode: &str) -> Self {
2660            let directory = tempfile::tempdir().unwrap();
2661            let program = directory.path().join("fake-ssh");
2662            let script = format!(
2663                r#"#!/bin/sh
2664root={root}
2665mode={mode}
2666operation=
2667socket=
2668log=
2669while [ "$#" -gt 0 ]; do
2670    case "$1" in
2671        -G) operation=config; shift ;;
2672        -F) shift 2 ;;
2673        -S) socket=$2; shift 2 ;;
2674        -E) log=$2; shift 2 ;;
2675        -O) operation=$2; shift 2 ;;
2676        *) shift ;;
2677    esac
2678done
2679if [ "$operation" = config ]; then
2680    printf '%s\n' "$$" >> "$root/capability-pids"
2681    case "$mode" in
2682        capability-stall) exec sleep 30 ;;
2683        old-client) printf 'hostname none\n' ;;
2684        *) printf 'hostname none\nforkafterauthentication no\n' ;;
2685    esac
2686    exit 0
2687fi
2688if [ "$operation" = check ]; then
2689    printf '%s\n' "$$" >> "$root/check-pids"
2690    case "$mode" in
2691        probe-stall) exec sleep 30 ;;
2692        probe-reject) exit 37 ;;
2693        *) exit 0 ;;
2694    esac
2695fi
2696printf '%s\n' "$$" >> "$root/master-pids"
2697case "$mode" in
2698    ready|old-client|probe-stall|probe-reject) : > "$socket" ;;
2699    fail) printf 'Permission denied (publickey).\n' > "$log"; exit 255 ;;
2700esac
2701exec sleep 30
2702"#,
2703                root = shell_escape_value(&directory.path().to_string_lossy())
2704                    .expect("fixture path must be shell-escapable"),
2705                mode = shell_escape_value(mode).expect("fixture mode must be shell-escapable"),
2706            );
2707            std::fs::write(&program, script).unwrap();
2708            std::fs::set_permissions(&program, std::fs::Permissions::from_mode(0o700)).unwrap();
2709            Self { directory, program }
2710        }
2711
2712        fn client(&self, timeout: Duration, reuse: bool) -> SshClient {
2713            SshClient::new(
2714                WorkerConfig {
2715                    host: "auth-stalled.invalid".to_owned(),
2716                    identity_file: "/nonexistent/rch/test-key".to_owned(),
2717                    ..WorkerConfig::default()
2718                },
2719                SshOptions {
2720                    connect_timeout: timeout,
2721                    control_master: reuse,
2722                    ..SshOptions::default()
2723                },
2724            )
2725        }
2726
2727        async fn connect(&self, client: &mut SshClient) -> Result<()> {
2728            client
2729                .connect_using(&self.program, Some(self.directory.path()))
2730                .await
2731        }
2732
2733        fn pids(&self, filename: &str) -> Vec<u32> {
2734            std::fs::read_to_string(self.directory.path().join(filename))
2735                .unwrap_or_default()
2736                .lines()
2737                .map(|line| line.parse().unwrap())
2738                .collect()
2739        }
2740
2741        async fn wait_for_pid(&self, filename: &str) {
2742            tokio::time::timeout(Duration::from_secs(3), async {
2743                while self.pids(filename).is_empty() {
2744                    tokio::time::sleep(Duration::from_millis(10)).await;
2745                }
2746            })
2747            .await
2748            .expect("fake SSH process must start");
2749        }
2750
2751        async fn assert_reaped(&self) {
2752            let pids: Vec<_> = self
2753                .pids("master-pids")
2754                .into_iter()
2755                .chain(self.pids("check-pids"))
2756                .chain(self.pids("capability-pids"))
2757                .collect();
2758            assert!(!pids.is_empty(), "test must launch a real SSH stand-in");
2759            tokio::time::timeout(Duration::from_secs(3), async {
2760                loop {
2761                    let any_alive = pids.iter().any(|pid| {
2762                        // kill -0 also sees zombies, so this checks reaping,
2763                        // not just sending a termination signal. sh supplies
2764                        // the same builtin on Linux and macOS.
2765                        std::process::Command::new("sh")
2766                            .args(["-c", &format!("kill -0 {pid} 2>/dev/null")])
2767                            .stdin(std::process::Stdio::null())
2768                            .stdout(std::process::Stdio::null())
2769                            .stderr(std::process::Stdio::null())
2770                            .status()
2771                            .unwrap()
2772                            .success()
2773                    });
2774                    let control_dir_exists =
2775                        std::fs::read_dir(self.directory.path())
2776                            .unwrap()
2777                            .any(|entry| {
2778                                entry
2779                                    .unwrap()
2780                                    .file_name()
2781                                    .to_string_lossy()
2782                                    .starts_with("rch-")
2783                            });
2784                    if !any_alive && !control_dir_exists {
2785                        break;
2786                    }
2787                    tokio::time::sleep(Duration::from_millis(10)).await;
2788                }
2789            })
2790            .await
2791            .expect("owned SSH children must be killed, reaped, and release their sockets");
2792        }
2793    }
2794
2795    #[test]
2796    fn foreground_master_keeps_host_identity_and_keepalive_policy() {
2797        let fake = FakeSsh::new("ready");
2798        let key = fake.directory.path().join("identity");
2799        std::fs::write(&key, "test identity").unwrap();
2800        for (policy, expected) in [
2801            (KnownHostsPolicy::Strict, "StrictHostKeyChecking=yes"),
2802            (KnownHostsPolicy::Add, "StrictHostKeyChecking=accept-new"),
2803            (KnownHostsPolicy::AcceptAll, "StrictHostKeyChecking=no"),
2804        ] {
2805            let mut client = fake.client(Duration::from_secs(7), true);
2806            client.config.identity_file = key.to_string_lossy().into_owned();
2807            client.options.known_hosts = policy;
2808            client.options.server_alive_interval = Some(Duration::from_secs(13));
2809            client.options.control_persist_idle = Some(Duration::from_secs(60));
2810            let command = client.master_command(
2811                &fake.program,
2812                "builder@worker",
2813                Path::new("/private/control/master"),
2814                Path::new("/private/control/log"),
2815                true,
2816            );
2817            let args: Vec<_> = command
2818                .as_std()
2819                .get_args()
2820                .map(|argument| argument.to_string_lossy().into_owned())
2821                .collect();
2822            for expected in [
2823                expected,
2824                "BatchMode=yes",
2825                "ConnectTimeout=7",
2826                "ServerAliveInterval=13",
2827                "IdentitiesOnly=yes",
2828                "ControlPersist=no",
2829                "ForkAfterAuthentication=no",
2830                "-M",
2831                "-N",
2832                "builder@worker",
2833            ] {
2834                assert!(args.iter().any(|argument| argument == expected), "{args:?}");
2835            }
2836            assert!(
2837                args.iter()
2838                    .any(|argument| argument == key.to_str().unwrap())
2839            );
2840            assert!(!args.iter().any(|argument| argument == "-f"));
2841            assert!(
2842                !args
2843                    .iter()
2844                    .any(|argument| argument.starts_with("IgnoreUnknown="))
2845            );
2846        }
2847
2848        let client = fake.client(Duration::from_secs(1), false);
2849        let command = client.master_command(
2850            &fake.program,
2851            "builder@worker",
2852            Path::new("/private/master"),
2853            Path::new("/private/log"),
2854            false,
2855        );
2856        assert!(
2857            command
2858                .as_std()
2859                .get_args()
2860                .all(|arg| arg != "IdentitiesOnly=yes")
2861        );
2862        assert!(
2863            command
2864                .as_std()
2865                .get_args()
2866                .all(|arg| arg != "ForkAfterAuthentication=no")
2867        );
2868    }
2869
2870    #[tokio::test]
2871    async fn capability_probe_deadline_and_cancellation_reap_before_authentication() {
2872        for cancelled in [false, true] {
2873            let fake = FakeSsh::new("capability-stall");
2874            let timeout = if cancelled {
2875                Duration::from_secs(30)
2876            } else {
2877                Duration::from_millis(200)
2878            };
2879            let mut client = fake.client(timeout, true);
2880            let program = fake.program.clone();
2881            let directory = fake.directory.path().to_owned();
2882            let connecting =
2883                tokio::spawn(async move { client.connect_using(&program, Some(&directory)).await });
2884            if cancelled {
2885                fake.wait_for_pid("capability-pids").await;
2886                connecting.abort();
2887                assert!(connecting.await.unwrap_err().is_cancelled());
2888            } else {
2889                let error = connecting.await.unwrap().unwrap_err();
2890                assert!(format!("{error:#}").contains("timed out"), "{error:#}");
2891            }
2892            assert!(fake.pids("master-pids").is_empty());
2893            assert_eq!(fake.pids("capability-pids").len(), 1);
2894            fake.assert_reaped().await;
2895        }
2896    }
2897
2898    #[tokio::test]
2899    async fn client_without_fork_option_remains_supported() {
2900        let fake = FakeSsh::new("old-client");
2901        let mut client = fake.client(Duration::from_secs(3), true);
2902        fake.connect(&mut client).await.unwrap();
2903        assert!(client.is_connected());
2904        client.disconnect().await.unwrap();
2905        fake.assert_reaped().await;
2906    }
2907
2908    #[tokio::test]
2909    async fn foreground_options_preserve_user_ignore_unknown_configuration() {
2910        let directory = Arc::new(tempfile::tempdir().unwrap());
2911        let config = directory.path().join("user-config");
2912        std::fs::write(
2913            &config,
2914            "IgnoreUnknown RchFixtureExtension\nRchFixtureExtension yes\n",
2915        )
2916        .unwrap();
2917        let deadline = Instant::now() + Duration::from_secs(3);
2918        let capability = ssh_can_disable_background(Path::new("ssh"), directory.clone(), deadline)
2919            .await
2920            .unwrap();
2921        let client = SshClient::new(WorkerConfig::default(), SshOptions::default());
2922        let master = client.master_command(
2923            Path::new("ssh"),
2924            "none",
2925            &directory.path().join("master"),
2926            &directory.path().join("log"),
2927            capability,
2928        );
2929        // -G inspects the real client's configuration without connecting. Use
2930        // precisely the master's options, including its backgrounding policy.
2931        let mut command = tokio::process::Command::new("ssh");
2932        command
2933            .args(["-G", "-F"])
2934            .arg(config)
2935            .args(master.as_std().get_args())
2936            .stdin(std::process::Stdio::null())
2937            .stdout(std::process::Stdio::null())
2938            .stderr(std::process::Stdio::null());
2939        let process = OwnedSshProcess::spawn(&mut command, Some(directory), None).unwrap();
2940        assert!(process.wait_until(deadline).await.unwrap().success());
2941    }
2942
2943    #[tokio::test]
2944    async fn authentication_deadline_kills_and_reaps_without_starting_fallback() {
2945        let fake = FakeSsh::new("auth-stall");
2946        let mut client = fake.client(Duration::from_millis(200), true);
2947        let start = Instant::now();
2948        let error = fake.connect(&mut client).await.unwrap_err();
2949        assert!(format!("{error:#}").contains("timed out"), "{error:#}");
2950        assert!(start.elapsed() < Duration::from_secs(2));
2951        assert!(!client.is_connected());
2952        assert_eq!(fake.pids("master-pids").len(), 1);
2953        fake.assert_reaped().await;
2954    }
2955
2956    #[tokio::test]
2957    async fn cancelling_authentication_kills_and_reaps_child() {
2958        let fake = FakeSsh::new("auth-stall");
2959        let mut client = fake.client(Duration::from_secs(30), true);
2960        let program = fake.program.clone();
2961        let directory = fake.directory.path().to_owned();
2962        let connecting =
2963            tokio::spawn(async move { client.connect_using(&program, Some(&directory)).await });
2964        fake.wait_for_pid("master-pids").await;
2965        connecting.abort();
2966        assert!(connecting.await.unwrap_err().is_cancelled());
2967        fake.assert_reaped().await;
2968    }
2969
2970    #[tokio::test]
2971    async fn stalled_readiness_probe_shares_deadline_and_is_reaped() {
2972        let fake = FakeSsh::new("probe-stall");
2973        let mut client = fake.client(Duration::from_millis(200), true);
2974        let error = fake.connect(&mut client).await.unwrap_err();
2975        assert!(format!("{error:#}").contains("timed out"), "{error:#}");
2976        assert_eq!(fake.pids("master-pids").len(), 1);
2977        assert_eq!(fake.pids("check-pids").len(), 1);
2978        fake.assert_reaped().await;
2979    }
2980
2981    #[tokio::test]
2982    async fn unsuccessful_readiness_probe_never_marks_client_connected() {
2983        let fake = FakeSsh::new("probe-reject");
2984        let mut client = fake.client(Duration::from_millis(200), false);
2985        assert!(fake.connect(&mut client).await.is_err());
2986        assert!(!client.is_connected());
2987        assert!(!fake.pids("check-pids").is_empty());
2988        fake.assert_reaped().await;
2989    }
2990
2991    #[tokio::test]
2992    async fn dropping_and_disconnecting_ready_client_reap_master() {
2993        let fake = FakeSsh::new("ready");
2994        let mut client = fake.client(Duration::from_secs(3), true);
2995        fake.connect(&mut client).await.unwrap();
2996        assert!(client.is_connected());
2997        client.disconnect().await.unwrap();
2998        assert!(!client.is_connected());
2999        fake.assert_reaped().await;
3000
3001        fake.connect(&mut client).await.unwrap();
3002        assert!(client.is_connected());
3003        drop(client);
3004        fake.assert_reaped().await;
3005        assert_eq!(fake.pids("master-pids").len(), 2);
3006    }
3007
3008    #[tokio::test]
3009    async fn failed_primary_retries_and_preserves_authentication_diagnostic() {
3010        let fake = FakeSsh::new("fail");
3011        let mut client = fake.client(Duration::from_secs(3), true);
3012        let error = fake.connect(&mut client).await.unwrap_err();
3013        assert!(
3014            format!("{error:#}").contains("Permission denied (publickey)."),
3015            "{error:#}"
3016        );
3017        assert_eq!(fake.pids("master-pids").len(), 2);
3018        fake.assert_reaped().await;
3019    }
3020
3021    #[tokio::test]
3022    async fn active_commands_prevent_idle_expiry_and_release_restarts_timer() {
3023        let mut command = tokio::process::Command::new("sleep");
3024        command.arg("30");
3025        let process =
3026            OwnedSshProcess::spawn(&mut command, None, Some(Duration::from_millis(100))).unwrap();
3027        let first = process.begin_use().unwrap();
3028        let second = process.begin_use().unwrap();
3029        tokio::time::sleep(Duration::from_millis(150)).await;
3030        first.complete();
3031        tokio::time::sleep(Duration::from_millis(150)).await;
3032        assert!(
3033            process.is_running(),
3034            "the remaining command still owns the master"
3035        );
3036        second.complete();
3037        assert!(
3038            process.is_running(),
3039            "expiry starts after the last command completes"
3040        );
3041        let status = process
3042            .wait_until(Instant::now() + Duration::from_secs(3))
3043            .await
3044            .unwrap();
3045        assert!(!status.success());
3046        assert!(!process.is_running());
3047        assert!(process.begin_use().is_err());
3048    }
3049
3050    #[tokio::test]
3051    async fn cancelled_command_retires_master_after_other_commands_finish() {
3052        let mut command = tokio::process::Command::new("sleep");
3053        command.arg("30");
3054        let process = OwnedSshProcess::spawn(&mut command, None, None).unwrap();
3055        let cancelled = process.begin_use().unwrap();
3056        let running = process.begin_use().unwrap();
3057        drop(cancelled);
3058        assert!(
3059            process.begin_use().is_err(),
3060            "a draining master rejects new work"
3061        );
3062        tokio::time::sleep(Duration::from_millis(100)).await;
3063        assert!(
3064            process.state.borrow().outcome.is_none(),
3065            "cancellation must not kill another active command's master"
3066        );
3067        running.complete();
3068        let status = process
3069            .wait_until(Instant::now() + Duration::from_secs(3))
3070            .await
3071            .unwrap();
3072        assert!(!status.success());
3073    }
3074
3075    #[tokio::test]
3076    async fn expired_pooled_connection_can_connect_again() {
3077        let fake = FakeSsh::new("ready");
3078        let mut client = fake.client(Duration::from_secs(3), true);
3079        client.options.control_persist_idle = Some(Duration::from_secs(1));
3080        fake.connect(&mut client).await.unwrap();
3081        client
3082            .session
3083            .as_ref()
3084            .unwrap()
3085            .process
3086            .wait_until(Instant::now() + Duration::from_secs(3))
3087            .await
3088            .unwrap();
3089        assert!(!client.is_connected());
3090        fake.connect(&mut client).await.unwrap();
3091        assert!(client.is_connected());
3092        client.disconnect().await.unwrap();
3093        fake.assert_reaped().await;
3094        assert_eq!(fake.pids("master-pids").len(), 2);
3095    }
3096}
3097
3098/// The client-owned lifetime policy for an SSH control master.
3099#[derive(Debug, Clone, PartialEq, Eq)]
3100enum ControlPersistMode {
3101    /// Close the control master when its client is dropped or disconnected.
3102    Closed,
3103    /// Keep a warm control master for the given idle seconds between commands.
3104    IdleFor(NonZeroUsize),
3105    /// Requested idle exceeded usize; caller falls back to Closed.
3106    TooLarge(u64),
3107}
3108
3109/// Decide idle expiry for an SSH session. A warm master is kept ONLY for the
3110/// explicit connection-reuse path (`control_master` + a configured non-zero idle);
3111/// every other case closes with its owning client. OpenSSH persistence is always
3112/// disabled so cancellation can kill and reap the actual master process.
3113fn control_persist_mode(control_master: bool, idle: Option<Duration>) -> ControlPersistMode {
3114    match idle {
3115        Some(idle) if control_master && !idle.is_zero() => match usize::try_from(idle.as_secs()) {
3116            Ok(secs) => match NonZeroUsize::new(secs) {
3117                Some(nonzero) => ControlPersistMode::IdleFor(nonzero),
3118                None => ControlPersistMode::Closed,
3119            },
3120            Err(_) => ControlPersistMode::TooLarge(idle.as_secs()),
3121        },
3122        _ => ControlPersistMode::Closed,
3123    }
3124}
3125
3126#[cfg(test)]
3127mod control_persist_tests {
3128    use super::*;
3129    use std::time::Duration;
3130
3131    #[test]
3132    fn non_mux_sessions_never_persist_forever() {
3133        // The per-call paths (control_master=false) must close after use.
3134        assert_eq!(
3135            control_persist_mode(false, None),
3136            ControlPersistMode::Closed
3137        );
3138        assert_eq!(
3139            control_persist_mode(false, Some(Duration::from_secs(60))),
3140            ControlPersistMode::Closed
3141        );
3142    }
3143
3144    #[test]
3145    fn mux_without_idle_closes() {
3146        assert_eq!(control_persist_mode(true, None), ControlPersistMode::Closed);
3147        assert_eq!(
3148            control_persist_mode(true, Some(Duration::from_secs(0))),
3149            ControlPersistMode::Closed
3150        );
3151    }
3152
3153    #[test]
3154    fn mux_with_idle_keeps_warm() {
3155        assert_eq!(
3156            control_persist_mode(true, Some(Duration::from_secs(60))),
3157            ControlPersistMode::IdleFor(NonZeroUsize::new(60).unwrap())
3158        );
3159    }
3160}