diff --git a/src/client/transport.rs b/src/client/transport.rs index 4807c31d..2320598b 100644 --- a/src/client/transport.rs +++ b/src/client/transport.rs @@ -144,3 +144,83 @@ pub(super) fn write_to_server( ) -> io::Result<()> { stream.send_client_message(msg) } + +#[cfg(all(test, unix))] +mod tests { + use super::*; + use crate::client::endpoint::EndpointTransport as _; + use interprocess::local_socket::traits::Listener as _; + use std::io::{Read as _, Write as _}; + use std::time::Instant; + + #[test] + fn upload_cancellation_preserves_pending_endpoint_download() { + let path = std::env::temp_dir().join(format!( + "herdr-cancel-{}-{}.sock", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let listener = crate::ipc::bind_private_local_listener(&path).unwrap(); + let client = crate::ipc::connect_local_stream(&path).unwrap(); + let mut bridge = listener.accept().unwrap(); + std::fs::remove_file(path).unwrap(); + drop(listener); + let mut reader_stream = client.try_clone().unwrap(); + let mut writer = endpoint::NativeEndpointTransport::with_lifetime(client, ()).unwrap(); + let stopped = writer.stop_handle(); + struct ForwardedInput(std::sync::mpsc::Sender>); + impl io::Write for ForwardedInput { + fn write(&mut self, bytes: &[u8]) -> io::Result { + self.0.send(bytes.to_vec()).unwrap(); + Ok(bytes.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + let (forwarded_tx, forwarded_rx) = std::sync::mpsc::channel(); + let cancel = crate::remote::bridge_upload_cancellation_for_test( + bridge.try_clone().unwrap(), + ForwardedInput(forwarded_tx), + ); + let message = ClientMessage::ClientShellFocus { focused: false }; + let mut expected = Vec::new(); + protocol::write_message(&mut expected, &message).unwrap(); + writer.send(&message).unwrap(); + let mut forwarded = Vec::new(); + while forwarded.len() < expected.len() { + forwarded.extend(forwarded_rx.recv_timeout(Duration::from_secs(3)).unwrap()); + } + assert_eq!(forwarded, expected); + cancel(); + + // A client write after upload cancellation must not stop the download reader. + writer + .send(&ClientMessage::ClientShellFocus { focused: true }) + .unwrap(); + let flushed = writer.flush(Instant::now() + Duration::from_secs(3)); + if flushed.is_ok() { + let received: ClientMessage = + protocol::read_message(&mut bridge, protocol::MAX_FRAME_SIZE).unwrap(); + assert_eq!(received, ClientMessage::ClientShellFocus { focused: true }); + } + const FINAL: &[u8] = b"pending-download: FINAL OUTPUT\n"; + bridge.write_all(FINAL).unwrap(); + drop(bridge); + let mut output = Vec::new(); + EndpointReader { + stream: &mut reader_stream, + stopped: &stopped, + } + .read_to_end(&mut output) + .unwrap(); + assert_eq!(output, FINAL); + assert!(flushed.is_ok(), "client write failed: {flushed:?}"); + assert!(!stopped.load(Ordering::Acquire)); + assert!(writer.take_error().is_none()); + } +} diff --git a/src/platform/mod.rs b/src/platform/mod.rs index d48ad4fa..335293e8 100644 --- a/src/platform/mod.rs +++ b/src/platform/mod.rs @@ -264,7 +264,7 @@ pub(crate) struct RemoteSshConfigPaths { #[cfg(unix)] mod unix_common; #[cfg(unix)] -pub(crate) use unix_common::{begin_cli_output, end_cli_output}; +pub(crate) use unix_common::{begin_cli_output, end_cli_output, RemoteBridgeWake}; mod client_state; pub(crate) use client_state::{create_private_state_file, replace_file, sync_parent_directory}; diff --git a/src/platform/unix_common.rs b/src/platform/unix_common.rs index 104f6b69..0df215d6 100644 --- a/src/platform/unix_common.rs +++ b/src/platform/unix_common.rs @@ -27,6 +27,50 @@ pub(crate) fn wait_client_stream_readable(stream: &crate::ipc::LocalStream) -> s Ok(()) } +pub(crate) struct RemoteBridgeWake { + reader: std::os::unix::net::UnixStream, + writer: std::os::unix::net::UnixStream, +} + +impl RemoteBridgeWake { + pub(crate) fn new() -> std::io::Result { + let (reader, writer) = std::os::unix::net::UnixStream::pair()?; + Ok(Self { reader, writer }) + } + + pub(crate) fn cancel(&self) -> std::io::Result<()> { + // EOF stays readable, including when cancellation precedes the wait. + self.writer.shutdown(std::net::Shutdown::Write) + } + + pub(crate) fn wait(&self, stream: &crate::ipc::LocalStream) -> std::io::Result<()> { + use std::os::fd::{AsFd as _, AsRawFd as _}; + let crate::ipc::LocalStream::UdSocket(stream) = stream; + let mut descriptors = [ + libc::pollfd { + fd: stream.as_fd().as_raw_fd(), + events: libc::POLLIN, + revents: 0, + }, + libc::pollfd { + fd: self.reader.as_raw_fd(), + events: libc::POLLIN, + revents: 0, + }, + ]; + loop { + // SAFETY: both descriptors remain borrowed and the array has two entries. + if unsafe { libc::poll(descriptors.as_mut_ptr(), 2, -1) } >= 0 { + return Ok(()); + } + let error = std::io::Error::last_os_error(); + if error.kind() != std::io::ErrorKind::Interrupted { + return Err(error); + } + } + } +} + pub(super) fn read_terminal_grid_size() -> std::io::Result<(u16, u16)> { crossterm::terminal::window_size().map(|size| (size.columns, size.rows)) } diff --git a/src/platform/windows.rs b/src/platform/windows.rs index 4deb1864..e58bacae 100644 --- a/src/platform/windows.rs +++ b/src/platform/windows.rs @@ -24,6 +24,25 @@ pub(crate) fn classify_child_exit(status: &portable_pty::ExitStatus) -> super::C } } +pub(crate) struct RemoteBridgeWake; + +impl RemoteBridgeWake { + pub(crate) fn new() -> std::io::Result { + Ok(Self) + } + + pub(crate) fn cancel(&self) -> std::io::Result<()> { + // The named-pipe reader checks its cancellation flag between peeks. + Ok(()) + } + + pub(crate) fn wait(&self, _stream: &crate::ipc::LocalStream) -> std::io::Result<()> { + // Synchronous named pipes still use peek-before-read polling on Windows. + std::thread::sleep(Duration::from_millis(1)); + Ok(()) + } +} + pub(crate) fn wait_client_stream_readable( _stream: &crate::ipc::LocalStream, ) -> std::io::Result<()> { diff --git a/src/remote/attach.rs b/src/remote/attach.rs index 10db88a6..46c0165d 100644 --- a/src/remote/attach.rs +++ b/src/remote/attach.rs @@ -1941,6 +1941,63 @@ fn write_managed_ssh_config() -> io::Result { }) } +struct BridgeUploadStop { + stopped: AtomicBool, + wake: crate::platform::RemoteBridgeWake, +} + +impl BridgeUploadStop { + fn new() -> io::Result { + Ok(Self { + stopped: AtomicBool::new(false), + wake: crate::platform::RemoteBridgeWake::new()?, + }) + } + + fn cancel(&self) { + if !self.stopped.swap(true, Ordering::AcqRel) { + if let Err(error) = self.wake.cancel() { + tracing::debug!(%error, "remote bridge read cancellation failed"); + } + } + } + + fn is_stopped(&self) -> bool { + self.stopped.load(Ordering::Acquire) + } +} + +#[cfg(all(test, unix))] +pub(crate) fn bridge_upload_cancellation_for_test( + stream: crate::ipc::LocalStream, + mut writer: impl io::Write + Send + 'static, +) -> impl FnOnce() { + stream.set_nonblocking(true).unwrap(); + let stop = Arc::new(BridgeUploadStop::new().unwrap()); + let worker_stop = Arc::clone(&stop); + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + let closed = AtomicBool::new(false); + let result = copy_local_stream_to_writer( + stream, + &mut writer, + &worker_stop, + &AtomicBool::new(false), + &closed, + ); + done_tx + .send((result, closed.load(Ordering::Acquire))) + .unwrap(); + }); + move || { + stop.cancel(); + let (result, closed) = done_rx.recv_timeout(Duration::from_secs(3)).unwrap(); + worker.join().unwrap(); + result.unwrap(); + assert!(!closed, "upload cancellation must not report peer EOF"); + } +} + fn bridge_connection( stream: crate::ipc::LocalStream, target: &str, @@ -1950,6 +2007,7 @@ fn bridge_connection( 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 { @@ -1994,7 +2052,6 @@ fn bridge_connection( let mut child_to_stream = stream; let connection_stop = Arc::new(AtomicBool::new(false)); - let upload_stop = Arc::new(AtomicBool::new(false)); let upload_failed = Arc::new(AtomicBool::new(false)); let download_done = Arc::new(AtomicBool::new(false)); let client_closed = Arc::new(AtomicBool::new(false)); @@ -2025,7 +2082,7 @@ fn bridge_connection( &download_bridge_stop, ); download_done_worker.store(true, Ordering::Release); - download_upload_stop.store(true, Ordering::Release); + download_upload_stop.cancel(); result }); @@ -2033,13 +2090,13 @@ fn bridge_connection( let (status_result, child_exited) = loop { match child.try_wait() { Ok(Some(status)) => { - upload_stop.store(true, Ordering::Release); + upload_stop.cancel(); break (Ok(status), true); } Ok(None) => {} Err(err) => { connection_stop.store(true, Ordering::Release); - upload_stop.store(true, Ordering::Release); + upload_stop.cancel(); let _ = child.kill(); let _ = child.wait(); break (Err(err), false); @@ -2047,7 +2104,7 @@ fn bridge_connection( } if bridge_stop.load(Ordering::Acquire) { connection_stop.store(true, Ordering::Release); - upload_stop.store(true, Ordering::Release); + upload_stop.cancel(); let _ = child.kill(); break (child.wait(), false); } @@ -2055,7 +2112,7 @@ fn bridge_connection( || upload_failed.load(Ordering::Acquire) || download_done.load(Ordering::Acquire) { - upload_stop.store(true, Ordering::Release); + upload_stop.cancel(); let stopped_at = stopped_at.get_or_insert_with(Instant::now); if stopped_at.elapsed() >= Duration::from_millis(250) { connection_stop.store(true, Ordering::Release); @@ -2065,7 +2122,7 @@ fn bridge_connection( } thread::sleep(BRIDGE_ACCEPT_POLL); }; - upload_stop.store(true, Ordering::Release); + upload_stop.cancel(); if !child_exited { connection_stop.store(true, Ordering::Release); } @@ -2144,21 +2201,29 @@ fn copy_reader_to_local_stream( fn copy_local_stream_to_writer( mut stream: crate::ipc::LocalStream, writer: &mut W, - connection_stop: &AtomicBool, + connection_stop: &BridgeUploadStop, bridge_stop: &AtomicBool, client_closed: &AtomicBool, ) -> io::Result { let mut buffer = [0_u8; 16 * 1024]; let mut total = 0; - while !connection_stop.load(Ordering::Acquire) && !bridge_stop.load(Ordering::Acquire) { + while !connection_stop.is_stopped() && !bridge_stop.load(Ordering::Acquire) { + #[cfg(all(test, unix))] + tests::UPLOAD_READ_ATTEMPTS.with(|attempts| { + if let Some(attempts) = attempts.borrow().as_ref() { + attempts.fetch_add(1, Ordering::Relaxed); + } + }); match crate::ipc::poll_local_stream_read_count(&mut stream, &mut buffer)? { crate::ipc::LocalStreamReadCount::Data(read) => { writer.write_all(&buffer[..read])?; writer.flush()?; total += read as u64; } - crate::ipc::LocalStreamReadCount::Pending => thread::sleep(BRIDGE_IO_POLL), + crate::ipc::LocalStreamReadCount::Pending => { + connection_stop.wake.wait(&stream)?; + } crate::ipc::LocalStreamReadCount::Closed => { client_closed.store(true, Ordering::Release); break; @@ -2246,6 +2311,153 @@ fn sanitize_path_component(input: &str) -> String { mod tests { use super::*; + #[cfg(unix)] + thread_local! { + pub(super) static UPLOAD_READ_ATTEMPTS: std::cell::RefCell>> = const { std::cell::RefCell::new(None) }; + } + + #[cfg(unix)] + fn upload_test_streams(name: &str) -> (crate::ipc::LocalStream, crate::ipc::LocalStream) { + let socket = local_forward_socket_path(name, "upload-test"); + let listener = crate::ipc::bind_private_local_listener(&socket).unwrap(); + let client = crate::ipc::connect_local_stream(&socket).unwrap(); + let server = listener.accept().unwrap(); + server.set_nonblocking(true).unwrap(); + drop(listener); + std::fs::remove_file(socket).unwrap(); + (client, server) + } + + #[cfg(unix)] + #[test] + fn bridge_upload_idle_waits_without_repeated_reads_and_cancels() { + use std::sync::atomic::AtomicUsize; + use std::sync::mpsc; + + let (mut client, stream) = upload_test_streams("idle"); + let attempts = Arc::new(AtomicUsize::new(0)); + let worker_attempts = Arc::clone(&attempts); + let stop = Arc::new(BridgeUploadStop::new().unwrap()); + let worker_stop = Arc::clone(&stop); + let (done_tx, done_rx) = mpsc::channel(); + let worker = thread::spawn(move || { + UPLOAD_READ_ATTEMPTS.with(|slot| *slot.borrow_mut() = Some(worker_attempts)); + let mut output = Vec::new(); + let closed = AtomicBool::new(false); + let result = copy_local_stream_to_writer( + stream, + &mut output, + &worker_stop, + &AtomicBool::new(false), + &closed, + ); + done_tx + .send((result, output, closed.load(Ordering::Acquire))) + .unwrap(); + }); + let deadline = Instant::now() + Duration::from_secs(5); + while attempts.load(Ordering::Relaxed) == 0 { + assert!(Instant::now() < deadline, "upload worker did not start"); + thread::sleep(Duration::from_millis(1)); + } + thread::sleep(Duration::from_millis(100)); + let idle_reads = attempts.load(Ordering::Relaxed); + client.write_all(b"pane input").unwrap(); + let deadline = Instant::now() + Duration::from_secs(5); + while attempts.load(Ordering::Relaxed) < idle_reads + 2 { + assert!( + Instant::now() < deadline, + "input did not wake the upload worker" + ); + thread::sleep(Duration::from_millis(1)); + } + thread::sleep(Duration::from_millis(100)); + let reads_after_input = attempts.load(Ordering::Relaxed); + stop.cancel(); + let (result, output, closed) = done_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + worker.join().unwrap(); + assert_eq!(result.unwrap(), 10); + assert_eq!(output, b"pane input"); + assert!(!closed, "cancellation is not a peer disconnect"); + assert_eq!(idle_reads, 1, "idle forwarding must wait, not retry reads"); + assert_eq!( + reads_after_input, 3, + "forwarding must sleep again after input" + ); + } + + #[cfg(unix)] + #[test] + fn bridge_upload_cancel_before_wait_preserves_download() { + use std::io::Read as _; + + let (mut client, stream) = upload_test_streams("cancel-before-wait"); + let mut download = stream.try_clone().unwrap(); + let stop = BridgeUploadStop::new().unwrap(); + stop.cancel(); + stop.cancel(); + let closed = AtomicBool::new(false); + let count = copy_local_stream_to_writer( + stream, + &mut Vec::new(), + &stop, + &AtomicBool::new(false), + &closed, + ) + .unwrap(); + assert_eq!(count, 0); + assert!(!closed.load(Ordering::Acquire)); + download.write_all(b"final frame").unwrap(); + let mut output = [0; 11]; + client.read_exact(&mut output).unwrap(); + assert_eq!(&output, b"final frame"); + } + + #[cfg(unix)] + #[test] + fn bridge_upload_cancel_between_stop_check_and_wait_is_retained() { + let (_client, stream) = upload_test_streams("cancel-before-poll"); + let stop = BridgeUploadStop::new().unwrap(); + assert!(!stop.is_stopped()); + stop.cancel(); + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + done_tx.send(stop.wake.wait(&stream)).unwrap(); + }); + done_rx + .recv_timeout(Duration::from_secs(2)) + .unwrap() + .unwrap(); + worker.join().unwrap(); + } + + #[cfg(unix)] + #[test] + fn bridge_upload_drains_input_before_peer_eof() { + let (mut client, stream) = upload_test_streams("drain"); + let payload = vec![b'x'; 1024 * 1024]; + let expected = payload.clone(); + let worker = thread::spawn(move || { + let stop = BridgeUploadStop::new().unwrap(); + let mut output = Vec::new(); + let closed = AtomicBool::new(false); + let count = copy_local_stream_to_writer( + stream, + &mut output, + &stop, + &AtomicBool::new(false), + &closed, + ) + .unwrap(); + assert!(closed.load(Ordering::Acquire)); + assert_eq!(count, output.len() as u64); + output + }); + client.write_all(&payload).unwrap(); + drop(client); + assert_eq!(worker.join().unwrap(), expected); + } + #[cfg(unix)] #[test] fn bridge_socket_is_user_only() {