diff --git a/src/app/dispatcher.rs b/src/app/dispatcher.rs index 1c81785b..a964f8a4 100644 --- a/src/app/dispatcher.rs +++ b/src/app/dispatcher.rs @@ -32,6 +32,11 @@ use bssh::{ ssh::{ CliTtyMode, SessionPolicy, SessionRequest, SshClient, client::ConnectionConfig, + control::{ + AttachOutcome, ControlCommand, ControlPathContext, ControlPolicy, ControlResponseKind, + SessionOpenRequest, attach_session, connect_control_socket, expand_control_path, + remove_stale_control_socket, send_control_command, start_control_master, + }, tokio_client::{AddressFamily, ProxyMode, SshConnectionConfigResolver}, }, }; @@ -79,6 +84,277 @@ fn build_ssh_connection_config_resolver( .with_stdio_forward(cli.stdio_forward.is_some()) } +#[derive(Debug, Clone)] +struct ResolvedControlInvocation { + policy: ControlPolicy, + path: PathBuf, + session_policy: SessionPolicy, + session_request: SessionOpenRequest, + forwarding_directives: Vec, + address_family: AddressFamily, + jump_spec: Option, +} + +fn resolve_control_invocation( + cli: &Cli, + ctx: &AppContext, + command: &str, +) -> Result> { + if !cli.is_ssh_mode() || ctx.nodes.len() != 1 { + anyhow::ensure!( + cli.control_command.is_none(), + "-O requires exactly one SSH destination" + ); + return Ok(None); + } + anyhow::ensure!( + cli.stdio_forward.is_none() || (cli.control_master == 0 && cli.control_command.is_none()), + "-W cannot be combined with connection multiplexing" + ); + let node = ctx + .nodes + .first() + .context("Connection multiplexing requires an SSH destination")?; + let effective = ctx.ssh_config.find_host_config(node.config_host()); + let policy = ControlPolicy::from_raw( + effective.control_master.as_deref(), + effective.control_path.as_deref(), + effective.control_persist.as_deref(), + )?; + let Some(template) = policy.path.as_deref() else { + anyhow::ensure!( + cli.control_command.is_none(), + "-O requires ControlPath (use --control-path or -o ControlPath=...)" + ); + return Ok(None); + }; + let resolver = build_ssh_connection_config_resolver( + cli, + ctx, + ctx.cluster_name.as_deref().or(cli.cluster.as_deref()), + ); + let resolved_connection = resolver.resolve_for_host(node.config_host()); + let jump_spec = + session_policy_jump_spec(resolved_connection.proxy_mode.as_ref()).map(str::to_string); + let local_host = whoami::hostname().unwrap_or_else(|_| "localhost".to_string()); + let home = dirs::home_dir().unwrap_or_default(); + let mut path_context = ControlPathContext::new( + local_host, + home, + node.host.clone(), + node.port, + node.username.clone(), + ); + if let Some(ProxyMode::Jump(jump)) = resolved_connection.proxy_mode.as_ref() { + path_context = path_context.with_jump_host(jump); + } + let path = expand_control_path(template, &path_context)?; + let session_policy = SessionPolicy::resolve_with_jump_spec( + &effective, + node, + (!command.is_empty()).then_some(command), + cli_tty_mode(cli), + std::io::stdin().is_terminal(), + jump_spec.as_deref(), + )?; + let mut remote_policy = session_policy.clone(); + remote_policy.local_command = None; + let terminal = remote_policy + .request_pty + .then(|| std::env::var("TERM").unwrap_or_else(|_| "xterm".to_string())); + let session_request = SessionOpenRequest::new(remote_policy, terminal)?; + Ok(Some(ResolvedControlInvocation { + policy, + path, + session_policy, + session_request, + forwarding_directives: resolved_connection.forwarding_plan.directives.clone(), + address_family: resolved_connection.address_family, + jump_spec, + })) +} + +async fn try_existing_control_master( + cli: &Cli, + control: &ResolvedControlInvocation, +) -> Result> { + if let Some(command) = cli.control_command.as_deref() { + let command = command.parse::()?; + let forwards = if matches!(command, ControlCommand::Forward | ControlCommand::Cancel) { + control.forwarding_directives.clone() + } else { + Vec::new() + }; + let response = + send_control_command(&control.path, command, forwards, control.address_family).await?; + match response { + ControlResponseKind::Alive { pid } => { + println!("Master running (pid={pid})"); + } + ControlResponseKind::Ok => {} + response => anyhow::bail!("unexpected control command response: {response:?}"), + } + return Ok(Some(EXIT_SUCCESS)); + } + if !control.policy.master.tries_existing() { + return Ok(None); + } + match attach_session( + &control.path, + control.session_request.clone(), + &control.session_policy, + ) + .await? + { + AttachOutcome::NoMaster => Ok(None), + AttachOutcome::ExitStatus(status) => Ok(Some(i32::try_from(status).unwrap_or(255))), + } +} + +async fn execute_initial_control_session( + client: &bssh::ssh::tokio_client::Client, + policy: &SessionPolicy, +) -> Result { + policy.run_local_command().await?; + if matches!(policy.request, SessionRequest::None) { + return Ok(0); + } + let (sender, mut receiver) = tokio::sync::mpsc::channel(128); + let output = tokio::spawn(async move { + use tokio::io::AsyncWriteExt as _; + + let mut stdout = tokio::io::stdout(); + let mut stderr = tokio::io::stderr(); + let mut stdout_open = true; + let mut stderr_open = true; + while let Some(event) = receiver.recv().await { + match event { + bssh::ssh::tokio_client::CommandOutput::StdOut(bytes) if stdout_open => { + if stdout.write_all(&bytes).await.is_err() { + stdout_open = false; + } else { + stdout.flush().await.ok(); + } + } + bssh::ssh::tokio_client::CommandOutput::StdErr(bytes) if stderr_open => { + if stderr.write_all(&bytes).await.is_err() { + stderr_open = false; + } else { + stderr.flush().await.ok(); + } + } + _ => {} + } + } + }); + let result = if policy.stdin_null { + client.execute_session_streaming(policy, sender).await + } else { + client + .execute_session_streaming_with_stdin(policy, sender) + .await + }; + output.await.context("Control-master output task failed")?; + result.map_err(anyhow::Error::from) +} + +async fn prepare_control_socket_for_master(path: &Path) -> Result<()> { + match connect_control_socket(path).await { + Ok(_) => anyhow::bail!( + "A control master is already running at '{}'", + path.display() + ), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::ConnectionRefused => { + anyhow::ensure!( + remove_stale_control_socket(path)?, + "Refusing to remove stale ControlPath '{}': it is not an owned Unix socket", + path.display() + ); + Ok(()) + } + Err(error) => Err(error).with_context(|| { + format!( + "Could not inspect existing ControlPath '{}'", + path.display() + ) + }), + } +} + +async fn handle_control_master( + cli: &Cli, + ctx: &AppContext, + control: &ResolvedControlInvocation, + ssh_password: Option>, +) -> Result { + anyhow::ensure!( + control.policy.master.creates_master(), + "internal error: non-master control policy reached master startup" + ); + prepare_control_socket_for_master(&control.path).await?; + let node = ctx + .nodes + .first() + .context("Connection multiplexing requires an SSH destination")?; + let effective_cluster_name = ctx.cluster_name.as_deref().or(cli.cluster.as_deref()); + let resolver = build_ssh_connection_config_resolver(cli, ctx, effective_cluster_name); + let resolved_connection = resolver.resolve_for_host(node.config_host()); + let key_path = determine_ssh_key_path( + cli, + &ctx.config, + &ctx.ssh_config, + Some(node.config_host()), + effective_cluster_name, + ); + #[cfg(target_os = "macos")] + let use_keychain = determine_use_keychain(&ctx.ssh_config, Some(node.config_host())); + let connection = ConnectionConfig { + key_path: key_path.as_deref(), + strict_mode: Some(ctx.strict_mode), + use_agent: cli.use_agent, + use_password: cli.password, + #[cfg(target_os = "macos")] + use_keychain, + timeout_seconds: cli.timeout, + connect_timeout_seconds: Some(cli.connect_timeout), + jump_hosts_spec: control.jump_spec.as_deref(), + ssh_connection_config: Some(&resolved_connection), + ssh_connection_config_resolver: Some(&resolver), + session_policy: Some(&control.session_policy), + ssh_password, + }; + let mut ssh_client = SshClient::new(node.host.clone(), node.port, node.username.clone()); + let client = ssh_client.connect_authenticated(&connection).await?; + let master = match start_control_master( + &control.path, + client.clone(), + control.policy.master.requires_confirmation(), + ) { + Ok(master) => master, + Err(error) => { + let _ = client.disconnect().await; + return Err(error); + } + }; + let initial_request_was_none = matches!(control.session_policy.request, SessionRequest::None); + let status = match execute_initial_control_session(&client, &control.session_policy).await { + Ok(status) => status, + Err(error) => { + if let Err(shutdown_error) = master.shutdown_immediately().await { + tracing::warn!( + "Initial control-master session failed and teardown also failed: {shutdown_error:#}" + ); + } + return Err(error); + } + }; + master + .finish_after_initial(control.policy.persist, initial_request_was_none) + .await?; + Ok(i32::try_from(status).unwrap_or(255)) +} + /// Decide whether `-S` (sudo-password) is meaningful for the given dispatch path. /// /// Only command execution can consume `SudoPassword`, because it monitors @@ -206,6 +482,16 @@ pub async fn dispatch_command(cli: &Cli, ctx: &AppContext) -> Result { ); } + // A passenger must attach before any key selection, password prompt, or + // network authentication. This is the invariant that guarantees repeated + // invocations reuse the master's one authenticated transport. + let control = resolve_control_invocation(cli, ctx, &command)?; + if let Some(control) = control.as_ref() + && let Some(exit_code) = try_existing_control_master(cli, control).await? + { + return Ok(exit_code); + } + // Calculate hostname for SSH config integration before deciding whether an // up-front password prompt is permitted by every actual target policy. let hostname_for_ssh_config = if cli.is_ssh_mode() { @@ -389,6 +675,11 @@ pub async fn dispatch_command(cli: &Cli, ctx: &AppContext) -> Result { unreachable!("CacheStats should be handled before dispatch") } None => { + if let Some(control) = control.as_ref() + && control.policy.master.creates_master() + { + return handle_control_master(cli, ctx, control, ssh_password).await; + } // Execute command (auto-exec or interactive shell). This path owns // its own exit code strategy (`ExitCodeStrategy`, selected by // `--require-all-success` / `--check-all-nodes`) and exits the diff --git a/src/cli/bssh.rs b/src/cli/bssh.rs index 1bdbc8c5..d96e9b32 100644 --- a/src/cli/bssh.rs +++ b/src/cli/bssh.rs @@ -269,6 +269,31 @@ pub struct Cli { help = "SSH options (e.g., -o StrictHostKeyChecking=no)")] pub ssh_options: Vec, + #[arg( + short = 'M', + long = "control-master", + action = clap::ArgAction::Count, + help = "Enable SSH connection sharing; repeat to require confirmation" + )] + pub control_master: u8, + + #[arg( + short = 'O', + long = "control-command", + value_name = "command", + value_parser = ["check", "forward", "cancel", "exit", "stop"], + conflicts_with = "stdio_forward", + help = "Send a control command to an existing connection master" + )] + pub control_command: Option, + + #[arg( + long = "control-path", + value_name = "path", + help = "Path template for the connection-sharing control socket" + )] + pub control_path: Option, + #[arg( short = 'c', long = "cipher", @@ -677,8 +702,21 @@ impl Cli { + usize::from(self.cipher.is_some()) + usize::from(self.macs.is_some()) + usize::from(self.subsystem) - + usize::from(self.stdin_null), + + usize::from(self.stdin_null) + + usize::from(self.control_master > 0) + + usize::from(self.control_path.is_some()), ); + if self.control_master > 0 { + let mode = if self.control_master == 1 { + "yes" + } else { + "ask" + }; + options.push(format!("ControlMaster={mode}")); + } + if let Some(path) = &self.control_path { + options.push(format!("ControlPath={}", path.display())); + } if let Some(cipher) = &self.cipher { options.push(format!("Ciphers={cipher}")); } @@ -946,6 +984,40 @@ mod tests { ); } + #[test] + fn multiplex_flags_preserve_openssh_priority_and_repetition() { + let once = Cli::try_parse_from([ + "bssh", + "-M", + "--control-path", + "/tmp/bssh-%C", + "-o", + "ControlMaster=no", + "target", + ]) + .unwrap(); + assert_eq!(once.control_master, 1); + assert_eq!( + once.ssh_config_overrides(), + [ + "ControlMaster=yes", + "ControlPath=/tmp/bssh-%C", + "ControlMaster=no" + ] + ); + + let repeated = + Cli::try_parse_from(["bssh", "-MM", "-Ocheck", "--control-path=/tmp/c", "target"]) + .unwrap(); + assert_eq!(repeated.control_master, 2); + assert_eq!(repeated.control_command.as_deref(), Some("check")); + assert_eq!(repeated.ssh_config_overrides()[0], "ControlMaster=ask"); + + let repeated_more = Cli::try_parse_from(["bssh", "-MMM", "target"]).unwrap(); + assert_eq!(repeated_more.control_master, 3); + assert_eq!(repeated_more.ssh_config_overrides()[0], "ControlMaster=ask"); + } + #[test] fn compatibility_flags_parse_repetition_modifiers_and_session_overrides() { let cli = Cli::try_parse_from([ diff --git a/src/cli/pdsh.rs b/src/cli/pdsh.rs index d39423bd..53c5f83b 100644 --- a/src/cli/pdsh.rs +++ b/src/cli/pdsh.rs @@ -317,6 +317,9 @@ impl PdshCli { require_all_success: false, check_all_nodes: false, ssh_options: Vec::new(), + control_master: 0, + control_command: None, + control_path: None, cipher: None, macs: None, subsystem: false, diff --git a/src/cli/ssh_args.rs b/src/cli/ssh_args.rs index 1341ee7a..d2a26ec2 100644 --- a/src/cli/ssh_args.rs +++ b/src/cli/ssh_args.rs @@ -101,10 +101,9 @@ fn scoped_second_pass_width(argument: &str, next: Option<&String>) -> Option {} - 'c' | 'm' | 'W' | 'F' | 'o' => { + 's' | 'n' | 'M' => {} + 'c' | 'm' | 'W' | 'F' | 'o' | 'O' => { let attached = position + short.len_utf8() < shorts.len(); return Some(usize::from(!attached && next.is_some()) + 1); } diff --git a/src/forwarding/manager.rs b/src/forwarding/manager.rs index 7edb1328..f91eae45 100644 --- a/src/forwarding/manager.rs +++ b/src/forwarding/manager.rs @@ -416,16 +416,22 @@ impl ForwardingManager { // Cancel the forwarding task session.cancel_token.cancel(); - // Wait for task to complete if it exists - if let Some(task) = session.task_handle.take() { - let _ = task.await; // Ignore join errors - } + // Wait for task completion so remote cancel failures reach callers. + let task_result = match session.task_handle.take() { + Some(task) => task + .await + .with_context(|| format!("Forwarding session {id} task failed to join"))?, + None => match &session.status { + ForwardingStatus::Failed(reason) => Err(anyhow::anyhow!(reason.clone())), + _ => Ok(()), + }, + }; session.status = ForwardingStatus::Stopped; session.updated_at = Instant::now(); tracing::info!("Stopped forwarding session {}", id); - Ok(()) + task_result } /// Stop all forwarding sessions @@ -482,7 +488,7 @@ impl ForwardingManager { /// Remove a forwarding session (must be stopped first) pub async fn remove_forwarding(&mut self, id: ForwardingId) -> Result<()> { // Ensure session is stopped first - let _ = self.stop_forwarding(id).await; + let stop_result = self.stop_forwarding(id).await; let mut sessions = self.sessions.write().await; sessions @@ -490,7 +496,7 @@ impl ForwardingManager { .ok_or_else(|| anyhow::anyhow!("Forwarding session {id} not found"))?; tracing::info!("Removed forwarding session {}", id); - Ok(()) + stop_result } /// Shutdown the ForwardingManager and all active sessions diff --git a/src/forwarding/mod.rs b/src/forwarding/mod.rs index 9e2ff110..89d0d97a 100644 --- a/src/forwarding/mod.rs +++ b/src/forwarding/mod.rs @@ -47,7 +47,7 @@ pub(crate) fn format_host_port(host: &str, port: u16) -> String { } /// One forwarding directive in the order OpenSSH obtained it. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum ForwardingDirective { Local(String), Remote(String), @@ -109,7 +109,7 @@ impl ForwardingPlan { } /// Port forwarding specification types -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum ForwardingType { /// Local port forwarding (-L) /// Format: [bind_address:]port:host:hostport @@ -138,7 +138,7 @@ pub enum ForwardingType { } /// SOCKS protocol version for dynamic forwarding -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum SocksVersion { V4, V5, diff --git a/src/forwarding/remote.rs b/src/forwarding/remote.rs index 726978c7..2602e696 100644 --- a/src/forwarding/remote.rs +++ b/src/forwarding/remote.rs @@ -325,9 +325,10 @@ impl RemoteForwarder { Self::cancel_request(&ssh_client, &address, allocated_port, request_timeout).await; registry.unregister(&target).await; if let Err(error) = cancel_result { - tracing::warn!( - "Failed to cancel remote forwarding {address}:{allocated_port}: {error}" - ); + let message = + format!("Failed to cancel remote forwarding {address}:{allocated_port}: {error}"); + send_status(ForwardingStatus::Failed(message.clone())); + anyhow::bail!(message); } send_status(ForwardingStatus::Stopped); Ok(()) diff --git a/src/forwarding/runtime.rs b/src/forwarding/runtime.rs index 0095ecd2..2c9491ee 100644 --- a/src/forwarding/runtime.rs +++ b/src/forwarding/runtime.rs @@ -5,12 +5,56 @@ use std::sync::Arc; use anyhow::{Context, Result}; use tokio::sync::Mutex; -use super::{ForwardingConfig, ForwardingManager, ForwardingPlan}; -use crate::ssh::tokio_client::Client; +use super::{ForwardingConfig, ForwardingId, ForwardingManager, ForwardingPlan, ForwardingType}; +use crate::ssh::tokio_client::{AddressFamily, Client}; + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ConfiguredForwarding { + spec: ForwardingType, + id: ForwardingId, +} + +#[derive(Default)] +struct ForwardingRuntimeState { + manager: Option, + configured: Vec, + address_family: Option, +} + +impl ForwardingRuntimeState { + fn forwarding_id(&self, spec: &ForwardingType) -> Option { + self.configured + .iter() + .find(|forwarding| &forwarding.spec == spec) + .map(|forwarding| forwarding.id) + } + + fn track(&mut self, spec: ForwardingType, id: ForwardingId) { + debug_assert!(self.forwarding_id(&spec).is_none()); + self.configured.push(ConfiguredForwarding { spec, id }); + } + + fn untrack(&mut self, id: ForwardingId) { + self.configured.retain(|forwarding| forwarding.id != id); + } + + async fn shutdown_if_idle(&mut self) -> Result<()> { + if !self.configured.is_empty() { + return Ok(()); + } + + let shutdown_result = match self.manager.take() { + Some(mut manager) => manager.shutdown().await, + None => Ok(()), + }; + self.address_family = None; + shutdown_result + } +} #[derive(Default)] pub struct ForwardingRuntime { - manager: Mutex>, + state: Mutex, } impl std::fmt::Debug for ForwardingRuntime { @@ -19,7 +63,10 @@ impl std::fmt::Debug for ForwardingRuntime { .debug_struct("ForwardingRuntime") .field( "active", - &self.manager.try_lock().is_ok_and(|slot| slot.is_some()), + &self + .state + .try_lock() + .is_ok_and(|state| state.manager.is_some()), ) .finish() } @@ -37,8 +84,8 @@ impl ForwardingRuntime { return Ok(()); } - let mut slot = self.manager.lock().await; - if slot.is_some() { + let mut state = self.state.lock().await; + if state.manager.is_some() { return Ok(()); } @@ -51,7 +98,7 @@ impl ForwardingRuntime { let mut manager = ForwardingManager::new(config); manager.start().await?; let transport = Arc::new(client.forwarding_transport_clone()); - let mut active = 0usize; + let mut configured = Vec::with_capacity(forwards.len()); for forward in forwards { let id = manager @@ -72,22 +119,213 @@ impl ForwardingRuntime { crate::diagnosticln!("Warning: forwarding {forward} failed: {error}"); continue; } - active += 1; + configured.push(ConfiguredForwarding { spec: forward, id }); } - if active == 0 { + if configured.is_empty() { manager.shutdown().await?; } else { - *slot = Some(manager); + state.manager = Some(manager); + state.configured = configured; + state.address_family = Some(plan.address_family); } Ok(()) } + /// Add forwarding requests to an already-authenticated control master. + /// + /// Existing identical requests are successful no-ops, matching OpenSSH's + /// multiplexing behavior. New requests are not reported as ready until + /// their local listener is bound or their remote forwarding request has + /// been acknowledged by the server. Independent requests are attempted + /// even if another request fails. + pub async fn add_control_forwardings( + &self, + client: &Client, + plan: &ForwardingPlan, + ) -> Result<()> { + let forwards = plan.parse()?; + if forwards.is_empty() { + return Ok(()); + } + + let mut state = self.state.lock().await; + if let Some(address_family) = state.address_family { + anyhow::ensure!( + address_family == plan.address_family, + "Cannot add forwarding with address family {:?}; the control master uses {:?}", + plan.address_family, + address_family + ); + } + + if state.manager.is_none() { + let config = ForwardingConfig { + auto_reconnect: false, + max_reconnect_attempts: 1, + address_family: plan.address_family, + ..ForwardingConfig::default() + }; + let mut manager = ForwardingManager::new(config); + manager + .start() + .await + .context("Failed to start forwarding manager for control request")?; + state.manager = Some(manager); + state.address_family = Some(plan.address_family); + } + + let transport = Arc::new(client.forwarding_transport_clone()); + let mut failures = Vec::new(); + for forward in forwards { + if state.forwarding_id(&forward).is_some() { + tracing::debug!(forward = %forward, "Control forwarding is already active"); + continue; + } + + let id_result = { + let manager = state + .manager + .as_mut() + .context("Forwarding manager disappeared during control request")?; + manager.add_forwarding(forward.clone()).await + }; + let id = match id_result { + Ok(id) => id, + Err(error) => { + failures.push(format!("{forward}: {error:#}")); + continue; + } + }; + let start_result = { + let manager = state + .manager + .as_ref() + .context("Forwarding manager disappeared during control request")?; + manager + .start_forwarding_and_wait(id, Arc::clone(&transport)) + .await + }; + match start_result { + Ok(()) => state.track(forward, id), + Err(error) => { + let cleanup_result = { + let manager = state + .manager + .as_mut() + .context("Forwarding manager disappeared during cleanup")?; + manager.remove_forwarding(id).await + }; + if let Err(cleanup_error) = cleanup_result { + tracing::warn!( + forwarding_id = %id, + "Failed to clean up rejected control forwarding: {cleanup_error:#}" + ); + } + failures.push(format!("{forward}: {error:#}")); + } + } + } + + state.shutdown_if_idle().await?; + + if failures.is_empty() { + Ok(()) + } else { + anyhow::bail!( + "One or more control forwarding requests failed: {}", + failures.join("; ") + ) + } + } + + /// Cancel forwarding requests previously configured on a control master. + /// + /// Cancellation is exact: an unregistered specification is an error. The + /// method waits for each forwarding task to finish before returning, which + /// includes the remote `cancel-tcpip-forward` round trip for `-R` requests. + /// Independent requests are attempted even if another request is missing + /// or fails to stop. + pub async fn cancel_control_forwardings(&self, plan: &ForwardingPlan) -> Result<()> { + let forwards = plan.parse()?; + if forwards.is_empty() { + return Ok(()); + } + + let mut state = self.state.lock().await; + let mut failures = Vec::new(); + for forward in forwards { + let Some(id) = state.forwarding_id(&forward) else { + failures.push(format!("{forward}: port not forwarded")); + continue; + }; + let remove_result = match state.manager.as_mut() { + Some(manager) => manager.remove_forwarding(id).await, + None => Err(anyhow::anyhow!("Forwarding manager is not running")), + }; + match remove_result { + Ok(()) => state.untrack(id), + Err(error) => failures.push(format!("{forward}: {error:#}")), + } + } + + state.shutdown_if_idle().await?; + + if failures.is_empty() { + Ok(()) + } else { + anyhow::bail!( + "One or more control forwarding cancellations failed: {}", + failures.join("; ") + ) + } + } + /// Deterministically stop listeners and cancel remote forwarding requests. pub async fn shutdown(&self) -> Result<()> { - if let Some(mut manager) = self.manager.lock().await.take() { - manager.shutdown().await?; + let mut state = self.state.lock().await; + let shutdown_result = match state.manager.take() { + Some(mut manager) => manager.shutdown().await, + None => Ok(()), + }; + state.configured.clear(); + state.address_family = None; + shutdown_result + } +} + +#[cfg(test)] +mod tests { + use std::net::{IpAddr, Ipv4Addr}; + + use super::*; + + fn local_forward(bind_port: u16, remote_port: u16) -> ForwardingType { + ForwardingType::Local { + bind_addr: IpAddr::V4(Ipv4Addr::LOCALHOST), + bind_port, + remote_host: "example.com".to_string(), + remote_port, } - Ok(()) + } + + #[test] + fn forwarding_tracking_is_exact_and_removable() { + let first = local_forward(3001, 80); + let second = local_forward(3002, 80); + let first_id = ForwardingId::new_v4(); + let second_id = ForwardingId::new_v4(); + let mut state = ForwardingRuntimeState::default(); + + state.track(first.clone(), first_id); + state.track(second.clone(), second_id); + + assert_eq!(state.forwarding_id(&first), Some(first_id)); + assert_eq!(state.forwarding_id(&second), Some(second_id)); + assert_eq!(state.forwarding_id(&local_forward(3001, 81)), None); + + state.untrack(first_id); + assert_eq!(state.forwarding_id(&first), None); + assert_eq!(state.forwarding_id(&second), Some(second_id)); } } diff --git a/src/ssh/client/command.rs b/src/ssh/client/command.rs index e7c3c105..4098fc91 100644 --- a/src/ssh/client/command.rs +++ b/src/ssh/client/command.rs @@ -32,6 +32,48 @@ use tokio::sync::mpsc::Sender; const DEFAULT_COMMAND_TIMEOUT_SECS: u64 = 300; impl SshClient { + /// Establish an authenticated SSH transport without closing it after one operation. + /// + /// Connection multiplexing keeps the returned client in the control master and + /// opens one channel per attached invocation. Callers own the transport lifetime + /// and must invoke [`crate::ssh::tokio_client::Client::disconnect`] exactly once + /// when the master terminates. + pub async fn connect_authenticated( + &mut self, + config: &ConnectionConfig<'_>, + ) -> Result { + let auth_method = self + .determine_auth_method( + config.key_path, + config.use_agent, + config.use_password, + #[cfg(target_os = "macos")] + config.use_keychain, + config.ssh_password.clone(), + config.ssh_connection_config, + ) + .await?; + let strict_mode = config + .strict_mode + .unwrap_or(StrictHostKeyChecking::AcceptNew); + self.establish_connection( + &auth_method, + strict_mode, + config.jump_hosts_spec, + config.key_path, + config.use_agent, + config.use_password, + config.connect_timeout_seconds, + config.ssh_connection_config, + config.ssh_connection_config_resolver, + config.ssh_password.clone(), + config + .session_policy + .map_or(crate::ssh::SessionPurpose::Bulk, |policy| policy.purpose()), + ) + .await + } + async fn finish_with_disconnect( &self, client: &crate::ssh::tokio_client::Client, @@ -167,41 +209,8 @@ impl SshClient { ) -> Result { tracing::debug!("Connecting to {}:{}", self.host, self.port); - // Determine authentication method based on parameters - let auth_method = self - .determine_auth_method( - config.key_path, - config.use_agent, - config.use_password, - #[cfg(target_os = "macos")] - config.use_keychain, - config.ssh_password.clone(), - config.ssh_connection_config, - ) - .await?; - - let strict_mode = config - .strict_mode - .unwrap_or(StrictHostKeyChecking::AcceptNew); - - // Create client connection - either direct or through jump hosts - let client = self - .establish_connection( - &auth_method, - strict_mode, - config.jump_hosts_spec, - config.key_path, - config.use_agent, - config.use_password, - config.connect_timeout_seconds, - config.ssh_connection_config, - config.ssh_connection_config_resolver, - config.ssh_password.clone(), - config - .session_policy - .map_or(crate::ssh::SessionPurpose::Bulk, |policy| policy.purpose()), - ) - .await?; + // Create client connection - either direct or through jump hosts. + let client = self.connect_authenticated(config).await?; tracing::debug!("Connected and authenticated successfully"); let operation = async { diff --git a/src/ssh/control/config.rs b/src/ssh/control/config.rs new file mode 100644 index 00000000..d5893e2c --- /dev/null +++ b/src/ssh/control/config.rs @@ -0,0 +1,395 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +use std::fmt; +use std::str::FromStr; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +/// OpenSSH `ControlMaster` connection-selection policy. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)] +#[non_exhaustive] +#[serde(rename_all = "snake_case")] +pub enum ControlMasterMode { + /// Try an existing master but never create one. + #[default] + No, + /// Create a master without prompting and do not attach to an existing one. + Yes, + /// Create a master and ask before accepting each shared request. + Ask, + /// Try an existing master and create one when none is reachable. + Auto, + /// `Auto` behavior with confirmation for shared requests. + AutoAsk, +} + +impl ControlMasterMode { + /// Whether an invocation should first try the configured control socket. + #[must_use] + pub const fn tries_existing(self) -> bool { + matches!(self, Self::No | Self::Auto | Self::AutoAsk) + } + + /// Whether a direct connection should publish a master socket. + #[must_use] + pub const fn creates_master(self) -> bool { + !matches!(self, Self::No) + } + + /// Whether a master must confirm incoming shared requests. + #[must_use] + pub const fn requires_confirmation(self) -> bool { + matches!(self, Self::Ask | Self::AutoAsk) + } +} + +impl FromStr for ControlMasterMode { + type Err = ControlConfigError; + + fn from_str(value: &str) -> Result { + match value.to_ascii_lowercase().as_str() { + "no" | "false" => Ok(Self::No), + "yes" | "true" => Ok(Self::Yes), + "ask" => Ok(Self::Ask), + "auto" => Ok(Self::Auto), + "autoask" => Ok(Self::AutoAsk), + _ => Err(ControlConfigError::InvalidControlMaster(value.to_string())), + } + } +} + +impl fmt::Display for ControlMasterMode { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(match self { + Self::No => "no", + Self::Yes => "yes", + Self::Ask => "ask", + Self::Auto => "auto", + Self::AutoAsk => "autoask", + }) + } +} + +/// Lifetime policy for an idle multiplexing master. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)] +#[non_exhaustive] +#[serde(rename_all = "snake_case")] +pub enum ControlPersist { + /// Exit when the initial client and all shared sessions are gone. + #[default] + Disabled, + /// Remain available until explicitly stopped. + Forever, + /// Exit after the connection has been idle for this duration. + Timeout(Duration), +} + +impl ControlPersist { + /// Whether the master outlives its initial client. + #[must_use] + pub const fn is_enabled(self) -> bool { + !matches!(self, Self::Disabled) + } + + /// Return the finite idle timeout, if one was configured. + #[must_use] + pub const fn timeout(self) -> Option { + match self { + Self::Timeout(duration) => Some(duration), + Self::Disabled | Self::Forever => None, + } + } +} + +impl FromStr for ControlPersist { + type Err = ControlConfigError; + + fn from_str(value: &str) -> Result { + match value.to_ascii_lowercase().as_str() { + "no" | "false" => Ok(Self::Disabled), + "yes" | "true" => Ok(Self::Forever), + _ => { + let seconds = parse_compound_duration(value)?; + if seconds == 0 { + Ok(Self::Forever) + } else { + Ok(Self::Timeout(Duration::from_secs(seconds))) + } + } + } + } +} + +impl fmt::Display for ControlPersist { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Disabled => formatter.write_str("no"), + Self::Forever => formatter.write_str("yes"), + Self::Timeout(duration) => write!(formatter, "{}", duration.as_secs()), + } + } +} + +/// Commands accepted by OpenSSH-compatible `-O` handling in issue #286. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[non_exhaustive] +#[serde(rename_all = "snake_case")] +pub enum ControlCommand { + Check, + Forward, + Cancel, + Exit, + Stop, +} + +impl FromStr for ControlCommand { + type Err = ControlConfigError; + + fn from_str(value: &str) -> Result { + match value { + "check" => Ok(Self::Check), + "forward" => Ok(Self::Forward), + "cancel" => Ok(Self::Cancel), + "exit" => Ok(Self::Exit), + "stop" => Ok(Self::Stop), + _ => Err(ControlConfigError::InvalidControlCommand(value.to_string())), + } + } +} + +impl fmt::Display for ControlCommand { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(match self { + Self::Check => "check", + Self::Forward => "forward", + Self::Cancel => "cancel", + Self::Exit => "exit", + Self::Stop => "stop", + }) + } +} + +/// Typed view over the three string-valued fields currently stored in +/// `SshHostConfig`. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)] +pub struct ControlPolicy { + pub master: ControlMasterMode, + /// Unexpanded path template. `None` includes explicit `ControlPath none`. + pub path: Option, + pub persist: ControlPersist, +} + +impl ControlPolicy { + /// Parse the resolved raw `SshHostConfig` values without changing that + /// structure's public representation. + pub fn from_raw( + control_master: Option<&str>, + control_path: Option<&str>, + control_persist: Option<&str>, + ) -> Result { + let master = control_master.map_or(Ok(ControlMasterMode::No), str::parse)?; + let path = control_path + .filter(|value| !value.eq_ignore_ascii_case("none")) + .map(str::to_string); + let persist = control_persist.map_or(Ok(ControlPersist::Disabled), str::parse)?; + Ok(Self { + master, + path, + persist, + }) + } + + /// Multiplexing has an operational socket only when `ControlPath` did not + /// resolve to unset or `none`. + #[must_use] + pub const fn is_socket_enabled(&self) -> bool { + self.path.is_some() + } +} + +/// Invalid user-authored control configuration. +#[derive(Debug, Clone, PartialEq, Eq, Error)] +#[non_exhaustive] +pub enum ControlConfigError { + #[error("invalid ControlMaster value '{0}'; expected yes, no, ask, auto, or autoask")] + InvalidControlMaster(String), + #[error( + "invalid ControlPersist value '{0}'; expected yes, no, or a duration such as 30s or 1h30m" + )] + InvalidControlPersist(String), + #[error("ControlPersist duration '{0}' is too large")] + ControlPersistOverflow(String), + #[error("invalid control command '{0}'; expected check, forward, cancel, exit, or stop")] + InvalidControlCommand(String), +} + +// OpenSSH's convtime() returns an int and rejects anything larger than +// INT_MAX. Keeping the same ceiling avoids platform-dependent timer behavior. +const MAX_CONTROL_PERSIST_SECONDS: u64 = i32::MAX as u64; + +fn parse_compound_duration(value: &str) -> Result { + if value.is_empty() { + return Err(ControlConfigError::InvalidControlPersist(value.to_string())); + } + + let bytes = value.as_bytes(); + let mut index = 0usize; + let mut total = 0u64; + + while index < bytes.len() { + let mut integer = 0u64; + let mut integer_digits = 0usize; + while index < bytes.len() && bytes[index].is_ascii_digit() { + integer = integer + .checked_mul(10) + .and_then(|number| number.checked_add(u64::from(bytes[index] - b'0'))) + .ok_or_else(|| ControlConfigError::ControlPersistOverflow(value.to_string()))?; + integer_digits += 1; + index += 1; + } + if integer_digits == 0 { + return Err(ControlConfigError::InvalidControlPersist(value.to_string())); + } + + let multiplier = if index == bytes.len() { + 1u64 + } else { + let unit = bytes[index]; + index += 1; + match unit { + b's' | b'S' => 1, + b'm' | b'M' => 60, + b'h' | b'H' => 3_600, + b'd' | b'D' => 86_400, + b'w' | b'W' => 604_800, + _ => { + return Err(ControlConfigError::InvalidControlPersist(value.to_string())); + } + } + }; + let seconds = integer + .checked_mul(multiplier) + .ok_or_else(|| ControlConfigError::ControlPersistOverflow(value.to_string()))?; + total = total + .checked_add(seconds) + .ok_or_else(|| ControlConfigError::ControlPersistOverflow(value.to_string()))?; + if total > MAX_CONTROL_PERSIST_SECONDS { + return Err(ControlConfigError::ControlPersistOverflow( + value.to_string(), + )); + } + } + + Ok(total) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn control_master_modes_expose_runtime_decisions() { + let cases = [ + ("no", ControlMasterMode::No, true, false, false), + ("yes", ControlMasterMode::Yes, false, true, false), + ("ask", ControlMasterMode::Ask, false, true, true), + ("auto", ControlMasterMode::Auto, true, true, false), + ("autoask", ControlMasterMode::AutoAsk, true, true, true), + ("TRUE", ControlMasterMode::Yes, false, true, false), + ("FALSE", ControlMasterMode::No, true, false, false), + ]; + for (value, expected, tries, creates, asks) in cases { + let parsed = value.parse::().expect("valid mode"); + assert_eq!(parsed, expected); + assert_eq!(parsed.tries_existing(), tries); + assert_eq!(parsed.creates_master(), creates); + assert_eq!(parsed.requires_confirmation(), asks); + } + assert!("invalid".parse::().is_err()); + } + + #[test] + fn control_persist_parses_boolean_and_compound_forms() { + let cases = [ + ("no", ControlPersist::Disabled), + ("false", ControlPersist::Disabled), + ("yes", ControlPersist::Forever), + ("true", ControlPersist::Forever), + ("YES", ControlPersist::Forever), + ("0", ControlPersist::Forever), + ("30", ControlPersist::Timeout(Duration::from_secs(30))), + ("1m", ControlPersist::Timeout(Duration::from_secs(60))), + ("1h30m", ControlPersist::Timeout(Duration::from_secs(5_400))), + ( + "1W2d3H4m5S", + ControlPersist::Timeout(Duration::from_secs(788_645)), + ), + ("1s2s", ControlPersist::Timeout(Duration::from_secs(3))), + ]; + for (value, expected) in cases { + assert_eq!( + value.parse::().expect("valid persist"), + expected + ); + } + } + + #[test] + fn control_persist_rejects_malformed_and_overflowing_values() { + for value in ["", "maybe", "1x", "1.5m", "1.9s", "1.", ".", "-1"] { + assert!(value.parse::().is_err(), "accepted {value}"); + } + assert!("2147483648s".parse::().is_err()); + assert!( + "999999999999999999999999999999999999999w" + .parse::() + .is_err() + ); + } + + #[test] + fn command_scope_is_deliberately_bounded() { + for command in ["check", "forward", "cancel", "exit", "stop"] { + assert_eq!( + command + .parse::() + .expect("supported command") + .to_string(), + command + ); + } + for command in ["proxy", "channels", "conninfo", "", "CHECK"] { + assert!(command.parse::().is_err()); + } + } + + #[test] + fn raw_policy_distinguishes_defaults_and_disabled_path() { + assert_eq!( + ControlPolicy::from_raw(None, None, None).expect("defaults"), + ControlPolicy::default() + ); + let disabled = + ControlPolicy::from_raw(Some("auto"), Some("NoNe"), Some("5m")).expect("valid policy"); + assert_eq!(disabled.master, ControlMasterMode::Auto); + assert_eq!(disabled.path, None); + assert_eq!( + disabled.persist, + ControlPersist::Timeout(Duration::from_secs(300)) + ); + assert!(!disabled.is_socket_enabled()); + + let enabled = + ControlPolicy::from_raw(None, Some("~/.ssh/cm-%C"), None).expect("valid path"); + assert_eq!(enabled.path.as_deref(), Some("~/.ssh/cm-%C")); + assert!(enabled.is_socket_enabled()); + } +} diff --git a/src/ssh/control/mod.rs b/src/ssh/control/mod.rs new file mode 100644 index 00000000..8feb9b79 --- /dev/null +++ b/src/ssh/control/mod.rs @@ -0,0 +1,40 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +//! Typed connection-multiplexing configuration and wire primitives. +//! +//! This module deliberately contains no CLI, dispatcher, or live SSH runtime +//! integration. It is the policy and protocol boundary those layers use. + +mod config; +mod path; +mod protocol; +mod runtime; +mod socket; + +pub use config::{ + ControlCommand, ControlConfigError, ControlMasterMode, ControlPersist, ControlPolicy, +}; +pub use path::{ + ControlPathContext, ControlPathError, expand_control_path, unix_socket_path_capacity, + validate_control_socket_path, +}; +pub use protocol::{ + CONTROL_PROTOCOL_VERSION, ControlData, ControlDataStream, ControlMessage, ControlOperation, + ControlProtocolError, ControlRequest, ControlResponse, ControlResponseKind, + MAX_CONTROL_DATA_BYTES, MAX_CONTROL_FRAME_BYTES, SessionOpenRequest, read_control_message, + write_control_message, +}; +pub use runtime::{ + AttachOutcome, RunningControlMaster, attach_session, send_control_command, start_control_master, +}; +#[cfg(unix)] +pub use socket::verify_same_user; +pub use socket::{ + ControlSocketGuard, bind_control_socket, connect_control_socket, remove_stale_control_socket, +}; diff --git a/src/ssh/control/path.rs b/src/ssh/control/path.rs new file mode 100644 index 00000000..35564226 --- /dev/null +++ b/src/ssh/control/path.rs @@ -0,0 +1,297 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +use std::fmt::Write as _; +use std::path::{Path, PathBuf}; + +use sha1::{Digest, Sha1}; +use thiserror::Error; + +/// Deterministic values used to expand an OpenSSH `ControlPath` template. +/// +/// Callers supply local host and home values explicitly so tests do not depend +/// on process-global environment or host OS identity. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ControlPathContext { + local_host: String, + home_dir: PathBuf, + remote_host: String, + remote_port: u16, + remote_user: String, + jump_host: String, +} + +impl ControlPathContext { + /// Build a context for a direct connection. + pub fn new( + local_host: impl Into, + home_dir: impl Into, + remote_host: impl Into, + remote_port: u16, + remote_user: impl Into, + ) -> Self { + Self { + local_host: local_host.into(), + home_dir: home_dir.into(), + remote_host: remote_host.into(), + remote_port, + remote_user: remote_user.into(), + jump_host: String::new(), + } + } + + /// Include the resolved ProxyJump spelling in `%C`, matching OpenSSH's + /// connection-hash identity. + #[must_use] + pub fn with_jump_host(mut self, jump_host: impl Into) -> Self { + self.jump_host = jump_host.into(); + self + } + + /// Full local hostname used by `%l`. + #[must_use] + pub fn local_host(&self) -> &str { + &self.local_host + } + + /// Effective remote hostname used by `%h`. + #[must_use] + pub fn remote_host(&self) -> &str { + &self.remote_host + } + + /// Effective remote username used by `%r`. + #[must_use] + pub fn remote_user(&self) -> &str { + &self.remote_user + } + + /// Effective remote port used by `%p`. + #[must_use] + pub const fn remote_port(&self) -> u16 { + self.remote_port + } + + fn connection_hash(&self) -> String { + let mut digest = Sha1::new(); + digest.update(self.local_host.as_bytes()); + digest.update(self.remote_host.as_bytes()); + digest.update(self.remote_port.to_string().as_bytes()); + digest.update(self.remote_user.as_bytes()); + digest.update(self.jump_host.as_bytes()); + + let mut output = String::with_capacity(40); + for byte in digest.finalize() { + let _ = write!(output, "{byte:02x}"); + } + output + } +} + +/// Expand the #286 ControlPath token subset and validate the resulting Unix +/// socket address before any connection or bind attempt. +pub fn expand_control_path( + template: &str, + context: &ControlPathContext, +) -> Result { + if template.is_empty() { + return Err(ControlPathError::Empty); + } + if template.as_bytes().contains(&0) { + return Err(ControlPathError::ContainsNul); + } + + let expanded_home = if template == "~" { + context.home_dir.to_string_lossy().into_owned() + } else if let Some(suffix) = template.strip_prefix("~/") { + context.home_dir.join(suffix).to_string_lossy().into_owned() + } else if template.starts_with('~') { + return Err(ControlPathError::UnsupportedTilde(template.to_string())); + } else { + template.to_string() + }; + + let connection_hash = context.connection_hash(); + let port = context.remote_port.to_string(); + let mut output = String::with_capacity(expanded_home.len() + connection_hash.len()); + let mut characters = expanded_home.chars(); + while let Some(character) = characters.next() { + if character != '%' { + output.push(character); + continue; + } + let token = characters.next().ok_or(ControlPathError::IncompleteToken)?; + let replacement = match token { + '%' => "%", + 'C' => &connection_hash, + 'h' => &context.remote_host, + 'p' => &port, + 'r' => &context.remote_user, + 'l' => &context.local_host, + _ => return Err(ControlPathError::UnsupportedToken(token)), + }; + output.push_str(replacement); + } + + let path = PathBuf::from(output); + validate_control_socket_path(&path)?; + Ok(path) +} + +/// Validate an already-expanded control socket path. +pub fn validate_control_socket_path(path: &Path) -> Result<(), ControlPathError> { + if path.as_os_str().is_empty() { + return Err(ControlPathError::Empty); + } + + #[cfg(unix)] + { + use std::os::unix::ffi::OsStrExt as _; + + let bytes = path.as_os_str().as_bytes(); + if bytes.contains(&0) { + return Err(ControlPathError::ContainsNul); + } + let capacity = unix_socket_path_capacity().unwrap_or(0); + // sockaddr_un.sun_path must retain one byte for its terminating NUL. + let maximum = capacity.saturating_sub(1); + if bytes.len() > maximum { + return Err(ControlPathError::TooLong { + path: path.to_path_buf(), + length: bytes.len(), + maximum, + }); + } + } + + #[cfg(not(unix))] + if path.to_string_lossy().contains('\0') { + return Err(ControlPathError::ContainsNul); + } + + Ok(()) +} + +/// Number of bytes in `sockaddr_un.sun_path`, including its terminating NUL. +/// Non-Unix targets return `None` and remain compile-safe without claiming +/// Unix-domain socket support. +#[must_use] +pub const fn unix_socket_path_capacity() -> Option { + #[cfg(unix)] + { + Some( + std::mem::size_of::() + - std::mem::offset_of!(libc::sockaddr_un, sun_path), + ) + } + #[cfg(not(unix))] + { + None + } +} + +/// Invalid or unrepresentable `ControlPath` values. +#[derive(Debug, Clone, PartialEq, Eq, Error)] +#[non_exhaustive] +pub enum ControlPathError { + #[error("ControlPath must not be empty")] + Empty, + #[error("ControlPath contains a NUL byte")] + ContainsNul, + #[error("ControlPath ends with an incomplete '%' token")] + IncompleteToken, + #[error( + "ControlPath contains unsupported token '%{0}'; supported tokens are %C, %h, %p, %r, %l, and %%" + )] + UnsupportedToken(char), + #[error( + "ControlPath '{0}' uses unsupported '~user' expansion; use '~/' for the current user's home" + )] + UnsupportedTilde(String), + #[error( + "ControlPath '{}' is {length} bytes, but this platform permits at most {maximum}; shorten the directory or use %C", + path.display() + )] + TooLong { + path: PathBuf, + length: usize, + maximum: usize, + }, +} + +#[cfg(test)] +mod tests { + use super::*; + + fn context() -> ControlPathContext { + ControlPathContext::new( + "workstation.example", + "/home/alice", + "server.example", + 2222, + "deploy", + ) + .with_jump_host("jump.example:22") + } + + #[test] + fn expands_required_tokens_tilde_and_literal_percent() { + let short_context = ControlPathContext::new("local", "/h", "host", 22, "user"); + let expanded = expand_control_path("~/cm-%h-%p-%r-%l-%%-%C", &short_context) + .expect("valid control path"); + let rendered = expanded.to_string_lossy(); + assert!(rendered.starts_with("/h/cm-host-22-user-local-%-")); + let hash = rendered.rsplit('-').next().expect("hash suffix"); + assert_eq!(hash.len(), 40); + assert!(hash.bytes().all(|byte| byte.is_ascii_hexdigit())); + } + + #[test] + fn connection_hash_is_stable_and_includes_jump_host() { + let direct = ControlPathContext::new("local", "/home/test", "remote", 22, "user"); + let direct_hash = expand_control_path("%C", &direct).expect("direct hash"); + let repeated = expand_control_path("%C", &direct).expect("repeated hash"); + assert_eq!(direct_hash, repeated); + + let jumped = direct.clone().with_jump_host("jump"); + assert_ne!( + direct_hash, + expand_control_path("%C", &jumped).expect("jump hash") + ); + } + + #[test] + fn rejects_ambiguous_or_malformed_templates() { + for template in ["", "~other/socket", "/tmp/%", "/tmp/%x", "/tmp/a\0b"] { + assert!( + expand_control_path(template, &context()).is_err(), + "accepted {template:?}" + ); + } + } + + #[cfg(unix)] + #[test] + fn reports_platform_socket_limit_before_bind() { + let capacity = unix_socket_path_capacity().expect("Unix capacity"); + let valid = PathBuf::from("a".repeat(capacity - 1)); + validate_control_socket_path(&valid).expect("maximum path"); + + let too_long = PathBuf::from("a".repeat(capacity)); + let error = validate_control_socket_path(&too_long).expect_err("overlong path"); + assert!(matches!( + error, + ControlPathError::TooLong { + length, + maximum, + .. + } if length == capacity && maximum == capacity - 1 + )); + assert!(error.to_string().contains("use %C")); + } +} diff --git a/src/ssh/control/protocol.rs b/src/ssh/control/protocol.rs new file mode 100644 index 00000000..34542798 --- /dev/null +++ b/src/ssh/control/protocol.rs @@ -0,0 +1,391 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +use std::io; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; + +use crate::forwarding::ForwardingDirective; +use crate::ssh::SessionPolicy; +use crate::ssh::tokio_client::AddressFamily; + +use super::ControlCommand; + +/// Version of bssh's length-prefixed JSON multiplexing protocol. +pub const CONTROL_PROTOCOL_VERSION: u32 = 1; +/// Largest accepted serialized control message. +pub const MAX_CONTROL_FRAME_BYTES: usize = 1024 * 1024; +/// Largest binary chunk carried in one data message. +pub const MAX_CONTROL_DATA_BYTES: usize = 64 * 1024; +const MAX_CONTROL_ENVIRONMENT_ENTRIES: usize = 4_096; + +/// One framed message on a bssh control socket. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] +#[serde(tag = "type", content = "body", rename_all = "snake_case")] +pub enum ControlMessage { + Request(ControlRequest), + Response(ControlResponse), + Data(ControlData), +} + +impl ControlMessage { + fn validate(&self) -> Result<(), ControlProtocolError> { + match self { + Self::Request(request) => request.validate(), + Self::Response(_) => Ok(()), + Self::Data(data) => data.validate(), + } + } +} + +/// Correlated client-to-master request. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ControlRequest { + pub request_id: u64, + pub operation: ControlOperation, +} + +impl ControlRequest { + /// Validate semantic limits that cannot be expressed by the JSON frame + /// length alone. + pub fn validate(&self) -> Result<(), ControlProtocolError> { + match &self.operation { + ControlOperation::Hello { version } if *version != CONTROL_PROTOCOL_VERSION => { + Err(ControlProtocolError::UnsupportedVersion { + received: *version, + supported: CONTROL_PROTOCOL_VERSION, + }) + } + ControlOperation::Hello { .. } => Ok(()), + ControlOperation::OpenSession(session) => session.validate(), + ControlOperation::Command { + command, forwards, .. + } => match command { + ControlCommand::Forward | ControlCommand::Cancel if forwards.is_empty() => { + Err(ControlProtocolError::MissingForwarding(*command)) + } + ControlCommand::Check | ControlCommand::Exit | ControlCommand::Stop + if !forwards.is_empty() => + { + Err(ControlProtocolError::UnexpectedForwarding(*command)) + } + _ => Ok(()), + }, + } + } +} + +/// Request operation understood by a bssh control master. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] +#[serde(tag = "operation", rename_all = "snake_case")] +pub enum ControlOperation { + Hello { + version: u32, + }, + OpenSession(SessionOpenRequest), + Command { + command: ControlCommand, + /// Raw directives in original CLI/config order. + forwards: Vec, + /// Address family used when the master parses raw forwarding specs. + address_family: AddressFamily, + }, +} + +/// Session policy transferred to the authenticated control master. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SessionOpenRequest { + pub policy: SessionPolicy, + /// Client-side terminal name used when `policy.request_pty` is true. + pub terminal: Option, +} + +impl SessionOpenRequest { + /// Build a wire request. `LocalCommand` must execute in the invoking client + /// process, never in the long-lived master. + pub fn new( + policy: SessionPolicy, + terminal: Option, + ) -> Result { + let request = Self { policy, terminal }; + request.validate()?; + Ok(request) + } + + /// Enforce limits before sending or accepting a session request. + pub fn validate(&self) -> Result<(), ControlProtocolError> { + if self.policy.local_command.is_some() { + return Err(ControlProtocolError::LocalCommandNotAllowed); + } + if self.policy.environment.len() > MAX_CONTROL_ENVIRONMENT_ENTRIES { + return Err(ControlProtocolError::TooManyEnvironmentEntries { + count: self.policy.environment.len(), + maximum: MAX_CONTROL_ENVIRONMENT_ENTRIES, + }); + } + Ok(()) + } +} + +/// Correlated master-to-client response. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ControlResponse { + pub request_id: u64, + pub kind: ControlResponseKind, +} + +/// Result variants needed by #286's session and five control commands. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] +#[serde(tag = "result", rename_all = "snake_case")] +pub enum ControlResponseKind { + Ok, + Alive { pid: u32 }, + SessionOpened { session_id: u64 }, + ExitStatus { session_id: u64, status: u32 }, + Error { code: String, message: String }, +} + +/// Byte stream represented by a data message. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[non_exhaustive] +#[serde(rename_all = "snake_case")] +pub enum ControlDataStream { + Stdin, + Stdout, + Stderr, +} + +/// Ordered session data or EOF indication. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ControlData { + pub session_id: u64, + pub stream: ControlDataStream, + pub sequence: u64, + pub payload: Vec, + pub eof: bool, +} + +impl ControlData { + /// Validate the per-message binary payload limit. + pub fn validate(&self) -> Result<(), ControlProtocolError> { + if self.payload.len() > MAX_CONTROL_DATA_BYTES { + return Err(ControlProtocolError::DataTooLarge { + length: self.payload.len(), + maximum: MAX_CONTROL_DATA_BYTES, + }); + } + Ok(()) + } +} + +/// Serialize and write one bounded big-endian-length-prefixed JSON message. +pub async fn write_control_message( + writer: &mut W, + message: &ControlMessage, +) -> Result<(), ControlProtocolError> +where + W: AsyncWrite + Unpin, +{ + message.validate()?; + let encoded = serde_json::to_vec(message)?; + if encoded.len() > MAX_CONTROL_FRAME_BYTES { + return Err(ControlProtocolError::FrameTooLarge { + length: encoded.len(), + maximum: MAX_CONTROL_FRAME_BYTES, + }); + } + let length = u32::try_from(encoded.len()).map_err(|_| ControlProtocolError::FrameTooLarge { + length: encoded.len(), + maximum: MAX_CONTROL_FRAME_BYTES, + })?; + writer.write_all(&length.to_be_bytes()).await?; + writer.write_all(&encoded).await?; + writer.flush().await?; + Ok(()) +} + +/// Read, decode, and validate one bounded big-endian-length-prefixed JSON +/// message without allocating an attacker-declared oversized frame. +pub async fn read_control_message(reader: &mut R) -> Result +where + R: AsyncRead + Unpin, +{ + let mut header = [0u8; 4]; + reader.read_exact(&mut header).await?; + let length = u32::from_be_bytes(header) as usize; + if length == 0 { + return Err(ControlProtocolError::EmptyFrame); + } + if length > MAX_CONTROL_FRAME_BYTES { + return Err(ControlProtocolError::FrameTooLarge { + length, + maximum: MAX_CONTROL_FRAME_BYTES, + }); + } + + let mut encoded = vec![0u8; length]; + reader.read_exact(&mut encoded).await?; + let message = serde_json::from_slice::(&encoded)?; + message.validate()?; + Ok(message) +} + +/// Framing, serialization, version, or semantic-limit failure. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum ControlProtocolError { + #[error("control socket I/O failed: {0}")] + Io(#[from] io::Error), + #[error("control message JSON is invalid: {0}")] + Json(#[from] serde_json::Error), + #[error("control protocol frame must not be empty")] + EmptyFrame, + #[error("control protocol frame is {length} bytes, exceeding the {maximum}-byte limit")] + FrameTooLarge { length: usize, maximum: usize }, + #[error("control data chunk is {length} bytes, exceeding the {maximum}-byte limit")] + DataTooLarge { length: usize, maximum: usize }, + #[error("control protocol version {received} is unsupported; this build supports {supported}")] + UnsupportedVersion { received: u32, supported: u32 }, + #[error( + "LocalCommand must execute in the invoking process and cannot be sent to a control master" + )] + LocalCommandNotAllowed, + #[error("session request contains {count} environment entries; maximum is {maximum}")] + TooManyEnvironmentEntries { count: usize, maximum: usize }, + #[error("control command '{0}' requires at least one forwarding directive")] + MissingForwarding(ControlCommand), + #[error("control command '{0}' does not accept forwarding directives")] + UnexpectedForwarding(ControlCommand), +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ssh::SessionRequest; + + fn session_policy() -> SessionPolicy { + SessionPolicy { + environment: vec![("LANG".into(), "C.UTF-8".into())], + local_command: None, + request_pty: false, + stdin_null: false, + request: SessionRequest::Exec("printf test".into()), + } + } + + #[tokio::test] + async fn request_round_trip_preserves_session_policy() { + let request = ControlMessage::Request(ControlRequest { + request_id: 7, + operation: ControlOperation::OpenSession( + SessionOpenRequest::new(session_policy(), Some("xterm-256color".into())) + .expect("wire-safe session"), + ), + }); + let (mut client, mut server) = tokio::io::duplex(16 * 1024); + let write = write_control_message(&mut client, &request); + let read = read_control_message(&mut server); + let (written, decoded) = tokio::join!(write, read); + written.expect("write succeeds"); + assert_eq!(decoded.expect("read succeeds"), request); + } + + #[tokio::test] + async fn data_round_trip_is_binary_safe() { + let message = ControlMessage::Data(ControlData { + session_id: 9, + stream: ControlDataStream::Stdout, + sequence: 3, + payload: vec![0, 255, b'\n', 0], + eof: true, + }); + let (mut client, mut server) = tokio::io::duplex(16 * 1024); + let write = write_control_message(&mut client, &message); + let read = read_control_message(&mut server); + let (written, decoded) = tokio::join!(write, read); + written.expect("write succeeds"); + assert_eq!(decoded.expect("read succeeds"), message); + } + + #[tokio::test] + async fn oversized_declared_frame_is_rejected_before_payload_read() { + let (mut client, mut server) = tokio::io::duplex(16); + client + .write_all(&((MAX_CONTROL_FRAME_BYTES as u32) + 1).to_be_bytes()) + .await + .expect("header write"); + let error = read_control_message(&mut server) + .await + .expect_err("oversized frame"); + assert!(matches!(error, ControlProtocolError::FrameTooLarge { .. })); + } + + #[tokio::test] + async fn oversized_data_is_rejected_before_serialization() { + let message = ControlMessage::Data(ControlData { + session_id: 1, + stream: ControlDataStream::Stdin, + sequence: 0, + payload: vec![0; MAX_CONTROL_DATA_BYTES + 1], + eof: false, + }); + let mut sink = tokio::io::sink(); + let error = write_control_message(&mut sink, &message) + .await + .expect_err("oversized data"); + assert!(matches!(error, ControlProtocolError::DataTooLarge { .. })); + } + + #[test] + fn local_command_cannot_cross_the_master_boundary() { + let mut policy = session_policy(); + policy.local_command = Some("touch /tmp/client-only".into()); + assert!(matches!( + SessionOpenRequest::new(policy, None), + Err(ControlProtocolError::LocalCommandNotAllowed) + )); + } + + #[test] + fn command_forwarding_shape_is_validated() { + let missing = ControlRequest { + request_id: 1, + operation: ControlOperation::Command { + command: ControlCommand::Forward, + forwards: Vec::new(), + address_family: AddressFamily::Any, + }, + }; + assert!(matches!( + missing.validate(), + Err(ControlProtocolError::MissingForwarding( + ControlCommand::Forward + )) + )); + + let unexpected = ControlRequest { + request_id: 2, + operation: ControlOperation::Command { + command: ControlCommand::Check, + forwards: vec![ForwardingDirective::Dynamic("1080".into())], + address_family: AddressFamily::Any, + }, + }; + assert!(matches!( + unexpected.validate(), + Err(ControlProtocolError::UnexpectedForwarding( + ControlCommand::Check + )) + )); + } +} diff --git a/src/ssh/control/runtime.rs b/src/ssh/control/runtime.rs new file mode 100644 index 00000000..16a05b4e --- /dev/null +++ b/src/ssh/control/runtime.rs @@ -0,0 +1,903 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. + +//! Live bssh control-master runtime. + +#[cfg(unix)] +mod unix { + use std::io::{self, BufRead as _, Write as _}; + use std::os::fd::AsRawFd as _; + use std::path::Path; + use std::sync::Arc; + use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; + use std::time::Duration; + + use anyhow::{Context, Result}; + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + use tokio::net::{UnixListener, UnixStream}; + use tokio::sync::{Mutex, Semaphore, mpsc, watch}; + use tokio::task::{JoinHandle, JoinSet}; + + use crate::forwarding::ForwardingPlan; + use crate::ssh::tokio_client::{Client, CommandOutput}; + + use super::super::{ + CONTROL_PROTOCOL_VERSION, ControlCommand, ControlData, ControlDataStream, ControlMessage, + ControlOperation, ControlPersist, ControlRequest, ControlResponse, ControlResponseKind, + MAX_CONTROL_DATA_BYTES, SessionOpenRequest, bind_control_socket, connect_control_socket, + read_control_message, verify_same_user, write_control_message, + }; + + const HELLO_REQUEST_ID: u64 = 0; + const OPERATION_REQUEST_ID: u64 = 1; + const SESSION_PIPE_CAPACITY: usize = 256 * 1024; + const CONTROL_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(3); + const CONTROL_CONFIRM_TIMEOUT_MILLIS: libc::c_int = 30_000; + const MAX_CONTROL_CLIENTS: usize = 128; + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub enum AttachOutcome { + NoMaster, + ExitStatus(u32), + } + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum MasterSignal { + Exit, + Stop, + } + + pub struct RunningControlMaster { + signal: mpsc::Sender, + active_sessions: watch::Receiver, + task: JoinHandle>, + } + + impl RunningControlMaster { + pub async fn shutdown_immediately(self) -> Result<()> { + let _ = self.signal.send(MasterSignal::Exit).await; + join_master(self.task).await + } + + /// Apply ControlPersist after the initial invocation finishes. + pub async fn finish_after_initial( + mut self, + persist: ControlPersist, + initial_request_was_none: bool, + ) -> Result<()> { + if matches!(persist, ControlPersist::Disabled) && !initial_request_was_none { + let _ = self.signal.send(MasterSignal::Stop).await; + return join_master(self.task).await; + } + let Some(timeout) = persist.timeout() else { + return join_master(self.task).await; + }; + loop { + if *self.active_sessions.borrow() != 0 { + tokio::select! { + result = &mut self.task => return flatten_master_join(result), + changed = self.active_sessions.changed() => { + if changed.is_err() { + return join_master(self.task).await; + } + } + } + continue; + } + let idle = tokio::time::sleep(timeout); + tokio::pin!(idle); + tokio::select! { + result = &mut self.task => return flatten_master_join(result), + changed = self.active_sessions.changed() => { + if changed.is_err() { + return join_master(self.task).await; + } + } + () = &mut idle => { + let _ = self.signal.send(MasterSignal::Stop).await; + return join_master(self.task).await; + } + } + } + } + } + + async fn join_master(task: JoinHandle>) -> Result<()> { + flatten_master_join(task.await) + } + + fn flatten_master_join(result: Result, tokio::task::JoinError>) -> Result<()> { + result.context("Control master task failed to join")? + } + + pub fn start_control_master( + path: &Path, + client: Client, + require_confirmation: bool, + ) -> Result { + let (listener, guard) = bind_control_socket(path)?; + let (signal_tx, signal_rx) = mpsc::channel(8); + let (active_tx, active_rx) = watch::channel(0usize); + let task_signal = signal_tx.clone(); + let task = tokio::spawn(async move { + serve_control_master( + listener, + guard, + client, + require_confirmation, + signal_rx, + task_signal, + active_tx, + ) + .await + }); + Ok(RunningControlMaster { + signal: signal_tx, + active_sessions: active_rx, + task, + }) + } + + pub async fn attach_session( + path: &Path, + session: SessionOpenRequest, + invoking_policy: &crate::ssh::SessionPolicy, + ) -> Result { + let mut stream = match tokio::time::timeout( + CONTROL_HANDSHAKE_TIMEOUT, + connect_control_socket(path), + ) + .await + { + Ok(Ok(stream)) => stream, + Err(_) => { + tracing::debug!(path = %path.display(), "Control master connect timed out; falling back"); + return Ok(AttachOutcome::NoMaster); + } + Ok(Err(error)) + if matches!( + error.kind(), + io::ErrorKind::NotFound | io::ErrorKind::ConnectionRefused + ) => + { + return Ok(AttachOutcome::NoMaster); + } + Ok(Err(error)) => { + return Err(error).with_context(|| { + format!("Could not connect to control master '{}'", path.display()) + }); + } + }; + match tokio::time::timeout(CONTROL_HANDSHAKE_TIMEOUT, perform_hello(&mut stream)).await { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::debug!(path = %path.display(), "Control master handshake failed; falling back: {error:#}"); + return Ok(AttachOutcome::NoMaster); + } + Err(_) => { + tracing::debug!(path = %path.display(), "Control master handshake timed out; falling back"); + return Ok(AttachOutcome::NoMaster); + } + } + invoking_policy.run_local_command().await?; + run_attached_session(stream, session) + .await + .map(AttachOutcome::ExitStatus) + } + + pub async fn send_control_command( + path: &Path, + command: ControlCommand, + forwards: Vec, + address_family: crate::ssh::tokio_client::AddressFamily, + ) -> Result { + let mut stream = + tokio::time::timeout(CONTROL_HANDSHAKE_TIMEOUT, connect_control_socket(path)) + .await + .context("Timed out connecting to control master")? + .with_context(|| format!("No control master is running at '{}'", path.display()))?; + tokio::time::timeout(CONTROL_HANDSHAKE_TIMEOUT, perform_hello(&mut stream)) + .await + .context("Timed out negotiating with control master")??; + write_control_message( + &mut stream, + &ControlMessage::Request(ControlRequest { + request_id: OPERATION_REQUEST_ID, + operation: ControlOperation::Command { + command, + forwards, + address_family, + }, + }), + ) + .await?; + let response = read_response(&mut stream, OPERATION_REQUEST_ID).await?; + match response { + ControlResponseKind::Error { code, message } => { + anyhow::bail!("control command failed ({code}): {message}") + } + response => Ok(response), + } + } + + async fn run_attached_session( + mut stream: UnixStream, + session: SessionOpenRequest, + ) -> Result { + write_control_message( + &mut stream, + &ControlMessage::Request(ControlRequest { + request_id: OPERATION_REQUEST_ID, + operation: ControlOperation::OpenSession(session.clone()), + }), + ) + .await?; + let opened = read_response(&mut stream, OPERATION_REQUEST_ID).await?; + let session_id = match opened { + ControlResponseKind::SessionOpened { session_id } => session_id, + ControlResponseKind::Error { code, message } => { + anyhow::bail!("control master rejected session ({code}): {message}") + } + response => { + anyhow::bail!("unexpected control response while opening session: {response:?}") + } + }; + let (mut reader, mut writer) = stream.into_split(); + let stdin_null = session.policy.stdin_null; + let input = tokio::spawn(async move { + let mut sequence = 0u64; + if !stdin_null { + let mut stdin = tokio::io::stdin(); + let mut buffer = vec![0u8; MAX_CONTROL_DATA_BYTES]; + loop { + let count = stdin.read(&mut buffer).await?; + if count == 0 { + break; + } + write_control_message( + &mut writer, + &ControlMessage::Data(ControlData { + session_id, + stream: ControlDataStream::Stdin, + sequence, + payload: buffer[..count].to_vec(), + eof: false, + }), + ) + .await + .map_err(io::Error::other)?; + sequence = sequence.saturating_add(1); + } + } + write_control_message( + &mut writer, + &ControlMessage::Data(ControlData { + session_id, + stream: ControlDataStream::Stdin, + sequence, + payload: Vec::new(), + eof: true, + }), + ) + .await + .map_err(io::Error::other) + }); + + let mut stdout = tokio::io::stdout(); + let mut stderr = tokio::io::stderr(); + let mut stdout_open = true; + let mut stderr_open = true; + let status = loop { + match read_control_message(&mut reader).await? { + ControlMessage::Data(data) if data.session_id == session_id => match data.stream { + ControlDataStream::Stdout if stdout_open => { + if stdout.write_all(&data.payload).await.is_err() { + stdout_open = false; + } else { + stdout.flush().await.ok(); + } + } + ControlDataStream::Stderr if stderr_open => { + if stderr.write_all(&data.payload).await.is_err() { + stderr_open = false; + } else { + stderr.flush().await.ok(); + } + } + _ => {} + }, + ControlMessage::Response(ControlResponse { + request_id: OPERATION_REQUEST_ID, + kind: + ControlResponseKind::ExitStatus { + session_id: response_session, + status, + }, + }) if response_session == session_id => break status, + ControlMessage::Response(ControlResponse { + request_id: OPERATION_REQUEST_ID, + kind: ControlResponseKind::Error { code, message }, + }) => anyhow::bail!("control session failed ({code}): {message}"), + message => anyhow::bail!("unexpected control session message: {message:?}"), + } + }; + input.abort(); + let _ = input.await; + Ok(status) + } + + async fn perform_hello(stream: &mut UnixStream) -> Result<()> { + write_control_message( + stream, + &ControlMessage::Request(ControlRequest { + request_id: HELLO_REQUEST_ID, + operation: ControlOperation::Hello { + version: CONTROL_PROTOCOL_VERSION, + }, + }), + ) + .await?; + match read_response(stream, HELLO_REQUEST_ID).await? { + ControlResponseKind::Ok => Ok(()), + ControlResponseKind::Error { code, message } => { + anyhow::bail!("control protocol handshake failed ({code}): {message}") + } + response => anyhow::bail!("unexpected control handshake response: {response:?}"), + } + } + + async fn read_response( + stream: &mut UnixStream, + request_id: u64, + ) -> Result { + match read_control_message(stream).await? { + ControlMessage::Response(response) if response.request_id == request_id => { + Ok(response.kind) + } + message => anyhow::bail!("unexpected control response: {message:?}"), + } + } + + #[allow(clippy::too_many_arguments)] + async fn serve_control_master( + listener: UnixListener, + guard: super::super::ControlSocketGuard, + client: Client, + require_confirmation: bool, + mut signal_rx: mpsc::Receiver, + signal_tx: mpsc::Sender, + active_tx: watch::Sender, + ) -> Result<()> { + let next_session = Arc::new(AtomicU64::new(1)); + let active = Arc::new(AtomicUsize::new(0)); + let confirmation = Arc::new(Mutex::new(())); + let handler_slots = Arc::new(Semaphore::new(MAX_CONTROL_CLIENTS)); + let mut handlers = JoinSet::new(); + let immediate = loop { + tokio::select! { + signal = signal_rx.recv() => match signal { + Some(MasterSignal::Exit) => break true, + Some(MasterSignal::Stop) | None => break false, + }, + accepted = listener.accept() => { + let (stream, _) = accepted.context("Control socket accept failed")?; + if let Err(error) = verify_same_user(&stream) { + tracing::warn!("Rejected control client: {error}"); + continue; + } + let Ok(handler_slot) = Arc::clone(&handler_slots).try_acquire_owned() else { + tracing::warn!("Rejected control client: {MAX_CONTROL_CLIENTS} handlers are already active"); + continue; + }; + let handler_client = client.clone(); + let handler_signal = signal_tx.clone(); + let handler_active = Arc::clone(&active); + let handler_active_tx = active_tx.clone(); + let handler_next_session = Arc::clone(&next_session); + let handler_confirmation = Arc::clone(&confirmation); + handlers.spawn(async move { + let _handler_slot = handler_slot; + if let Err(error) = handle_control_connection( + stream, + handler_client, + require_confirmation, + handler_confirmation, + handler_signal, + handler_active, + handler_active_tx, + handler_next_session, + ) + .await + { + tracing::debug!("Control client disconnected with error: {error:#}"); + } + }); + } + Some(joined) = handlers.join_next(), if !handlers.is_empty() => { + if let Err(error) = joined { + tracing::warn!("Control client task failed to join: {error}"); + } + } + () = tokio::time::sleep(Duration::from_secs(1)) => { + if client.is_closed() { + break true; + } + } + } + }; + drop(listener); + drop(guard); + if immediate { + handlers.abort_all(); + } + while let Some(result) = handlers.join_next().await { + if !immediate && let Err(error) = result { + tracing::warn!("Control client task failed while draining: {error}"); + } + } + client + .disconnect() + .await + .context("Could not disconnect control-master SSH transport") + } + + #[allow(clippy::too_many_arguments)] + async fn handle_control_connection( + mut stream: UnixStream, + client: Client, + require_confirmation: bool, + confirmation: Arc>, + signal: mpsc::Sender, + active: Arc, + active_tx: watch::Sender, + next_session: Arc, + ) -> Result<()> { + let hello = read_initial_control_message(&mut stream, "hello").await?; + let request_id = match hello { + ControlMessage::Request(ControlRequest { + request_id, + operation: ControlOperation::Hello { version }, + }) if version == CONTROL_PROTOCOL_VERSION => request_id, + message => anyhow::bail!("expected protocol hello, received {message:?}"), + }; + send_response(&mut stream, request_id, ControlResponseKind::Ok).await?; + let request = match read_initial_control_message(&mut stream, "request").await? { + ControlMessage::Request(request) => request, + message => anyhow::bail!("expected control request, received {message:?}"), + }; + match request.operation { + ControlOperation::Hello { .. } => { + send_error( + &mut stream, + request.request_id, + "duplicate_hello", + "hello was already completed", + ) + .await + } + ControlOperation::OpenSession(session) => { + if require_confirmation { + let approved = confirm_attach(Arc::clone(&confirmation)).await?; + if !approved { + return send_error( + &mut stream, + request.request_id, + "permission_denied", + "control master did not approve the shared session", + ) + .await; + } + } + handle_remote_session( + stream, + request.request_id, + session, + client, + active, + active_tx, + next_session.fetch_add(1, Ordering::Relaxed), + ) + .await + } + ControlOperation::Command { + command, + forwards, + address_family, + } => { + let plan = ForwardingPlan { + directives: forwards, + clear_all: false, + exit_on_failure: true, + address_family, + }; + let result = match command { + ControlCommand::Check => { + return send_response( + &mut stream, + request.request_id, + ControlResponseKind::Alive { + pid: std::process::id(), + }, + ) + .await; + } + ControlCommand::Forward => client.add_control_forwardings(&plan).await, + ControlCommand::Cancel => client.cancel_control_forwardings(&plan).await, + ControlCommand::Exit | ControlCommand::Stop => Ok(()), + }; + if let Err(error) = result { + return send_error( + &mut stream, + request.request_id, + "forwarding_failed", + &error.to_string(), + ) + .await; + } + send_response(&mut stream, request.request_id, ControlResponseKind::Ok).await?; + match command { + ControlCommand::Exit => { + let _ = signal.send(MasterSignal::Exit).await; + } + ControlCommand::Stop => { + let _ = signal.send(MasterSignal::Stop).await; + } + _ => {} + } + Ok(()) + } + } + } + + async fn read_initial_control_message( + stream: &mut UnixStream, + phase: &str, + ) -> Result { + tokio::time::timeout(CONTROL_HANDSHAKE_TIMEOUT, read_control_message(stream)) + .await + .with_context(|| format!("Control client timed out during {phase}"))? + .map_err(anyhow::Error::from) + } + + async fn confirm_attach(lock: Arc>) -> Result { + let _guard = lock.lock().await; + tokio::task::spawn_blocking(|| -> Result { + let tty = std::fs::OpenOptions::new() + .read(true) + .write(true) + .open("/dev/tty") + .context("ControlMaster ask mode requires /dev/tty")?; + let mut writer = tty.try_clone().context("Could not clone /dev/tty")?; + writer.write_all(b"Allow shared SSH session? [y/N] ")?; + writer.flush()?; + let mut descriptor = libc::pollfd { + fd: tty.as_raw_fd(), + events: libc::POLLIN, + revents: 0, + }; + // SAFETY: `descriptor` points to one initialized pollfd whose file + // descriptor remains owned by `tty` for the duration of this call. + let ready = unsafe { + libc::poll( + std::ptr::from_mut(&mut descriptor), + 1, + CONTROL_CONFIRM_TIMEOUT_MILLIS, + ) + }; + if ready < 0 { + return Err(io::Error::last_os_error()).context("Could not poll /dev/tty"); + } + if ready == 0 || descriptor.revents & libc::POLLIN == 0 { + return Ok(false); + } + let mut answer = String::new(); + io::BufReader::new(tty).read_line(&mut answer)?; + Ok(matches!( + answer.trim().to_ascii_lowercase().as_str(), + "y" | "yes" + )) + }) + .await + .context("Control-master confirmation task failed")? + } + + struct ActiveSessionGuard { + active: Arc, + updates: watch::Sender, + } + + impl ActiveSessionGuard { + fn new(active: Arc, updates: watch::Sender) -> Self { + let count = active.fetch_add(1, Ordering::AcqRel) + 1; + updates.send_replace(count); + Self { active, updates } + } + } + + impl Drop for ActiveSessionGuard { + fn drop(&mut self) { + let count = self.active.fetch_sub(1, Ordering::AcqRel) - 1; + self.updates.send_replace(count); + } + } + + async fn handle_remote_session( + mut stream: UnixStream, + request_id: u64, + session: SessionOpenRequest, + client: Client, + active: Arc, + active_tx: watch::Sender, + session_id: u64, + ) -> Result<()> { + let _active = ActiveSessionGuard::new(active, active_tx); + send_response( + &mut stream, + request_id, + ControlResponseKind::SessionOpened { session_id }, + ) + .await?; + let (mut reader, mut writer) = stream.into_split(); + let (input_reader, mut input_writer) = tokio::io::duplex(SESSION_PIPE_CAPACITY); + let input = tokio::spawn(async move { + loop { + match read_control_message(&mut reader).await? { + ControlMessage::Data(data) + if data.session_id == session_id + && matches!(data.stream, ControlDataStream::Stdin) => + { + if !data.payload.is_empty() { + input_writer.write_all(&data.payload).await?; + } + if data.eof { + input_writer.shutdown().await?; + return Ok::<(), anyhow::Error>(()); + } + } + message => anyhow::bail!("unexpected session input message: {message:?}"), + } + } + }); + let (output_tx, mut output_rx) = mpsc::channel(128); + let policy = session.policy; + let terminal = session.terminal; + let execution_client = client.clone(); + let mut execution = tokio::spawn(async move { + execution_client + .execute_session_streaming_with_input_and_terminal( + &policy, + terminal.as_deref(), + output_tx, + input_reader, + ) + .await + }); + let mut status = None; + let mut execution_joined = false; + let mut output_open = true; + let mut sequence = 0u64; + let relay_result = async { + loop { + if status.is_some() && !output_open { + break; + } + tokio::select! { + result = &mut execution, if !execution_joined => { + execution_joined = true; + status = Some(result.context("Remote session task failed to join")??); + } + output = output_rx.recv(), if output_open => match output { + Some(CommandOutput::StdOut(payload)) => { + write_control_message( + &mut writer, + &ControlMessage::Data(ControlData { + session_id, + stream: ControlDataStream::Stdout, + sequence, + payload: payload.to_vec(), + eof: false, + }), + ).await?; + sequence = sequence.saturating_add(1); + } + Some(CommandOutput::StdErr(payload)) => { + write_control_message( + &mut writer, + &ControlMessage::Data(ControlData { + session_id, + stream: ControlDataStream::Stderr, + sequence, + payload: payload.to_vec(), + eof: false, + }), + ).await?; + sequence = sequence.saturating_add(1); + } + Some(CommandOutput::ExitCode(_)) => {} + None => output_open = false, + } + } + } + Ok::<(), anyhow::Error>(()) + } + .await; + input.abort(); + let _ = input.await; + if !execution_joined { + execution.abort(); + let _ = execution.await; + } + relay_result?; + send_response( + &mut writer, + request_id, + ControlResponseKind::ExitStatus { + session_id, + status: status.context("remote session ended without exit status")?, + }, + ) + .await + } + + async fn send_response( + writer: &mut W, + request_id: u64, + kind: ControlResponseKind, + ) -> Result<()> + where + W: tokio::io::AsyncWrite + Unpin, + { + write_control_message( + writer, + &ControlMessage::Response(ControlResponse { request_id, kind }), + ) + .await?; + Ok(()) + } + + async fn send_error(writer: &mut W, request_id: u64, code: &str, message: &str) -> Result<()> + where + W: tokio::io::AsyncWrite + Unpin, + { + send_response( + writer, + request_id, + ControlResponseKind::Error { + code: code.to_string(), + message: message.to_string(), + }, + ) + .await + } + + #[cfg(test)] + mod tests { + use tempfile::TempDir; + + use crate::ssh::{SessionPolicy, SessionRequest}; + + use super::*; + + fn no_session_request() -> SessionOpenRequest { + SessionOpenRequest::new( + SessionPolicy { + environment: Vec::new(), + local_command: None, + request_pty: false, + stdin_null: true, + request: SessionRequest::None, + }, + None, + ) + .expect("valid session request") + } + + #[tokio::test(start_paused = true)] + async fn silent_control_socket_falls_back_after_bounded_handshake() { + let directory = TempDir::new().expect("temporary directory"); + let path = directory.path().join("silent-control"); + let (listener, _guard) = bind_control_socket(&path).expect("control listener"); + let silent = tokio::spawn(async move { + let (_stream, _) = listener.accept().await.expect("accepted control client"); + std::future::pending::<()>().await; + }); + + let request = no_session_request(); + let outcome = attach_session(&path, request.clone(), &request.policy) + .await + .expect("silent master should fall back"); + assert_eq!(outcome, AttachOutcome::NoMaster); + silent.abort(); + let _ = silent.await; + } + + #[tokio::test] + async fn master_dying_during_hello_falls_back_without_replaying_a_session() { + let directory = TempDir::new().expect("temporary directory"); + let path = directory.path().join("dead-control"); + let (listener, _guard) = bind_control_socket(&path).expect("control listener"); + let dying = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accepted control client"); + drop(stream); + }); + + let request = no_session_request(); + let outcome = attach_session(&path, request.clone(), &request.policy) + .await + .expect("dead master should fall back"); + assert_eq!(outcome, AttachOutcome::NoMaster); + dying.await.expect("dying listener task"); + } + } +} + +#[cfg(unix)] +pub use unix::{ + AttachOutcome, RunningControlMaster, attach_session, send_control_command, start_control_master, +}; + +#[cfg(not(unix))] +mod unsupported { + use std::path::Path; + + use anyhow::Result; + + use crate::forwarding::ForwardingDirective; + use crate::ssh::tokio_client::{AddressFamily, Client}; + + use super::super::{ControlCommand, ControlPersist, ControlResponseKind, SessionOpenRequest}; + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub enum AttachOutcome { + NoMaster, + ExitStatus(u32), + } + + pub struct RunningControlMaster; + + impl RunningControlMaster { + pub async fn shutdown_immediately(self) -> Result<()> { + anyhow::bail!("connection multiplexing requires Unix-domain sockets") + } + + pub async fn finish_after_initial( + self, + _persist: ControlPersist, + _initial_request_was_none: bool, + ) -> Result<()> { + anyhow::bail!("connection multiplexing requires Unix-domain sockets") + } + } + + pub fn start_control_master( + _path: &Path, + _client: Client, + _require_confirmation: bool, + ) -> Result { + anyhow::bail!("connection multiplexing requires Unix-domain sockets") + } + + pub async fn attach_session( + _path: &Path, + _session: SessionOpenRequest, + _invoking_policy: &crate::ssh::SessionPolicy, + ) -> Result { + Ok(AttachOutcome::NoMaster) + } + + pub async fn send_control_command( + _path: &Path, + _command: ControlCommand, + _forwards: Vec, + _address_family: AddressFamily, + ) -> Result { + anyhow::bail!("connection multiplexing requires Unix-domain sockets") + } +} + +#[cfg(not(unix))] +pub use unsupported::{ + AttachOutcome, RunningControlMaster, attach_session, send_control_command, start_control_master, +}; diff --git a/src/ssh/control/socket.rs b/src/ssh/control/socket.rs new file mode 100644 index 00000000..5677eabc --- /dev/null +++ b/src/ssh/control/socket.rs @@ -0,0 +1,285 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. + +//! Secure publication and cleanup of a Unix control socket. + +#[cfg(unix)] +mod unix { + use std::fs::{self, Metadata}; + use std::io; + use std::os::unix::fs::{FileTypeExt as _, MetadataExt as _, PermissionsExt as _}; + use std::path::{Path, PathBuf}; + + use anyhow::{Context, Result}; + use tokio::net::{UnixListener, UnixStream}; + + use super::super::validate_control_socket_path; + + #[derive(Debug)] + pub struct ControlSocketGuard { + path: PathBuf, + device: u64, + inode: u64, + } + + struct TemporarySocketGuard(Option); + + impl TemporarySocketGuard { + fn new(path: PathBuf) -> Self { + Self(Some(path)) + } + + fn disarm(&mut self) { + self.0 = None; + } + } + + impl Drop for TemporarySocketGuard { + fn drop(&mut self) { + if let Some(path) = self.0.take() { + let _ = fs::remove_file(path); + } + } + } + + impl ControlSocketGuard { + fn published(path: PathBuf, metadata: &Metadata) -> Self { + Self { + path, + device: metadata.dev(), + inode: metadata.ino(), + } + } + + fn still_owns_path(&self) -> bool { + fs::symlink_metadata(&self.path).is_ok_and(|metadata| { + metadata.file_type().is_socket() + && metadata.dev() == self.device + && metadata.ino() == self.inode + }) + } + } + + impl Drop for ControlSocketGuard { + fn drop(&mut self) { + if self.still_owns_path() + && let Err(error) = fs::remove_file(&self.path) + && error.kind() != io::ErrorKind::NotFound + { + tracing::warn!( + path = %self.path.display(), + "Could not remove owned control socket: {error}" + ); + } + } + } + + pub async fn connect_control_socket(path: &Path) -> io::Result { + validate_control_socket_path(path) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?; + let stream = UnixStream::connect(path).await?; + verify_same_user(&stream)?; + Ok(stream) + } + + pub fn verify_same_user(stream: &UnixStream) -> io::Result<()> { + let peer_uid = stream.peer_cred()?.uid(); + // SAFETY: geteuid() has no preconditions and only reads process credentials. + let effective_uid = unsafe { libc::geteuid() }; + if peer_uid != effective_uid { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!( + "control socket peer uid {peer_uid} does not match current uid {effective_uid}" + ), + )); + } + Ok(()) + } + + /// Remove a refused socket only when it is still the exact same socket owned + /// by this uid. Regular files, symlinks, replacement sockets, and foreign + /// sockets are never removed. + pub fn remove_stale_control_socket(path: &Path) -> Result { + let first = match fs::symlink_metadata(path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(false), + Err(error) => return Err(error).context("Could not inspect stale control socket"), + }; + // SAFETY: geteuid() has no preconditions and only reads process credentials. + let effective_uid = unsafe { libc::geteuid() }; + if !first.file_type().is_socket() || first.uid() != effective_uid { + return Ok(false); + } + let second = fs::symlink_metadata(path) + .with_context(|| format!("Could not re-check control socket '{}'", path.display()))?; + if !second.file_type().is_socket() + || second.uid() != effective_uid + || first.dev() != second.dev() + || first.ino() != second.ino() + { + return Ok(false); + } + fs::remove_file(path).with_context(|| { + format!("Could not remove stale control socket '{}'", path.display()) + })?; + Ok(true) + } + + /// Bind with a private temporary pathname, chmod it before publication, and + /// hard-link it into place. No client can reach the socket while its mode is + /// broader than 0600. + pub fn bind_control_socket(path: &Path) -> Result<(UnixListener, ControlSocketGuard)> { + validate_control_socket_path(path)?; + let parent = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .unwrap_or(Path::new(".")); + let mut last_collision = None; + for attempt in 0..32_u32 { + let temporary = parent.join(format!(".bssh-mux-{}-{attempt}", std::process::id())); + if validate_control_socket_path(&temporary).is_err() { + anyhow::bail!( + "ControlPath directory '{}' leaves no room for a safe temporary socket; shorten it or use %C", + parent.display() + ); + } + let listener = match std::os::unix::net::UnixListener::bind(&temporary) { + Ok(listener) => listener, + Err(error) if error.kind() == io::ErrorKind::AddrInUse => { + last_collision = Some(error); + continue; + } + Err(error) => { + return Err(error).with_context(|| { + format!("Could not bind control socket '{}'", temporary.display()) + }); + } + }; + let mut cleanup_temporary = TemporarySocketGuard::new(temporary.clone()); + fs::set_permissions(&temporary, fs::Permissions::from_mode(0o600)) + .context("Could not restrict control socket permissions")?; + fs::hard_link(&temporary, path).with_context(|| { + format!( + "Could not publish control socket '{}'; another master may already own it", + path.display() + ) + })?; + let metadata = + fs::symlink_metadata(path).context("Could not inspect published control socket")?; + fs::remove_file(&temporary) + .context("Could not remove temporary control socket link")?; + cleanup_temporary.disarm(); + listener + .set_nonblocking(true) + .context("Could not make control socket nonblocking")?; + let listener = UnixListener::from_std(listener) + .context("Could not register control socket with Tokio")?; + return Ok(( + listener, + ControlSocketGuard::published(path.to_path_buf(), &metadata), + )); + } + Err(last_collision.unwrap_or_else(|| { + io::Error::new(io::ErrorKind::AddrInUse, "temporary socket collision") + })) + .context("Could not allocate a temporary control socket") + } +} + +#[cfg(unix)] +pub use unix::{ + ControlSocketGuard, bind_control_socket, connect_control_socket, remove_stale_control_socket, + verify_same_user, +}; + +#[cfg(all(test, unix))] +mod tests { + use std::fs; + use std::os::unix::fs::{FileTypeExt as _, PermissionsExt as _}; + + use tempfile::TempDir; + + use super::*; + + #[tokio::test] + async fn published_socket_is_private_same_user_and_owned_cleanup_is_exact() { + let directory = TempDir::new().unwrap(); + let path = directory.path().join("control"); + let (listener, guard) = bind_control_socket(&path).unwrap(); + let metadata = fs::symlink_metadata(&path).unwrap(); + assert!(metadata.file_type().is_socket()); + assert_eq!(metadata.permissions().mode() & 0o777, 0o600); + + let connect = connect_control_socket(&path); + let accept = listener.accept(); + let (client, accepted) = tokio::join!(connect, accept); + client.unwrap(); + verify_same_user(&accepted.unwrap().0).unwrap(); + + drop(guard); + assert!(!path.exists()); + } + + #[test] + fn stale_cleanup_preserves_regular_files_and_symlinks() { + let directory = TempDir::new().unwrap(); + let regular = directory.path().join("regular"); + fs::write(®ular, b"do not remove").unwrap(); + assert!(!remove_stale_control_socket(®ular).unwrap()); + assert_eq!(fs::read(®ular).unwrap(), b"do not remove"); + + let symlink = directory.path().join("symlink"); + std::os::unix::fs::symlink(®ular, &symlink).unwrap(); + assert!(!remove_stale_control_socket(&symlink).unwrap()); + assert!( + fs::symlink_metadata(&symlink) + .unwrap() + .file_type() + .is_symlink() + ); + } + + #[test] + fn stale_owned_socket_is_removed() { + let directory = TempDir::new().unwrap(); + let path = directory.path().join("stale"); + let listener = std::os::unix::net::UnixListener::bind(&path).unwrap(); + drop(listener); + assert!(remove_stale_control_socket(&path).unwrap()); + assert!(!path.exists()); + } +} + +#[cfg(not(unix))] +mod unsupported { + use std::io; + use std::path::Path; + + use anyhow::Result; + + #[derive(Debug)] + pub struct ControlSocketGuard; + + pub async fn connect_control_socket(_path: &Path) -> io::Result<()> { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "connection multiplexing requires Unix-domain sockets", + )) + } + + pub fn bind_control_socket(_path: &Path) -> Result<((), ControlSocketGuard)> { + anyhow::bail!("connection multiplexing requires Unix-domain sockets") + } + + pub fn remove_stale_control_socket(_path: &Path) -> Result { + Ok(false) + } +} + +#[cfg(not(unix))] +pub use unsupported::{ + ControlSocketGuard, bind_control_socket, connect_control_socket, remove_stale_control_socket, +}; diff --git a/src/ssh/mod.rs b/src/ssh/mod.rs index 57844836..61c7a2bd 100644 --- a/src/ssh/mod.rs +++ b/src/ssh/mod.rs @@ -15,6 +15,7 @@ pub mod auth; pub mod client; pub mod config_cache; +pub mod control; pub mod handler; pub mod known_hosts; pub mod pool; @@ -28,6 +29,9 @@ pub mod keychain_macos; pub use auth::AuthContext; pub use client::SshClient; pub use config_cache::{CacheConfig, CacheStats, GLOBAL_CACHE, SshConfigCache}; +pub use control::{ + ControlCommand, ControlMasterMode, ControlPathContext, ControlPersist, ControlPolicy, +}; pub use handler::BsshHandler; pub use pool::ConnectionPool; pub use session_policy::{CliTtyMode, SessionPolicy, SessionPurpose, SessionRequest}; diff --git a/src/ssh/session_policy.rs b/src/ssh/session_policy.rs index 9e509ccf..dc577b7c 100644 --- a/src/ssh/session_policy.rs +++ b/src/ssh/session_policy.rs @@ -49,7 +49,7 @@ pub enum SessionPurpose { Bulk, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum SessionRequest { Exec(String), Shell, @@ -57,7 +57,7 @@ pub enum SessionRequest { None, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct SessionPolicy { pub environment: Vec<(String, String)>, pub local_command: Option, diff --git a/src/ssh/ssh_config/parser/options/control.rs b/src/ssh/ssh_config/parser/options/control.rs index 29e25f0a..b6af4890 100644 --- a/src/ssh/ssh_config/parser/options/control.rs +++ b/src/ssh/ssh_config/parser/options/control.rs @@ -17,10 +17,53 @@ //! Handles control socket configuration options for connection multiplexing //! including ControlMaster, ControlPath, and ControlPersist settings. -use crate::ssh::ssh_config::security::validate_control_path; +use crate::ssh::control::{ControlMasterMode, ControlPersist}; use crate::ssh::ssh_config::types::SshHostConfig; use anyhow::Result; +fn exactly_one_value<'a>(keyword: &str, args: &'a [String], line_number: usize) -> Result<&'a str> { + if args.len() != 1 || args[0].is_empty() { + anyhow::bail!("{keyword} expects exactly one value at line {line_number}"); + } + Ok(&args[0]) +} + +fn validate_control_master(value: &str, line_number: usize) -> Result<()> { + value + .parse::() + .map(|_| ()) + .map_err(|error| anyhow::anyhow!("{error} at line {line_number}")) +} + +fn validate_control_path(value: &str, line_number: usize) -> Result<()> { + if value.contains('\0') { + anyhow::bail!("ControlPath contains a NUL byte at line {line_number}"); + } + + let mut tokens = value.char_indices(); + while let Some((_, character)) = tokens.next() { + if character != '%' { + continue; + } + let Some((_, token)) = tokens.next() else { + anyhow::bail!("ControlPath has an incomplete '%' token at line {line_number}"); + }; + if !matches!(token, '%' | 'C' | 'h' | 'l' | 'p' | 'r') { + anyhow::bail!( + "ControlPath contains unknown substitution token '%{token}' at line {line_number}" + ); + } + } + Ok(()) +} + +fn validate_control_persist(value: &str, line_number: usize) -> Result<()> { + value + .parse::() + .map(|_| ()) + .map_err(|error| anyhow::anyhow!("{error} at line {line_number}")) +} + /// Parse control socket SSH configuration options pub(super) fn parse_control_option( host: &mut SshHostConfig, @@ -30,28 +73,121 @@ pub(super) fn parse_control_option( ) -> Result<()> { match keyword { "controlmaster" => { - if args.is_empty() { - anyhow::bail!("ControlMaster requires a value at line {line_number}"); - } - host.control_master = Some(args[0].clone()); + let value = exactly_one_value("ControlMaster", args, line_number)?; + validate_control_master(value, line_number)?; + host.control_master = Some(value.to_string()); } "controlpath" => { - if args.is_empty() { - anyhow::bail!("ControlPath requires a value at line {line_number}"); - } - let path = args[0].clone(); - // ControlPath has different validation - it allows SSH substitution patterns - validate_control_path(&path, line_number)?; - host.control_path = Some(path); + let value = exactly_one_value("ControlPath", args, line_number)?; + validate_control_path(value, line_number)?; + host.control_path = Some(value.to_string()); } "controlpersist" => { - if args.is_empty() { - anyhow::bail!("ControlPersist requires a value at line {line_number}"); - } - host.control_persist = Some(args[0].clone()); + let value = exactly_one_value("ControlPersist", args, line_number)?; + validate_control_persist(value, line_number)?; + host.control_persist = Some(value.to_string()); } _ => unreachable!("Unexpected keyword in parse_control_option: {}", keyword), } Ok(()) } + +#[cfg(test)] +fn parse_for_test(keyword: &str, values: &[&str]) -> Result { + let mut host = SshHostConfig::default(); + let values = values + .iter() + .map(|value| (*value).to_string()) + .collect::>(); + parse_control_option(&mut host, keyword, &values, 7)?; + Ok(host) +} + +#[test] +fn control_options_require_exactly_one_value() { + for keyword in ["controlmaster", "controlpath", "controlpersist"] { + assert!(parse_for_test(keyword, &[]).is_err(), "{keyword}"); + assert!(parse_for_test(keyword, &[""]).is_err(), "{keyword}"); + assert!( + parse_for_test(keyword, &["yes", "extra"]).is_err(), + "{keyword}" + ); + } +} + +#[test] +fn validates_control_master_modes_case_insensitively() { + for value in [ + "yes", "true", "no", "false", "auto", "ask", "autoask", "AUTO", + ] { + assert!(parse_for_test("controlmaster", &[value]).is_ok(), "{value}"); + } + for value in ["", "maybe", "auto-ask", "1"] { + assert!( + parse_for_test("controlmaster", &[value]).is_err(), + "{value}" + ); + } +} + +#[test] +fn accepts_control_path_literals_and_valid_templates() { + for value in [ + "none", + "~/.ssh/control-%C-%h-%p-%r-%l", + "/tmp/control path;literal|&`$(still-a-path)-%%", + ] { + assert!(parse_for_test("controlpath", &[value]).is_ok(), "{value:?}"); + } + for value in ["/tmp/control-%x", "/tmp/control-%", "/tmp/control\0path"] { + assert!( + parse_for_test("controlpath", &[value]).is_err(), + "{value:?}" + ); + } +} + +#[test] +fn validates_openssh_control_persist_time_grammar() { + let valid = [ + "yes", + "TRUE", + "no", + "False", + "0", + "2s", + "3m", + "1m30", + "1H30m", + "1w2d3h4m5s", + "1s2s", + "2147483647", + "3550w5d3h14m7s", + ]; + for value in valid { + assert!( + parse_for_test("controlpersist", &[value]).is_ok(), + "{value}" + ); + } + + let invalid = [ + "", + "-1", + "trout", + "1.5m", + "1.s", + ".5s", + "1m0.5s", + "1e3", + "2147483648", + "3550w5d3h14m8s", + ]; + for value in invalid { + assert!( + parse_for_test("controlpersist", &[value]).is_err(), + "{value}" + ); + } +} diff --git a/src/ssh/ssh_config/parser/options/support.rs b/src/ssh/ssh_config/parser/options/support.rs index de57c331..69b568f6 100644 --- a/src/ssh/ssh_config/parser/options/support.rs +++ b/src/ssh/ssh_config/parser/options/support.rs @@ -15,6 +15,7 @@ pub(super) enum RuntimeConsumer { Proxy, Session, Forwarding, + Multiplexing, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -25,7 +26,8 @@ pub(super) struct KeywordSpec { use KeywordSupport::{Runtime, Unimplemented}; use RuntimeConsumer::{ - Authentication, Forwarding, HostVerification, NodeResolution, Proxy, Session, Transport, + Authentication, Forwarding, HostVerification, Multiplexing, NodeResolution, Proxy, Session, + Transport, }; pub(super) const ACCEPTED_KEYWORDS: &[(&str, &str, KeywordSupport)] = &[ @@ -198,9 +200,9 @@ pub(super) const ACCEPTED_KEYWORDS: &[(&str, &str, KeywordSupport)] = &[ ("proxyjump", "proxyjump", Runtime(Proxy)), ("proxycommand", "proxycommand", Runtime(Proxy)), ("proxyusefdpass", "proxyusefdpass", Runtime(Proxy)), - ("controlmaster", "controlmaster", Unimplemented), - ("controlpath", "controlpath", Unimplemented), - ("controlpersist", "controlpersist", Unimplemented), + ("controlmaster", "controlmaster", Runtime(Multiplexing)), + ("controlpath", "controlpath", Runtime(Multiplexing)), + ("controlpersist", "controlpersist", Runtime(Multiplexing)), ("sendenv", "sendenv", Runtime(Session)), ("setenv", "setenv", Runtime(Session)), ("requesttty", "requesttty", Runtime(Session)), @@ -299,8 +301,8 @@ mod tests { use std::collections::HashSet; const ACCEPTED_SPELLING_COUNT: usize = 107; - const RUNTIME_SPELLING_COUNT: usize = 51; - const UNIMPLEMENTED_SPELLING_COUNT: usize = 56; + const RUNTIME_SPELLING_COUNT: usize = 54; + const UNIMPLEMENTED_SPELLING_COUNT: usize = 53; #[test] fn accepted_keywords_and_aliases_have_one_consistent_classification() { @@ -384,6 +386,9 @@ mod tests { ("proxyjump", Proxy), ("proxycommand", Proxy), ("proxyusefdpass", Proxy), + ("controlmaster", Multiplexing), + ("controlpath", Multiplexing), + ("controlpersist", Multiplexing), ("sendenv", Session), ("setenv", Session), ("requesttty", Session), @@ -432,9 +437,6 @@ mod tests { "forwardx11timeout", "forwardx11trusted", "connecttimeout", - "controlmaster", - "controlpath", - "controlpersist", "escapechar", "loglevel", "syslogfacility", diff --git a/src/ssh/ssh_config/security/mod.rs b/src/ssh/ssh_config/security/mod.rs index fb2d790e..91c07cc2 100644 --- a/src/ssh/ssh_config/security/mod.rs +++ b/src/ssh/ssh_config/security/mod.rs @@ -22,7 +22,9 @@ mod path_validation; mod string_validation; pub use path_validation::secure_validate_path; -pub use string_validation::{validate_control_path, validate_executable_string}; +#[cfg(test)] +pub use string_validation::validate_control_path; +pub use string_validation::validate_executable_string; #[cfg(test)] mod tests; diff --git a/src/ssh/ssh_config/security/string_validation.rs b/src/ssh/ssh_config/security/string_validation.rs index d859ceea..73186238 100644 --- a/src/ssh/ssh_config/security/string_validation.rs +++ b/src/ssh/ssh_config/security/string_validation.rs @@ -265,6 +265,7 @@ fn validate_local_executable_command( /// # Returns /// * `Ok(())` if the path is safe /// * `Err(anyhow::Error)` if the path contains dangerous patterns +#[cfg(test)] pub fn validate_control_path(path: &str, line_number: usize) -> Result<()> { // ControlPath "none" is a special case to disable control path if path == "none" { diff --git a/src/ssh/tokio_client/address_family.rs b/src/ssh/tokio_client/address_family.rs index de281847..57657312 100644 --- a/src/ssh/tokio_client/address_family.rs +++ b/src/ssh/tokio_client/address_family.rs @@ -32,7 +32,9 @@ use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; /// address is a candidate, tried in resolver order. The other two variants are /// hard constraints; there is no fallback to the other family, matching /// OpenSSH. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Hash, Default, serde::Serialize, serde::Deserialize, +)] pub enum AddressFamily { /// No constraint (`AddressFamily any`, neither `-4` nor `-6`). #[default] diff --git a/src/ssh/tokio_client/connection.rs b/src/ssh/tokio_client/connection.rs index d05e6f9e..1f31230b 100644 --- a/src/ssh/tokio_client/connection.rs +++ b/src/ssh/tokio_client/connection.rs @@ -1944,6 +1944,26 @@ impl Client { result } + pub(crate) async fn add_control_forwardings( + &self, + plan: &ForwardingPlan, + ) -> Result<(), super::Error> { + self.forwarding_runtime + .add_control_forwardings(self, plan) + .await + .map_err(|error| super::Error::PortForwardRequestFailed(format!("{error:#}"))) + } + + pub(crate) async fn cancel_control_forwardings( + &self, + plan: &ForwardingPlan, + ) -> Result<(), super::Error> { + self.forwarding_runtime + .cancel_control_forwardings(plan) + .await + .map_err(|error| super::Error::PortForwardRequestFailed(format!("{error:#}"))) + } + pub(crate) fn remote_forward_registry(&self) -> RemoteForwardRegistry { self.remote_forward_registry.clone() } diff --git a/src/ssh/tokio_client/session.rs b/src/ssh/tokio_client/session.rs index bb369121..320a7085 100644 --- a/src/ssh/tokio_client/session.rs +++ b/src/ssh/tokio_client/session.rs @@ -44,19 +44,33 @@ impl Client { .await } - async fn execute_session_streaming_with_input( + pub(crate) async fn execute_session_streaming_with_input( &self, policy: &SessionPolicy, sender: Sender, input: R, ) -> Result + where + R: AsyncRead + Unpin, + { + self.execute_session_streaming_with_input_and_terminal(policy, None, sender, input) + .await + } + + pub(crate) async fn execute_session_streaming_with_input_and_terminal( + &self, + policy: &SessionPolicy, + terminal: Option<&str>, + sender: Sender, + input: R, + ) -> Result where R: AsyncRead + Unpin, { if matches!(policy.request, SessionRequest::None) { return Ok(0); } - let channel = self.open_policy_channel(policy).await?; + let channel = self.open_policy_channel(policy, terminal).await?; self.drain_policy_channel_with_input(channel, sender, input) .await } @@ -83,6 +97,7 @@ impl Client { async fn open_policy_channel( &self, policy: &SessionPolicy, + terminal: Option<&str>, ) -> Result, super::Error> { let channel = self .connection_handle @@ -91,9 +106,17 @@ impl Client { .map_err(|source| self.session_error_or(super::Error::ChannelOpen { source }))?; if policy.request_pty { - let terminal = std::env::var("TERM").unwrap_or_else(|_| "xterm".to_string()); + let inherited_terminal; + let terminal = match terminal { + Some(terminal) => terminal, + None => { + inherited_terminal = + std::env::var("TERM").unwrap_or_else(|_| "xterm".to_string()); + &inherited_terminal + } + }; channel - .request_pty(true, &terminal, 80, 24, 0, 0, &[]) + .request_pty(true, terminal, 80, 24, 0, 0, &[]) .await .map_err(|source| { self.session_error_or(super::Error::CommandExecution { diff --git a/tests/control_multiplexing_live_test.rs b/tests/control_multiplexing_live_test.rs new file mode 100644 index 00000000..98fa9f4c --- /dev/null +++ b/tests/control_multiplexing_live_test.rs @@ -0,0 +1,328 @@ +use std::net::{Ipv4Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use anyhow::Result; +use bssh::forwarding::ForwardingDirective; +use bssh::ssh::control::{ + AttachOutcome, ControlCommand, ControlPersist, ControlResponseKind, SessionOpenRequest, + attach_session, send_control_command, start_control_master, +}; +use bssh::ssh::tokio_client::{ + AddressFamily, AuthMethod, Client, ServerCheckMethod, SshConnectionConfig, +}; +use bssh::ssh::{SessionPolicy, SessionRequest}; +use russh::keys::{Algorithm, PrivateKey}; +use russh::server::{self, Msg, Server, Session}; +use russh::{Channel, ChannelId}; +use tempfile::TempDir; +use tokio::net::TcpListener; +use tokio::task::JoinHandle; +use tokio::time::timeout; + +const TEST_TIMEOUT: Duration = Duration::from_secs(5); + +#[derive(Default)] +struct ServerState { + authentications: AtomicUsize, + sessions: AtomicUsize, +} + +#[derive(Clone)] +struct MultiplexTestServer { + state: Arc, +} + +impl Server for MultiplexTestServer { + type Handler = Self; + + fn new_client(&mut self, _peer_addr: Option) -> Self::Handler { + self.clone() + } +} + +impl server::Handler for MultiplexTestServer { + type Error = anyhow::Error; + + async fn auth_password( + &mut self, + _user: &str, + _password: &str, + ) -> Result { + self.state.authentications.fetch_add(1, Ordering::SeqCst); + Ok(server::Auth::Accept) + } + + async fn channel_open_session( + &mut self, + _channel: Channel, + reply: server::ChannelOpenHandle, + _session: &mut Session, + ) -> Result<(), Self::Error> { + self.state.sessions.fetch_add(1, Ordering::SeqCst); + reply.accept().await; + Ok(()) + } + + async fn exec_request( + &mut self, + channel: ChannelId, + data: &[u8], + session: &mut Session, + ) -> Result<(), Self::Error> { + session.channel_success(channel)?; + session.data(channel, data.to_vec())?; + if data == b"slow" { + tokio::time::sleep(Duration::from_millis(250)).await; + } + session.exit_status_request(channel, 0)?; + session.eof(channel)?; + session.close(channel)?; + Ok(()) + } +} + +struct RunningSshServer { + address: SocketAddr, + state: Arc, + task: JoinHandle>, +} + +impl RunningSshServer { + async fn start() -> Self { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .expect("bind SSH test server"); + let address = listener.local_addr().expect("SSH server address"); + let state = Arc::new(ServerState::default()); + let key = PrivateKey::random(&mut rand::rng(), Algorithm::Ed25519) + .expect("generate SSH test host key"); + let config = Arc::new(server::Config { + keys: vec![key], + auth_rejection_time: Duration::ZERO, + auth_rejection_time_initial: Some(Duration::ZERO), + ..Default::default() + }); + let mut server = MultiplexTestServer { + state: Arc::clone(&state), + }; + let task = tokio::spawn(async move { server.run_on_socket(config, &listener).await }); + Self { + address, + state, + task, + } + } + + async fn shutdown(self) { + self.task.abort(); + let _ = self.task.await; + } +} + +fn session(command: &str) -> (SessionPolicy, SessionOpenRequest) { + let policy = SessionPolicy { + environment: Vec::new(), + local_command: None, + request_pty: false, + stdin_null: true, + request: SessionRequest::Exec(command.to_string()), + }; + let request = SessionOpenRequest::new(policy.clone(), None).expect("wire-safe session"); + (policy, request) +} + +#[tokio::test] +async fn two_passengers_share_exactly_one_authenticated_transport() { + let ssh = RunningSshServer::start().await; + let client = Client::connect_with_ssh_config( + ssh.address, + "test", + AuthMethod::with_password("test"), + ServerCheckMethod::NoCheck, + &SshConnectionConfig::default(), + ) + .await + .expect("authenticate control master"); + assert_eq!(ssh.state.authentications.load(Ordering::SeqCst), 1); + + let directory = TempDir::new().expect("control tempdir"); + let path = directory.path().join("mux"); + let master = start_control_master(&path, client, false).expect("start control master"); + + for command in ["first", "second"] { + let (policy, request) = session(command); + let outcome = timeout(TEST_TIMEOUT, attach_session(&path, request, &policy)) + .await + .expect("passenger timed out") + .expect("passenger failed"); + assert_eq!(outcome, AttachOutcome::ExitStatus(0)); + } + assert_eq!(ssh.state.authentications.load(Ordering::SeqCst), 1); + assert_eq!(ssh.state.sessions.load(Ordering::SeqCst), 2); + + let check = send_control_command(&path, ControlCommand::Check, Vec::new(), AddressFamily::Any) + .await + .expect("check control master"); + assert!(matches!(check, ControlResponseKind::Alive { .. })); + + let reserved = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .expect("reserve local forwarding port"); + let forwarding_port = reserved.local_addr().unwrap().port(); + drop(reserved); + let directive = ForwardingDirective::Local(format!("127.0.0.1:{forwarding_port}:127.0.0.1:1")); + let forward = send_control_command( + &path, + ControlCommand::Forward, + vec![directive.clone()], + AddressFamily::V4, + ) + .await + .expect("add control forwarding"); + assert_eq!(forward, ControlResponseKind::Ok); + assert!( + TcpListener::bind((Ipv4Addr::LOCALHOST, forwarding_port)) + .await + .is_err(), + "forward command must not reply before its listener is ready" + ); + let cancel = send_control_command( + &path, + ControlCommand::Cancel, + vec![directive], + AddressFamily::V4, + ) + .await + .expect("cancel control forwarding"); + assert_eq!(cancel, ControlResponseKind::Ok); + TcpListener::bind((Ipv4Addr::LOCALHOST, forwarding_port)) + .await + .expect("cancel command must release its listener before replying"); + + let (slow_policy, slow_request) = session("slow"); + let slow_path = path.clone(); + let passenger = + tokio::spawn(async move { attach_session(&slow_path, slow_request, &slow_policy).await }); + timeout(TEST_TIMEOUT, async { + while ssh.state.sessions.load(Ordering::SeqCst) != 3 { + tokio::task::yield_now().await; + } + }) + .await + .expect("slow passenger did not open"); + + send_control_command(&path, ControlCommand::Stop, Vec::new(), AddressFamily::Any) + .await + .expect("stop control master"); + timeout(Duration::from_millis(100), async { + while path.exists() { + tokio::task::yield_now().await; + } + }) + .await + .expect("stop did not unlink the listener promptly"); + assert!( + !passenger.is_finished(), + "stop must not cancel an active passenger" + ); + assert_eq!( + passenger.await.unwrap().unwrap(), + AttachOutcome::ExitStatus(0) + ); + timeout( + TEST_TIMEOUT, + master.finish_after_initial(ControlPersist::Forever, true), + ) + .await + .expect("master stop timed out") + .expect("master stop failed"); + ssh.shutdown().await; +} + +#[tokio::test] +async fn exit_immediately_cancels_active_passengers() { + let ssh = RunningSshServer::start().await; + let client = Client::connect_with_ssh_config( + ssh.address, + "test", + AuthMethod::with_password("test"), + ServerCheckMethod::NoCheck, + &SshConnectionConfig::default(), + ) + .await + .expect("authenticate control master"); + let directory = TempDir::new().expect("control tempdir"); + let path = directory.path().join("mux-exit"); + let master = start_control_master(&path, client, false).expect("start control master"); + let (policy, request) = session("slow"); + let passenger_path = path.clone(); + let passenger = + tokio::spawn(async move { attach_session(&passenger_path, request, &policy).await }); + timeout(TEST_TIMEOUT, async { + while ssh.state.sessions.load(Ordering::SeqCst) != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("slow passenger did not open"); + + send_control_command(&path, ControlCommand::Exit, Vec::new(), AddressFamily::Any) + .await + .expect("exit control master"); + timeout( + TEST_TIMEOUT, + master.finish_after_initial(ControlPersist::Forever, true), + ) + .await + .expect("master exit timed out") + .expect("master exit failed"); + assert!(!path.exists()); + assert!( + passenger.await.unwrap().is_err(), + "exit must cancel an active passenger" + ); + assert_eq!(ssh.state.authentications.load(Ordering::SeqCst), 1); + ssh.shutdown().await; +} + +#[tokio::test] +async fn control_persist_timeout_resets_after_a_new_passenger() { + let ssh = RunningSshServer::start().await; + let client = Client::connect_with_ssh_config( + ssh.address, + "test", + AuthMethod::with_password("test"), + ServerCheckMethod::NoCheck, + &SshConnectionConfig::default(), + ) + .await + .expect("authenticate control master"); + let directory = TempDir::new().expect("control tempdir"); + let path = directory.path().join("mux-persist"); + let master = start_control_master(&path, client, false).expect("start control master"); + let finish = tokio::spawn( + master.finish_after_initial(ControlPersist::Timeout(Duration::from_millis(150)), false), + ); + + tokio::time::sleep(Duration::from_millis(90)).await; + let (policy, request) = session("reset"); + assert_eq!( + attach_session(&path, request, &policy).await.unwrap(), + AttachOutcome::ExitStatus(0) + ); + tokio::time::sleep(Duration::from_millis(90)).await; + assert!( + !finish.is_finished(), + "a passenger must reset the idle timeout" + ); + timeout(TEST_TIMEOUT, finish) + .await + .expect("persist timeout did not expire") + .expect("persist task failed") + .expect("persist shutdown failed"); + assert!(!path.exists()); + assert_eq!(ssh.state.authentications.load(Ordering::SeqCst), 1); + ssh.shutdown().await; +}