From 26a830dcac519eca48e1eb254d7f377d23721861 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 1 Sep 2026 04:47:41 +0900 Subject: [PATCH] feat(ssh): add connection multiplexing runtime Add a private Unix control socket protocol that reuses one authenticated SSH transport for passenger sessions and runtime forwarding control. Implement typed ControlMaster, ControlPath, and ControlPersist policy resolution, bounded same-user handshakes, stale-socket handling, exact forwarding registration, and check, forward, cancel, exit, and stop commands. Cover authentication reuse, persist timer reset, graceful and immediate shutdown, socket safety, parser boundaries, and forwarding teardown with unit and live SSH tests. Refs #286 --- src/app/dispatcher.rs | 291 ++++++ src/cli/bssh.rs | 74 +- src/cli/pdsh.rs | 3 + src/cli/ssh_args.rs | 11 +- src/forwarding/manager.rs | 20 +- src/forwarding/mod.rs | 6 +- src/forwarding/remote.rs | 7 +- src/forwarding/runtime.rs | 264 ++++- src/ssh/client/command.rs | 79 +- src/ssh/control/config.rs | 395 ++++++++ src/ssh/control/mod.rs | 40 + src/ssh/control/path.rs | 297 ++++++ src/ssh/control/protocol.rs | 391 ++++++++ src/ssh/control/runtime.rs | 903 ++++++++++++++++++ src/ssh/control/socket.rs | 285 ++++++ src/ssh/mod.rs | 4 + src/ssh/session_policy.rs | 4 +- src/ssh/ssh_config/parser/options/control.rs | 168 +++- src/ssh/ssh_config/parser/options/support.rs | 20 +- src/ssh/ssh_config/security/mod.rs | 4 +- .../ssh_config/security/string_validation.rs | 1 + src/ssh/tokio_client/address_family.rs | 4 +- src/ssh/tokio_client/connection.rs | 20 + src/ssh/tokio_client/session.rs | 31 +- tests/control_multiplexing_live_test.rs | 328 +++++++ 25 files changed, 3549 insertions(+), 101 deletions(-) create mode 100644 src/ssh/control/config.rs create mode 100644 src/ssh/control/mod.rs create mode 100644 src/ssh/control/path.rs create mode 100644 src/ssh/control/protocol.rs create mode 100644 src/ssh/control/runtime.rs create mode 100644 src/ssh/control/socket.rs create mode 100644 tests/control_multiplexing_live_test.rs 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; +}