From 76aacf46313ee8277fbf0ac8a5aa4094252ceea0 Mon Sep 17 00:00:00 2001 From: Jonathan Liebig Date: Tue, 13 Oct 2026 20:19:29 +0200 Subject: [PATCH] fix: cancel and clean up endpoint connection attempts refs #4721 --- src/client/endpoint/supervisor.rs | 236 ++++++++++++++- src/client/endpoint/writer.rs | 1 + src/client/handshake.rs | 18 +- src/client/mod.rs | 1 + src/client/terminal_sessions.rs | 2 + src/ipc.rs | 41 +++ src/platform/mod.rs | 2 +- src/platform/unix_common.rs | 51 ++++ src/platform/windows.rs | 145 ++++++++++ src/remote/attach.rs | 210 ++++++++++---- src/remote/process.rs | 465 +++++++++++++++++++++++++----- src/remote/saved.rs | 61 +++- 12 files changed, 1087 insertions(+), 146 deletions(-) diff --git a/src/client/endpoint/supervisor.rs b/src/client/endpoint/supervisor.rs index 4be72784e..417cd7323 100644 --- a/src/client/endpoint/supervisor.rs +++ b/src/client/endpoint/supervisor.rs @@ -53,6 +53,8 @@ struct ReconnectState { attempts: u32, next_attempt: Option, in_flight: bool, + cancelled: Arc, + completed: Arc, generation: Option, ready_since: Option, } @@ -64,14 +66,31 @@ impl ReconnectState { attempts: 0, next_attempt: Some(now), in_flight: false, + cancelled: Arc::new(AtomicBool::new(false)), + completed: Arc::new(AtomicBool::new(true)), generation: None, ready_since: None, } } } +impl Drop for ReconnectState { + fn drop(&mut self) { + self.cancelled.store(true, Ordering::Release); + } +} + +struct AttemptCompletion(Arc); + +impl Drop for AttemptCompletion { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } +} + pub(crate) struct EndpointSupervisors { endpoints: HashMap, + retired: HashMap>, next_generation: u64, shutdown: Arc, } @@ -90,6 +109,7 @@ impl EndpointSupervisors { .collect(); Self { endpoints, + retired: HashMap::new(), next_generation: 2, shutdown: Arc::new(AtomicBool::new(false)), } @@ -101,7 +121,13 @@ impl EndpointSupervisors { if generation.is_some() { state.next_attempt = None; } - self.endpoints.insert(ClientEndpointId::Local, state); + if let Some(previous) = self.endpoints.insert(ClientEndpointId::Local, state) { + if !previous.completed.load(Ordering::Acquire) { + self.retired + .entry(ClientEndpointId::Local) + .or_insert_with(|| Arc::clone(&previous.completed)); + } + } } pub(crate) fn reconcile_profiles( @@ -110,6 +136,7 @@ impl EndpointSupervisors { now: Instant, ) -> Vec { let mut retired = Vec::new(); + let retiring = &mut self.retired; self.endpoints.retain(|endpoint_id, state| { let ConnectTarget::Ssh(previous) = &state.target else { return true; @@ -122,6 +149,12 @@ impl EndpointSupervisors { }); if !keep { retired.push(endpoint_id.clone()); + // A stale Connected event still owns its bridge until the client rejects it. + if !state.completed.load(Ordering::Acquire) { + retiring + .entry(endpoint_id.clone()) + .or_insert_with(|| Arc::clone(&state.completed)); + } } keep }); @@ -141,11 +174,21 @@ impl EndpointSupervisors { options: EndpointConnectOptions, event_tx: &tokio::sync::mpsc::Sender, ) { + self.retired + .retain(|_, completed| !completed.load(Ordering::Acquire)); for (endpoint_id, state) in &mut self.endpoints { - if state.in_flight || state.next_attempt.is_none_or(|deadline| deadline > now) { + if state.in_flight + || !state.completed.load(Ordering::Acquire) + || self.retired.contains_key(endpoint_id) + || state.next_attempt.is_none_or(|deadline| deadline > now) + { continue; } state.in_flight = true; + state.cancelled = Arc::new(AtomicBool::new(false)); + let cancelled = Arc::clone(&state.cancelled); + state.completed = Arc::new(AtomicBool::new(false)); + let completed = Arc::clone(&state.completed); state.next_attempt = None; let generation = self.next_generation; state.generation = Some(generation); @@ -156,11 +199,19 @@ impl EndpointSupervisors { let shutdown = self.shutdown.clone(); tokio::spawn(async move { if shutdown.load(Ordering::Acquire) { + completed.store(true, Ordering::Release); return; } let task_endpoint_id = endpoint_id.clone(); let result = tokio::task::spawn_blocking(move || { - connect_once(&target, options, endpoint_id, generation) + connect_once( + &target, + options, + endpoint_id, + generation, + cancelled, + completed, + ) }) .await; let event = match result { @@ -296,10 +347,24 @@ fn connect_once( options: EndpointConnectOptions, endpoint_id: ClientEndpointId, generation: u64, + cancelled: Arc, + completed: Arc, ) -> Result { + let completion = AttemptCompletion(completed); + if cancelled.load(Ordering::Acquire) { + return Err(std::io::Error::new( + std::io::ErrorKind::ConnectionAborted, + "connection attempt cancelled", + )); + } let (mut stream, lifetime): (_, Box) = match target { ConnectTarget::Local(path) => { - let stream = crate::ipc::connect_local_stream(path).map_err(|error| { + let stream = crate::ipc::connect_local_stream_cancellable( + path, + &cancelled, + Instant::now() + super::super::handshake::LOCAL_HANDSHAKE_READ_TIMEOUT, + ) + .map_err(|error| { // An absent Local socket is transient, unlike a missing SSH install. if error.kind() == std::io::ErrorKind::NotFound { std::io::Error::new( @@ -317,6 +382,8 @@ fn connect_once( profile.id.as_str(), &profile.target, &profile.session, + generation, + Arc::clone(&cancelled), )?; (connected.stream, Box::new(connected.bridge)) } @@ -336,8 +403,15 @@ fn connect_once( options.mouse_capture, false, matches!(target, ConnectTarget::Local(_)), + &cancelled, ) .map_err(handshake_error)?; + if cancelled.load(Ordering::Acquire) { + return Err(std::io::Error::new( + std::io::ErrorKind::ConnectionAborted, + "connection attempt cancelled", + )); + } if handshake.encoding != RenderEncoding::SemanticFrame { return Err(std::io::Error::new( std::io::ErrorKind::InvalidData, @@ -357,7 +431,8 @@ fn connect_once( )); } let reader = stream.try_clone()?; - let writer = NativeEndpointTransport::with_lifetime(stream, lifetime)?; + // The completion barrier follows the bridge through queued and active transports. + let writer = NativeEndpointTransport::with_lifetime(stream, (lifetime, completion))?; Ok(EndpointSupervisorEvent::Connected { endpoint_id, generation, @@ -403,6 +478,157 @@ fn retry_delay(attempt: u32) -> Duration { #[cfg(test)] mod tests { use super::*; + + fn connect_options() -> EndpointConnectOptions { + EndpointConnectOptions { + cols: 80, + rows: 24, + cell_width_px: 0, + cell_height_px: 0, + pixel_geometry_exact: false, + pixel_mouse: false, + surface_size: ClientSurfaceSize { cols: 80, rows: 24 }, + endpoint_keybindings: false, + mouse_capture: false, + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn replacement_waits_for_queued_connection_teardown() { + use crate::protocol::endpoint::{EndpointServerWelcome, ENDPOINT_WELCOME_KIND}; + use crate::protocol::{self, ClientMessage, ServerMessage}; + use interprocess::local_socket::traits::{Listener as _, Stream as _}; + let path = std::env::temp_dir().join(format!("hretired-{}.sock", std::process::id())); + let listener = crate::ipc::bind_private_local_listener(&path).unwrap(); + let now = Instant::now(); + let mut supervisors = EndpointSupervisors::new(&[], now); + supervisors.add_local(path.clone(), None, now); + let (events, mut received) = tokio::sync::mpsc::channel(4); + supervisors.spawn_due(now, connect_options(), &events); + let mut peer = listener.accept().unwrap(); + let _: ClientMessage = protocol::read_message(&mut peer, protocol::MAX_FRAME_SIZE).unwrap(); + protocol::write_message( + &mut peer, + &ServerMessage::EndpointControl { + kind: ENDPOINT_WELCOME_KIND.into(), + data: serde_json::to_string(&EndpointServerWelcome::compatible(vec![ + "client_shell.surface.set".into(), + ])) + .unwrap(), + }, + ) + .unwrap(); + let queued = tokio::time::timeout(Duration::from_secs(2), received.recv()) + .await + .unwrap() + .unwrap(); + let EndpointSupervisorEvent::Connected { generation, .. } = &queued else { + panic!("handshake failed"); + }; + let old_generation = *generation; + // Repeated catalog replacements must keep the original unfinished barrier. + for _ in 0..3 { + supervisors.add_local(path.clone(), None, now); + supervisors.spawn_due(now, connect_options(), &events); + assert!(supervisors.endpoints[&ClientEndpointId::Local] + .generation + .is_none()); + assert_eq!(supervisors.retired.len(), 1); + } + assert!(!supervisors.record_status( + &ClientEndpointId::Local, + old_generation, + ClientEndpointStatus::Online, + now + )); + peer.set_nonblocking(true).unwrap(); + assert!(matches!( + crate::ipc::poll_local_stream_read_count(&mut peer, &mut [0; 1]).unwrap(), + crate::ipc::LocalStreamReadCount::Pending + )); + drop(queued); + let deadline = Instant::now() + Duration::from_secs(2); + loop { + supervisors.spawn_due(now, connect_options(), &events); + if supervisors.endpoints[&ClientEndpointId::Local] + .generation + .is_some() + { + break; + } + assert!( + Instant::now() < deadline, + "retired transport did not release its attempt" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert!(supervisors.retired.is_empty()); + assert!(matches!( + crate::ipc::poll_local_stream_read_count(&mut peer, &mut [0; 1]).unwrap(), + crate::ipc::LocalStreamReadCount::Closed + )); + let _replacement = listener.accept().unwrap(); + drop(supervisors); + drop(listener); + std::fs::remove_file(path).unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn retiring_or_shutting_down_a_silent_attempt_closes_its_connection() { + use crate::protocol::{self, ClientMessage}; + use interprocess::local_socket::traits::{Listener as _, Stream as _}; + for shutdown in [false, true] { + let path = std::env::temp_dir() + .join(format!("hcancel-{}-{shutdown}.sock", std::process::id())); + let listener = crate::ipc::bind_private_local_listener(&path).unwrap(); + let now = Instant::now(); + let mut supervisors = EndpointSupervisors::new(&[], now); + supervisors.add_local(path.clone(), None, now); + let (events, mut received) = tokio::sync::mpsc::channel(4); + supervisors.spawn_due(now, connect_options(), &events); + let mut peer = listener.accept().unwrap(); + let _: ClientMessage = + protocol::read_message(&mut peer, protocol::MAX_FRAME_SIZE).unwrap(); + drop(listener); + std::fs::remove_file(path).unwrap(); + let started = Instant::now(); + if shutdown { + drop(supervisors); + } else { + supervisors.add_local(PathBuf::from("replacement"), None, now); + let event = tokio::time::timeout(Duration::from_secs(2), received.recv()) + .await + .unwrap() + .unwrap(); + let EndpointSupervisorEvent::Status { + endpoint_id, + generation, + status, + .. + } = event + else { + panic!("cancelled attempt connected"); + }; + assert!(!supervisors.record_status(&endpoint_id, generation, status, now)); + } + peer.set_nonblocking(true).unwrap(); + let deadline = Instant::now() + Duration::from_secs(2); + loop { + if matches!( + crate::ipc::poll_local_stream_read_count(&mut peer, &mut [0; 1]).unwrap(), + crate::ipc::LocalStreamReadCount::Closed + ) { + break; + } + assert!( + Instant::now() < deadline, + "cancelled attempt retained its stream" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert!(started.elapsed() < Duration::from_secs(1)); + } + } use crate::client::endpoint::ProfileId; fn profile() -> super::super::SavedSshEndpoint { diff --git a/src/client/endpoint/writer.rs b/src/client/endpoint/writer.rs index f7d67454a..38c1261c4 100644 --- a/src/client/endpoint/writer.rs +++ b/src/client/endpoint/writer.rs @@ -75,6 +75,7 @@ impl NativeEndpointTransport { break; } } + drop(stream); })?; Ok(Self { sender, diff --git a/src/client/handshake.rs b/src/client/handshake.rs index 76d75c71a..7f2dcfc1f 100644 --- a/src/client/handshake.rs +++ b/src/client/handshake.rs @@ -103,14 +103,14 @@ fn direct_graphics_profile_allowed() -> bool { fn read_handshake_response( stream: &mut LocalStream, read_timeout: Duration, + stopped: &AtomicBool, ) -> Result { let deadline = Instant::now() + read_timeout; crate::ipc::set_local_stream_polling(stream, true).map_err(ClientError::ConnectionFailed)?; - let stopped = AtomicBool::new(false); let response = protocol::read_message( &mut super::transport::EndpointReader { stream, - stopped: &stopped, + stopped, deadline: Some(deadline), }, MAX_FRAME_SIZE, @@ -146,6 +146,7 @@ pub(crate) fn probe_endpoint_negotiation( false, false, false, // This probe uses an SSH bridge, not a shared filesystem. + &AtomicBool::new(false), ) .map_err(io::Error::other)?; Ok(super::endpoint::EndpointNegotiation::new( @@ -165,6 +166,9 @@ pub(super) struct HandshakePixelCapabilities { /// shells use the stable endpoint generation and negotiate whole codecs without /// comparing Herdr build versions. `local_transport` means the endpoint shares the /// client's filesystem, not merely that its bridge exposes a local socket. +// Cancellation extends the existing handshake call; regrouping unrelated display options +// would broaden this lifecycle fix across all client entry points. +#[allow(clippy::too_many_arguments)] pub(super) fn do_handshake( stream: &mut LocalStream, cols: u16, @@ -177,6 +181,7 @@ pub(super) fn do_handshake( mouse_capture: bool, surface_active: bool, local_transport: bool, + stopped: &AtomicBool, ) -> Result { stream .set_nonblocking(false) @@ -231,7 +236,7 @@ pub(super) fn do_handshake( } else { handshake_read_timeout() }; - let welcome = read_handshake_response(stream, read_timeout)?; + let welcome = read_handshake_response(stream, read_timeout, stopped)?; if endpoint_shell { let ServerMessage::EndpointControl { kind, data } = welcome else { @@ -356,7 +361,11 @@ mod tests { } }); let start = Instant::now(); - let result = read_handshake_response(&mut client, Duration::from_millis(100)); + let result = read_handshake_response( + &mut client, + Duration::from_millis(100), + &AtomicBool::new(false), + ); let elapsed = start.elapsed(); let _ = release.send(()); peer.join().unwrap(); @@ -430,6 +439,7 @@ mod tests { false, true, true, + &AtomicBool::new(false), ); peer.join().unwrap(); match outcome { diff --git a/src/client/mod.rs b/src/client/mod.rs index 14482134a..c4538d109 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -296,6 +296,7 @@ fn run_client_with_mode( loop_config.mouse_capture_active, true, !is_remote_client_process(), + &AtomicBool::new(false), ) .map_err(|error| io::Error::other(error.to_string()))?; if federated diff --git a/src/client/terminal_sessions.rs b/src/client/terminal_sessions.rs index 2cc7784b1..29a597a4b 100644 --- a/src/client/terminal_sessions.rs +++ b/src/client/terminal_sessions.rs @@ -1,4 +1,5 @@ use std::io::{self, BufRead, Write as _}; +use std::sync::atomic::AtomicBool; use base64::Engine; use interprocess::local_socket::traits::Stream as _; @@ -104,6 +105,7 @@ fn connect_terminal_session_stream( false, true, true, + &AtomicBool::new(false), ) { Ok(handshake) if handshake.encoding == RenderEncoding::TerminalAnsi => {} Ok(handshake) => { diff --git a/src/ipc.rs b/src/ipc.rs index f740b17f6..f0737dbd8 100644 --- a/src/ipc.rs +++ b/src/ipc.rs @@ -51,6 +51,47 @@ pub(crate) fn connect_local_stream(path: &Path) -> io::Result { } } +pub(crate) fn connect_local_stream_cancellable( + path: &Path, + stopped: &std::sync::atomic::AtomicBool, + deadline: std::time::Instant, +) -> io::Result { + use std::sync::atomic::Ordering; + use std::time::{Duration, Instant}; + loop { + if stopped.load(Ordering::Acquire) { + return Err(io::Error::new( + io::ErrorKind::ConnectionAborted, + "connection attempt cancelled", + )); + } + let remaining = deadline + .checked_duration_since(Instant::now()) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::TimedOut, "endpoint connection timed out") + })?; + let timeout = remaining.min(Duration::from_millis(100)); + #[cfg(unix)] + let result = { + use interprocess::local_socket::{ConnectOptions, GenericFilePath, ToFsName as _}; + ConnectOptions::new() + .name(path.to_fs_name::()?) + .wait_mode(interprocess::ConnectWaitMode::Timeout(timeout)) + .connect_sync() + }; + #[cfg(windows)] + let result = crate::platform::connect_local_stream_with_timeout(path, timeout); + match result { + Err(error) + if matches!( + error.kind(), + io::ErrorKind::TimedOut | io::ErrorKind::WouldBlock + ) => {} + result => return result, + } + } +} + pub(crate) fn bind_local_listener(path: &Path) -> io::Result { #[cfg(unix)] { diff --git a/src/platform/mod.rs b/src/platform/mod.rs index 0349d7c6f..3d480b9f3 100644 --- a/src/platform/mod.rs +++ b/src/platform/mod.rs @@ -445,7 +445,7 @@ pub(crate) mod unix_image_files; #[cfg(unix)] pub(crate) use unix_common::{ begin_cli_output, end_cli_output, forward_remote_bridge_stdio, ignore_server_hangup, - local_stream_peer_description, spawn_server_signal_monitor, RemoteBridgeWake, + local_stream_peer_description, spawn_server_signal_monitor, ChildPipe, RemoteBridgeWake, }; mod client_state; diff --git a/src/platform/unix_common.rs b/src/platform/unix_common.rs index 8e0757e0b..ce27ff33f 100644 --- a/src/platform/unix_common.rs +++ b/src/platform/unix_common.rs @@ -1,5 +1,56 @@ use std::path::{Path, PathBuf}; +pub(crate) struct ChildPipe(std::fs::File); + +impl ChildPipe { + pub(crate) fn new_reader(pipe: impl Into) -> std::io::Result { + use std::os::fd::AsRawFd as _; + let file = std::fs::File::from(pipe.into()); + let fd = file.as_raw_fd(); + let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) }; + if flags < 0 || unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) } < 0 { + return Err(std::io::Error::last_os_error()); + } + Ok(Self(file)) + } + + pub(crate) fn new_writer(pipe: impl Into) -> std::io::Result { + Self::new_reader(pipe) + } + + pub(crate) fn wait(&self, writing: bool) -> std::io::Result<()> { + use std::os::fd::AsRawFd as _; + let mut descriptor = libc::pollfd { + fd: self.0.as_raw_fd(), + events: if writing { libc::POLLOUT } else { libc::POLLIN }, + revents: 0, + }; + if unsafe { libc::poll(&mut descriptor, 1, 50) } < 0 { + let error = std::io::Error::last_os_error(); + if error.kind() != std::io::ErrorKind::Interrupted { + return Err(error); + } + } + Ok(()) + } +} + +impl std::io::Read for ChildPipe { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + std::io::Read::read(&mut self.0, buffer) + } +} + +impl std::io::Write for ChildPipe { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + std::io::Write::write(&mut self.0, buffer) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + pub(crate) fn classify_child_exit(status: &portable_pty::ExitStatus) -> super::ChildExitReason { if status.signal().is_some() { super::ChildExitReason::Interrupted diff --git a/src/platform/windows.rs b/src/platform/windows.rs index 3701c25ed..feb6d5398 100644 --- a/src/platform/windows.rs +++ b/src/platform/windows.rs @@ -23,6 +23,79 @@ pub(crate) use notifications::{ static ALLOW_UNELEVATED_CLIENTS: OnceLock = OnceLock::new(); +pub(crate) struct ChildPipe(std::fs::File); + +impl ChildPipe { + pub(crate) fn new_reader(pipe: impl Into) -> std::io::Result { + // Rust's stdout/stderr handles lack permission to change pipe mode. Peek before + // reading instead; each pipe has exactly one reader. + Ok(Self(std::fs::File::from(pipe.into()))) + } + + pub(crate) fn new_writer(pipe: impl Into) -> std::io::Result { + use windows_sys::Win32::System::Pipes::{SetNamedPipeHandleState, PIPE_NOWAIT}; + let file = std::fs::File::from(pipe.into()); + let mode = PIPE_NOWAIT; + if unsafe { SetNamedPipeHandleState(file.as_raw_handle(), &mode, null_mut(), null_mut()) } + == 0 + { + return Err(std::io::Error::last_os_error()); + } + Ok(Self(file)) + } + + pub(crate) fn wait(&self, _writing: bool) -> std::io::Result<()> { + std::thread::sleep(Duration::from_millis(2)); + Ok(()) + } +} + +impl std::io::Read for ChildPipe { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + use windows_sys::Win32::System::Pipes::PeekNamedPipe; + if buffer.is_empty() { + return Ok(0); + } + let mut available = 0; + if unsafe { + PeekNamedPipe( + self.0.as_raw_handle(), + null_mut(), + 0, + null_mut(), + &mut available, + null_mut(), + ) + } == 0 + { + let error = std::io::Error::last_os_error(); + return if crate::ipc::is_connection_closed_error(&error) { + Ok(0) + } else { + Err(error) + }; + } + if available == 0 { + return Err(std::io::ErrorKind::WouldBlock.into()); + } + let count = buffer.len().min(available as usize); + self.0.read(&mut buffer[..count]) + } +} + +impl std::io::Write for ChildPipe { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + match self.0.write(buffer) { + Ok(0) if !buffer.is_empty() => Err(std::io::ErrorKind::WouldBlock.into()), + result => result, + } + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + pub(crate) fn allow_unelevated_clients() { let _ = ALLOW_UNELEVATED_CLIENTS.set(true); } @@ -42,6 +115,20 @@ pub(crate) fn probe_local_server(path: &std::path::Path) -> std::io::Result<()> .map_err(local_server_connection_error) } +pub(crate) fn connect_local_stream_with_timeout( + path: &std::path::Path, + timeout: Duration, +) -> std::io::Result { + use interprocess::os::windows::named_pipe::{pipe_mode::Bytes, DuplexPipeStream}; + let name = format!(r"\\.\pipe\{}", path.to_string_lossy()); + DuplexPipeStream::::connect_by_path_with_wait_mode( + name.as_str(), + interprocess::ConnectWaitMode::Timeout(timeout), + ) + .map(|pipe| crate::ipc::LocalStream::NamedPipe(pipe.into())) + .map_err(local_server_connection_error) +} + pub(crate) fn local_server_security_descriptor( ) -> std::io::Result { use windows_sys::Win32::Security::{ @@ -3344,6 +3431,64 @@ mod tests { AllocConsole, FreeConsole, GetConsoleProcessList, GetConsoleWindow, }; + #[test] + fn busy_endpoint_connect_obeys_cancellation_and_deadline() { + use std::os::windows::io::{FromRawHandle, OwnedHandle}; + use std::sync::atomic::{AtomicBool, Ordering}; + use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE; + use windows_sys::Win32::Storage::FileSystem::PIPE_ACCESS_DUPLEX; + use windows_sys::Win32::System::Pipes::CreateNamedPipeW; + let path = std::path::PathBuf::from(format!("herdr-busy-{}", std::process::id())); + let name = super::wide_null(&format!(r"\\.\pipe\{}", path.display())); + let handle = unsafe { + CreateNamedPipeW( + name.as_ptr(), + PIPE_ACCESS_DUPLEX, + 0, + 1, + 512, + 512, + 0, + std::ptr::null(), + ) + }; + assert_ne!(handle, INVALID_HANDLE_VALUE); + let _listener = unsafe { OwnedHandle::from_raw_handle(handle) }; + let _occupant = + super::connect_local_stream_with_timeout(&path, Duration::from_millis(500)).unwrap(); + for cancel in [false, true] { + let stopped = Arc::new(AtomicBool::new(false)); + let worker_stop = Arc::clone(&stopped); + let cancel_worker = thread::spawn(move || { + if cancel { + thread::sleep(Duration::from_millis(50)); + worker_stop.store(true, Ordering::Release); + } + }); + let start = Instant::now(); + let result = crate::ipc::connect_local_stream_cancellable( + &path, + &stopped, + start + + if cancel { + Duration::from_secs(5) + } else { + Duration::from_millis(200) + }, + ); + cancel_worker.join().unwrap(); + assert_eq!( + result.unwrap_err().kind(), + if cancel { + std::io::ErrorKind::ConnectionAborted + } else { + std::io::ErrorKind::TimedOut + } + ); + assert!(start.elapsed() < Duration::from_secs(1)); + } + } + #[test] fn local_resources_authorize_account_without_admin_rights() { use interprocess::local_socket::traits::Listener as _; diff --git a/src/remote/attach.rs b/src/remote/attach.rs index da0e120a5..68699b075 100644 --- a/src/remote/attach.rs +++ b/src/remote/attach.rs @@ -1,6 +1,11 @@ //! Remote thin-client launcher over SSH command stdio. -use super::{args::*, process::wait_with_output_timeout, restart_policy::*, shell_quote}; +use super::{ + args::*, + process::{wait_with_output_timeout, CancellablePipe, OwnedSshChild, PipeControl}, + restart_policy::*, + shell_quote, +}; use base64::Engine as _; use std::collections::BTreeMap; use std::fs::{self, File}; @@ -9,7 +14,6 @@ use std::path::{Path, PathBuf}; use std::process::{Child, Command, Output, Stdio}; use interprocess::local_socket::traits::Listener as _; -#[cfg(all(test, unix))] use interprocess::local_socket::traits::Stream as _; use interprocess::local_socket::ListenerNonblockingMode; use interprocess::TryClone as _; @@ -81,6 +85,7 @@ pub(crate) fn run_remote(remote: RemoteLaunch) -> io::Result<()> { session_name.clone(), remote_ssh.options(), false, + Arc::new(AtomicBool::new(false)), )?; run_client_process( @@ -776,6 +781,7 @@ pub(super) struct RemoteSsh { session_name: String, managed_config: Option, noninteractive: bool, + pub(super) cancelled: Arc, } impl RemoteSsh { @@ -795,6 +801,7 @@ impl RemoteSsh { session_name, managed_config, noninteractive: false, + cancelled: Arc::new(AtomicBool::new(false)), } } @@ -843,8 +850,11 @@ impl RemoteSsh { } fn sh_output(&self, script: &str) -> io::Result { + if self.cancelled.load(Ordering::Acquire) { + return Err(super::process::cancelled()); + } let script = posix_remote_output_command(script); - let mut child = self + let child = self .command() .arg("/bin/sh -s") .stdin(Stdio::piped()) @@ -859,20 +869,19 @@ impl RemoteSsh { )?); } - let write_result = if let Some(mut stdin) = child.stdin.take() { - stdin.write_all(script.as_bytes()) - } else { - Err(io::Error::new( - io::ErrorKind::BrokenPipe, - "ssh bootstrap stdin missing", - )) - }; - let output = wait_with_output_timeout(child, NONINTERACTIVE_SSH_COMMAND_TIMEOUT)?; - write_result?; + let output = wait_with_output_timeout( + child, + Some(script.into_bytes()), + NONINTERACTIVE_SSH_COMMAND_TIMEOUT, + Arc::clone(&self.cancelled), + )?; normalize_remote_output(output) } fn framed_user_shell_output(&self, remote_command: &str) -> io::Result { + if self.cancelled.load(Ordering::Acquire) { + return Err(super::process::cancelled()); + } let mut command = self.command(); command // Windows OpenSSH can still read the console with stdin redirected to NUL. @@ -882,7 +891,12 @@ impl RemoteSsh { .stdout(Stdio::piped()) .stderr(Stdio::piped()); let output = if self.noninteractive { - wait_with_output_timeout(command.spawn()?, NONINTERACTIVE_SSH_COMMAND_TIMEOUT) + wait_with_output_timeout( + command.spawn()?, + None, + NONINTERACTIVE_SSH_COMMAND_TIMEOUT, + Arc::clone(&self.cancelled), + ) } else { output_with_forwarded_stderr(command.spawn()?, None) }?; @@ -2159,6 +2173,7 @@ fn probe_remote_endpoint( ssh.session_name.clone(), ssh.options(), true, + Arc::clone(&ssh.cancelled), )?; let mut stream = crate::ipc::connect_local_stream(&path)?; // Use the saved client's noninteractive path. This metadata-only attachment never @@ -2772,6 +2787,7 @@ impl SshStdioBridge { session_name: String, ssh_options: Option<&ManagedSshOptions>, noninteractive: bool, + should_stop: Arc, ) -> io::Result { Self::start_command( target, @@ -2785,6 +2801,7 @@ impl SshStdioBridge { local_socket, ssh_options, noninteractive, + should_stop, ) } @@ -2794,6 +2811,7 @@ impl SshStdioBridge { local_socket: PathBuf, ssh_options: Option<&ManagedSshOptions>, noninteractive: bool, + should_stop: Arc, ) -> io::Result { crate::ipc::prepare_socket_path(&local_socket, |path| { format!("remote bridge is already listening at {}", path.display()) @@ -2811,7 +2829,6 @@ impl SshStdioBridge { return Err(err); } - let should_stop = Arc::new(AtomicBool::new(false)); let thread_stop = Arc::clone(&should_stop); let thread_ssh_options = ssh_options.cloned(); let (failure_tx, failure_rx) = mpsc::sync_channel(1); @@ -3035,14 +3052,13 @@ pub(crate) fn bridge_upload_cancellation_for_test( } fn bridge_connection( - mut stream: crate::ipc::LocalStream, + stream: crate::ipc::LocalStream, target: &str, remote_command: &str, ssh_options: Option<&ManagedSshOptions>, noninteractive: bool, bridge_stop: &Arc, ) -> io::Result<()> { - let upload_stop = Arc::new(BridgeUploadStop::new()?); let mut command = Command::new("ssh"); apply_managed_ssh_options(&mut command, ssh_options); if noninteractive { @@ -3060,38 +3076,47 @@ fn bridge_connection( Stdio::inherit() }); - let mut child = command - .spawn() - .map_err(|err| io::Error::new(err.kind(), format!("failed to start ssh bridge: {err}")))?; - let mut child_stdin = match child.stdin.take() { - Some(stdin) => stdin, - None => return terminate_bridge_child(child, "ssh bridge stdin missing"), - }; - let child_stdout = match child.stdout.take() { - Some(stdout) => stdout, - None => return terminate_bridge_child(child, "ssh bridge stdout missing"), - }; - let stderr_reader = if noninteractive { - let Some(child_stderr) = child.stderr.take() else { - return terminate_bridge_child(child, "ssh bridge stderr missing"); - }; - Some(thread::spawn(move || capture_ssh_stderr(child_stderr))) - } else { - None - }; - let stream_to_child = match stream.try_clone() { - Ok(stream) => stream, - Err(err) => { - let _ = child.kill(); - let _ = child.wait(); - return Err(err); - } - }; - if let Err(err) = crate::ipc::set_local_stream_polling(&mut stream, true) { - let _ = child.kill(); - let _ = child.wait(); - return Err(err); - } + bridge_child(stream, command, bridge_stop) +} + +fn bridge_child( + stream: crate::ipc::LocalStream, + mut command: Command, + bridge_stop: &Arc, +) -> io::Result<()> { + let upload_stop = Arc::new(BridgeUploadStop::new()?); + + let mut child = + OwnedSshChild(command.spawn().map_err(|err| { + io::Error::new(err.kind(), format!("failed to start ssh bridge: {err}")) + })?); + let pipe_control = PipeControl::new(Arc::clone(bridge_stop)); + let stdin = child + .stdin + .take() + .ok_or_else(|| io::Error::other("ssh bridge stdin missing"))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| io::Error::other("ssh bridge stdout missing"))?; + let mut child_stdin = CancellablePipe::new( + crate::platform::ChildPipe::new_writer(stdin)?, + &pipe_control, + ); + let child_stdout = CancellablePipe::new( + crate::platform::ChildPipe::new_reader(stdout)?, + &pipe_control, + ); + let stderr = child + .stderr + .take() + .map(crate::platform::ChildPipe::new_reader) + .transpose()? + .map(|pipe| CancellablePipe::new(pipe, &pipe_control)); + let stream_to_child = stream.try_clone()?; + // Both directions poll their own reads; writes must also remain cancellable on Windows. + stream.set_nonblocking(true)?; + let stderr_reader = stderr.map(|stderr| thread::spawn(move || capture_ssh_stderr(stderr))); let mut child_to_stream = stream; let connection_stop = Arc::new(AtomicBool::new(false)); @@ -3137,6 +3162,7 @@ fn bridge_connection( match child.try_wait() { Ok(Some(status)) => { upload_stop.cancel(); + pipe_control.finish_reads(); break (Ok(status), true); } Ok(None) => {} @@ -3171,6 +3197,7 @@ fn bridge_connection( upload_stop.cancel(); if !child_exited { connection_stop.store(true, Ordering::Release); + pipe_control.stop(); } let upload_result = upload .join() @@ -3274,12 +3301,6 @@ fn discard_remote_output_preamble(reader: &mut impl io::BufRead) -> io::Result<( } } -fn terminate_bridge_child(mut child: std::process::Child, message: &'static str) -> io::Result<()> { - let _ = child.kill(); - let _ = child.wait(); - Err(io::Error::new(io::ErrorKind::BrokenPipe, message)) -} - fn copy_reader_to_local_stream( reader: &mut R, stream: &mut crate::ipc::LocalStream, @@ -3456,6 +3477,79 @@ fn sanitize_path_component(input: &str) -> String { mod tests { use super::*; + #[test] + fn bridge_cancellation_releases_stalled_preamble_upload_and_download() { + for mode in ["bridge-preamble", "bridge-input", "bridge-output"] { + let socket = local_forward_socket_path(mode, "cancellation-test"); + let listener = crate::ipc::bind_private_local_listener(&socket).unwrap(); + let mut client = crate::ipc::connect_local_stream(&socket).unwrap(); + let server = listener.accept().unwrap(); + drop(listener); + fs::remove_file(socket).unwrap(); + client.set_nonblocking(true).unwrap(); + let fixture = super::super::process::tests::PipeFixture::new(mode); + let mut command = fixture.command(); + command + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + let stop = Arc::new(AtomicBool::new(false)); + let worker_stop = Arc::clone(&stop); + let (done_tx, done_rx) = mpsc::channel(); + let worker = thread::spawn(move || { + done_tx + .send(bridge_child(server, command, &worker_stop)) + .unwrap(); + }); + let descendant = fixture.ready(); + // Fill upload buffers while the fixture deliberately never reads stdin. + let deadline = Instant::now() + Duration::from_millis(150); + while Instant::now() < deadline { + if mode == "bridge-input" { + let _ = client.write(&[b'x'; 16 * 1024]); + } + thread::sleep(Duration::from_millis(1)); + } + stop.store(true, Ordering::Release); + let result = done_rx + .recv_timeout(Duration::from_secs(5)) + .expect("bridge worker completes"); + worker.join().unwrap(); + assert!(result.is_ok(), "{mode}: {result:?}"); + assert!( + crate::platform::process_exists(descendant), + "{mode}: cleanup waited for or killed the independent holder" + ); + } + } + + #[test] + fn bridge_drains_final_output_while_descendant_retains_pipes() { + let socket = local_forward_socket_path("final-output", "cancellation-test"); + let listener = crate::ipc::bind_private_local_listener(&socket).unwrap(); + let mut client = crate::ipc::connect_local_stream(&socket).unwrap(); + let server = listener.accept().unwrap(); + drop(listener); + fs::remove_file(socket).unwrap(); + let fixture = super::super::process::tests::PipeFixture::new("bridge-final"); + let mut command = fixture.command(); + command + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + let worker = + thread::spawn(move || bridge_child(server, command, &Arc::new(AtomicBool::new(false)))); + let descendant = fixture.ready(); + let mut output = Vec::new(); + client.read_to_end(&mut output).unwrap(); + worker.join().unwrap().unwrap(); + assert!( + crate::platform::process_exists(descendant), + "cleanup waited for or killed the independent holder" + ); + assert!(String::from_utf8_lossy(&output).contains("final child output")); + } + #[test] fn remote_client_preferences_cross_process_handoff() { let root = @@ -3675,6 +3769,7 @@ mod tests { "default".to_string(), None, false, + Arc::new(AtomicBool::new(false)), ) .expect("start bridge listener"); @@ -3796,6 +3891,7 @@ mod tests { "default".to_string(), None, false, + Arc::new(AtomicBool::new(false)), ) .expect("start bridge listener"); let started = Instant::now(); @@ -3990,6 +4086,7 @@ mod tests { session_name: crate::session::DEFAULT_SESSION_NAME.into(), managed_config: Some(managed_config), noninteractive: false, + cancelled: Arc::new(AtomicBool::new(false)), }; let command = ssh.command(); @@ -4061,6 +4158,7 @@ mod tests { session_name: crate::session::DEFAULT_SESSION_NAME.into(), managed_config: Some(managed_config), noninteractive: false, + cancelled: Arc::new(AtomicBool::new(false)), }; let args = ssh .command() @@ -4140,6 +4238,7 @@ mod tests { session_name: "probe-options".into(), managed_config: Some(managed_config), noninteractive: false, + cancelled: Arc::new(AtomicBool::new(false)), }; let remote = RemoteHerdr::for_platform(RemotePlatform { os: "linux", @@ -4271,6 +4370,7 @@ mod tests { session_name: crate::session::DEFAULT_SESSION_NAME.into(), managed_config: None, noninteractive: false, + cancelled: Arc::new(AtomicBool::new(false)), }; let command = ssh.command(); diff --git a/src/remote/process.rs b/src/remote/process.rs index d3db5c9f3..3e7f9eb25 100644 --- a/src/remote/process.rs +++ b/src/remote/process.rs @@ -1,87 +1,414 @@ -use std::io::{self, Read as _}; -use std::process::Output; +use std::io::{self, Read as _, Write as _}; +use std::process::{Child, Output}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; use std::thread; use std::time::{Duration, Instant}; -const POLL_INTERVAL: Duration = Duration::from_millis(50); +const POLL_INTERVAL: Duration = Duration::from_millis(10); -pub(super) fn wait_with_output_timeout( - mut child: std::process::Child, - timeout: Duration, -) -> io::Result { - let stdout = child - .stdout - .take() - .ok_or_else(|| io::Error::other("SSH command stdout was not captured"))?; - let stderr = child - .stderr - .take() - .ok_or_else(|| io::Error::other("SSH command stderr was not captured"))?; - let stdout = thread::spawn(move || { - let mut stdout = stdout; - let mut bytes = Vec::new(); - stdout.read_to_end(&mut bytes).map(|_| bytes) - }); - let stderr = thread::spawn(move || { - let mut stderr = stderr; - let mut bytes = Vec::new(); - stderr.read_to_end(&mut bytes).map(|_| bytes) - }); - let started = Instant::now(); - let status = loop { - match child.try_wait() { - Ok(Some(status)) => break status, - Ok(None) => {} - Err(error) => { - let _ = child.kill(); - let _ = child.wait(); - let _ = stdout.join(); - let _ = stderr.join(); - return Err(error); +pub(super) fn cancelled() -> io::Error { + io::Error::new( + io::ErrorKind::ConnectionAborted, + "SSH connection attempt cancelled", + ) +} + +pub(super) struct OwnedSshChild(pub(super) Child); + +impl std::ops::Deref for OwnedSshChild { + type Target = Child; + fn deref(&self) -> &Child { + &self.0 + } +} + +impl std::ops::DerefMut for OwnedSshChild { + fn deref_mut(&mut self) -> &mut Child { + &mut self.0 + } +} + +impl Drop for OwnedSshChild { + fn drop(&mut self) { + // Only this SSH invocation is ours; a multiplex master may serve other sessions. + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +pub(super) struct PipeControl { + cancelled: Arc, + stopped: AtomicBool, + finished: AtomicBool, +} + +impl PipeControl { + pub(super) fn new(cancelled: Arc) -> Arc { + Arc::new(Self { + cancelled, + stopped: AtomicBool::new(false), + finished: AtomicBool::new(false), + }) + } + + pub(super) fn stop(&self) { + self.stopped.store(true, Ordering::Release); + } + pub(super) fn finish_reads(&self) { + self.finished.store(true, Ordering::Release); + } + + fn check(&self) -> io::Result<()> { + if self.stopped.load(Ordering::Acquire) || self.cancelled.load(Ordering::Acquire) { + Err(cancelled()) + } else { + Ok(()) + } + } +} + +pub(super) struct CancellablePipe { + pipe: crate::platform::ChildPipe, + control: Arc, +} + +impl CancellablePipe { + pub(super) fn new(pipe: crate::platform::ChildPipe, control: &Arc) -> Self { + Self { + pipe, + control: Arc::clone(control), + } + } +} + +impl io::Read for CancellablePipe { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + loop { + if self.control.check().is_err() { + return Ok(0); + } + match self.pipe.read(buffer) { + Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + // Drain buffered output, then stop even if a descendant kept the pipe. + if self.control.finished.load(Ordering::Acquire) { + return Ok(0); + } + self.pipe.wait(false)?; + } + result => return result, } } - if started.elapsed() >= timeout { - let _ = child.kill(); - let _ = child.wait(); - let _ = stdout.join(); - let _ = stderr.join(); - return Err(io::Error::new( - io::ErrorKind::TimedOut, - "noninteractive SSH command timed out", - )); + } +} + +impl io::Write for CancellablePipe { + fn write(&mut self, buffer: &[u8]) -> io::Result { + loop { + self.control.check()?; + if self.control.finished.load(Ordering::Acquire) { + return Err(io::ErrorKind::BrokenPipe.into()); + } + match self.pipe.write(buffer) { + Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => self.pipe.wait(true)?, + result => return result, + } } - thread::sleep(POLL_INTERVAL); - }; - let stdout = stdout - .join() - .map_err(|_| io::Error::other("SSH stdout reader panicked"))??; - let stderr = stderr - .join() - .map_err(|_| io::Error::other("SSH stderr reader panicked"))??; - Ok(Output { - status, - stdout, - stderr, + } + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +pub(super) fn wait_with_output_timeout( + child: Child, + input: Option>, + timeout: Duration, + cancelled: Arc, +) -> io::Result { + let mut child = OwnedSshChild(child); + let control = PipeControl::new(cancelled); + let stdout = crate::platform::ChildPipe::new_reader( + child + .0 + .stdout + .take() + .ok_or_else(|| io::Error::other("SSH command stdout was not captured"))?, + )?; + let stderr = crate::platform::ChildPipe::new_reader( + child + .0 + .stderr + .take() + .ok_or_else(|| io::Error::other("SSH command stderr was not captured"))?, + )?; + let stdin = child + .0 + .stdin + .take() + .map(crate::platform::ChildPipe::new_writer) + .transpose()?; + let mut stdout = CancellablePipe::new(stdout, &control); + let mut stderr = CancellablePipe::new(stderr, &control); + let stdin = stdin.map(|pipe| CancellablePipe::new(pipe, &control)); + thread::scope(|scope| { + let stdout = scope.spawn(move || { + let mut bytes = Vec::new(); + stdout.read_to_end(&mut bytes).map(|_| bytes) + }); + let stderr = scope.spawn(move || { + let mut bytes = Vec::new(); + stderr.read_to_end(&mut bytes).map(|_| bytes) + }); + let stdin = scope.spawn(move || match (stdin, input) { + (Some(mut stdin), Some(input)) => stdin.write_all(&input), + (None, Some(_)) => Err(io::Error::other("ssh bootstrap stdin missing")), + _ => Ok(()), + }); + let deadline = Instant::now() + timeout; + let mut status = None; + let result = loop { + if let Err(error) = control.check() { + break Err(error); + } + if status.is_none() { + match child.0.try_wait() { + Ok(Some(exited)) => { + status = Some(exited); + control.finish_reads(); + } + Ok(None) => {} + Err(error) => break Err(error), + } + } + if let Some(status) = status { + if stdout.is_finished() && stderr.is_finished() && stdin.is_finished() { + break Ok(status); + } + } + if Instant::now() >= deadline { + break Err(io::Error::new( + io::ErrorKind::TimedOut, + "noninteractive SSH command timed out", + )); + } + thread::sleep(POLL_INTERVAL); + }; + if result.is_err() { + control.stop(); + } + drop(child); + let stdout = stdout + .join() + .map_err(|_| io::Error::other("SSH stdout reader panicked")); + let stderr = stderr + .join() + .map_err(|_| io::Error::other("SSH stderr reader panicked")); + let stdin = stdin + .join() + .map_err(|_| io::Error::other("SSH stdin writer panicked")); + let status = result?; + let stdout = stdout??; + let stderr = stderr??; + // A rejected command can close stdin early; its SSH error is more useful. + if status.success() { + stdin??; + } + Ok(Output { + status, + stdout, + stderr, + }) }) } -#[cfg(all(test, unix))] -mod tests { +#[cfg(test)] +pub(in crate::remote) mod tests { use super::*; use std::process::{Command, Stdio}; - #[test] - fn timeout_kills_the_child() { - let mut command = Command::new("sh"); + pub(in crate::remote) struct PipeFixture { + root: std::path::PathBuf, + mode: &'static str, + } + + impl PipeFixture { + pub(in crate::remote) fn new(mode: &'static str) -> Self { + static NEXT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + let root = std::env::temp_dir().join(format!( + "herdr-pipes-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + std::fs::create_dir(&root).unwrap(); + Self { root, mode } + } + + pub(in crate::remote) fn command(&self) -> Command { + let mut command = fixture_command(self.mode); + command.env("HERDR_PIPE_FIXTURE_ROOT", &self.root); + command + } + + pub(in crate::remote) fn ready(&self) -> u32 { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if let Some(pid) = std::fs::read_to_string(self.root.join("pid")) + .ok() + .and_then(|pid| pid.parse().ok()) + { + return pid; + } + assert!(Instant::now() < deadline, "pipe fixture did not start"); + thread::sleep(POLL_INTERVAL); + } + } + } + + impl Drop for PipeFixture { + fn drop(&mut self) { + let _ = std::fs::write(self.root.join("release"), b""); + if let Some(pid) = std::fs::read_to_string(self.root.join("pid")) + .ok() + .and_then(|pid| pid.parse().ok()) + { + let deadline = Instant::now() + Duration::from_secs(5); + while crate::platform::process_exists(pid) && Instant::now() < deadline { + thread::sleep(POLL_INTERVAL); + } + } + let _ = std::fs::remove_dir_all(&self.root); + } + } + + pub(in crate::remote) fn fixture_command(mode: &str) -> Command { + let mut command = Command::new(std::env::current_exe().unwrap()); command - .arg("-c") - .arg("exec sleep 10") + .args([ + "--exact", + "remote::process::tests::pipe_fixture", + "--nocapture", + ]) + .env("HERDR_PIPE_FIXTURE", mode); + crate::platform::configure_background_command(&mut command); + command + } + + #[test] + // This bounded fixture deliberately leaves inherited pipes open after its parent exits. + #[allow(clippy::zombie_processes)] + fn pipe_fixture() { + let Ok(mode) = std::env::var("HERDR_PIPE_FIXTURE") else { + return; + }; + let root = std::path::PathBuf::from(std::env::var_os("HERDR_PIPE_FIXTURE_ROOT").unwrap()); + let release = root.join("release"); + thread::spawn(move || { + let deadline = Instant::now() + Duration::from_secs(30); + while !release.exists() && Instant::now() < deadline { + thread::sleep(POLL_INTERVAL); + } + std::process::exit(0); + }); + if mode == "holder" { + // Acknowledge only after libtest's startup output, before the parent can close + // the captured pipes. Otherwise the harness itself can exit on a broken pipe. + std::fs::write(root.join("pid"), std::process::id().to_string()).unwrap(); + } else { + fixture_command("holder").spawn().unwrap(); + // Keep the parent alive and the pipe empty until the holder is actually ready. + while !root.join("pid").exists() { + thread::sleep(POLL_INTERVAL); + } + if mode.starts_with("bridge-") && mode != "bridge-preamble" { + println!("herdr-remote-output-ready:1"); + } + if mode == "bridge-output" { + io::stdout() + .write_all(&vec![b'x'; 2 * 1024 * 1024]) + .unwrap(); + } + println!("final child output"); + eprintln!("Permission denied (publickey)"); + } + if mode != "exit" && mode != "bridge-final" { + thread::sleep(Duration::from_secs(30)); + } + } + + #[test] + fn child_exit_drains_output_without_waiting_for_inherited_pipes() { + let fixture = PipeFixture::new("exit"); + let child = fixture + .command() .stdout(Stdio::piped()) - .stderr(Stdio::piped()); - let started = Instant::now(); - let error = wait_with_output_timeout(command.spawn().unwrap(), Duration::from_millis(25)) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let descendant = fixture.ready(); + let output = wait_with_output_timeout( + child, + None, + Duration::from_secs(5), + Arc::new(AtomicBool::new(false)), + ) + .unwrap(); + assert!(output.status.success()); + assert!( + crate::platform::process_exists(descendant), + "cleanup must not kill a shared descendant" + ); + assert!(String::from_utf8_lossy(&output.stdout).contains("final child output")); + assert!(String::from_utf8_lossy(&output.stderr).contains("Permission denied (publickey)")); + } + + #[test] + fn cancellation_and_timeout_reap_child_with_blocked_pipes() { + for cancel in [false, true] { + let fixture = PipeFixture::new("wait"); + let child = fixture + .command() + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let descendant = fixture.ready(); + let stopped = Arc::new(AtomicBool::new(false)); + let pid = child.id(); + let cancel_stop = Arc::clone(&stopped); + let cancel_thread = thread::spawn(move || { + if cancel { + thread::sleep(Duration::from_millis(100)); + cancel_stop.store(true, Ordering::Release); + } + }); + let error = wait_with_output_timeout( + child, + Some(vec![b'x'; 1024 * 1024]), + Duration::from_millis(200), + stopped, + ) .unwrap_err(); - assert_eq!(error.kind(), io::ErrorKind::TimedOut); - assert!(started.elapsed() < Duration::from_secs(1)); + cancel_thread.join().unwrap(); + assert!( + !crate::platform::process_exists(pid), + "owned child was not reaped" + ); + assert_eq!( + error.kind(), + if cancel { + io::ErrorKind::ConnectionAborted + } else { + io::ErrorKind::TimedOut + } + ); + assert!( + crate::platform::process_exists(descendant), + "cleanup waited for or killed the independent holder" + ); + } } } diff --git a/src/remote/saved.rs b/src/remote/saved.rs index 4d8bd74b6..5dfde111b 100644 --- a/src/remote/saved.rs +++ b/src/remote/saved.rs @@ -1,5 +1,10 @@ use std::io; use std::path::PathBuf; +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, +}; +use std::time::{Duration, Instant}; use super::attach::{find_installed_remote_herdr, RemoteSsh, SshStdioBridge}; @@ -16,11 +21,17 @@ pub(crate) fn connect_saved_ssh( profile_id: &str, target: &str, session: &str, + generation: u64, + cancelled: Arc, ) -> io::Result { - let ssh = validated_saved_ssh(profile_id, target, session)?; + let mut ssh = validated_saved_ssh(profile_id, target, session)?; + ssh.cancelled = Arc::clone(&cancelled); let remote_herdr = find_installed_remote_herdr(&ssh)?; let metadata = remote_herdr.machine_metadata(); - let path = saved_bridge_path(profile_id); + if cancelled.load(Ordering::Acquire) { + return Err(super::process::cancelled()); + } + let path = saved_bridge_path(profile_id, generation); let bridge = SshStdioBridge::start( target.to_owned(), remote_herdr, @@ -28,8 +39,13 @@ pub(crate) fn connect_saved_ssh( session.to_owned(), ssh.options(), true, + Arc::clone(&cancelled), + )?; + let stream = crate::ipc::connect_local_stream_cancellable( + &path, + &cancelled, + Instant::now() + Duration::from_secs(5), )?; - let stream = crate::ipc::connect_local_stream(&path)?; if let Some(metadata) = metadata { crate::client::endpoint::SshMetadataCache::new(profile_id, target, session)? .store(&metadata); @@ -82,6 +98,7 @@ impl SavedSshApiBridge { path.clone(), ssh.options(), true, + Arc::new(AtomicBool::new(false)), )?; Ok(Self { path, @@ -145,10 +162,10 @@ pub(crate) fn saved_ssh_failure_needs_attention(error: &io::Error) -> bool { .any(|needle| message.contains(needle)) } -fn saved_bridge_path(profile_id: &str) -> PathBuf { +fn saved_bridge_path(profile_id: &str, generation: u64) -> PathBuf { let pid = std::process::id(); - let readable = format!("herdr-ssh-{pid}-{profile_id}.sock"); - let short = format!("herdr-s-{pid}-{}.sock", &profile_id[..16]); + let readable = format!("herdr-ssh-{pid}-{profile_id}-{generation:x}.sock"); + let short = format!("herdr-s-{pid}-{}-{generation:x}.sock", &profile_id[..16]); crate::platform::remote_bridge_endpoint_path(&readable, &short) } @@ -179,12 +196,32 @@ mod tests { use super::*; #[test] - fn bridge_paths_use_profile_identity_not_target_or_session() { - let first = saved_bridge_path("0123456789abcdef0123456789abcdef"); - let second = saved_bridge_path("fedcba9876543210fedcba9876543210"); - assert_ne!(first, second); - assert!(!first.to_string_lossy().contains("example.com")); - assert!(!first.to_string_lossy().contains("default")); + fn retiring_old_bridge_does_not_remove_or_block_replacement_listener() { + use interprocess::local_socket::traits::Listener as _; + use std::io::{Read as _, Write as _}; + let profile = "0123456789abcdef0123456789abcdef"; + let first = saved_bridge_path(profile, 1); + let second = saved_bridge_path(profile, 2); + let old = SshStdioBridge::start_command( + "unused".into(), + "unused".into(), + first.clone(), + None, + true, + Arc::new(AtomicBool::new(false)), + ) + .unwrap(); + let replacement = crate::ipc::bind_private_local_listener(&second).unwrap(); + drop(old); + assert!(!first.exists()); + let mut client = crate::ipc::connect_local_stream(&second).unwrap(); + let mut server = replacement.accept().unwrap(); + client.write_all(b"replacement").unwrap(); + let mut bytes = [0; 11]; + server.read_exact(&mut bytes).unwrap(); + assert_eq!(&bytes, b"replacement"); + drop(replacement); + std::fs::remove_file(second).unwrap(); } #[test]