feat: manage multiple ssh machines from one client (#3670)

1 files changed, 400 insertions, 0 deletions

diff --git a/src/client/endpoint/writer.rs b/src/client/endpoint/writer.rs
new file mode 100644
index 00000000..973ff234
--- /dev/null
+++ b/src/client/endpoint/writer.rs
@@ -0,0 +1,400 @@
+use std::io;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::{mpsc, Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use interprocess::local_socket::traits::Stream as _;
+
+use super::EndpointTransport;
+use crate::ipc::LocalStream;
+use crate::protocol::ClientMessage;
+
+const MAX_QUEUED_MESSAGES: usize = 256;
+const MAX_QUEUED_BYTES: usize = 2 * crate::protocol::MAX_GRAPHICS_FRAME_SIZE;
+const WRITE_TIMEOUT: Duration = Duration::from_secs(5);
+const IO_POLL_INTERVAL: Duration = Duration::from_millis(2);
+
+enum WriterCommand {
+    Frame(Vec<u8>),
+    Flush(mpsc::Sender<()>),
+}
+
+/// The UI only enqueues complete frames. A worker owns partial writes, cancellation, and the
+/// bridge lifetime, so neither socket backpressure nor bridge teardown can block other endpoints.
+pub(crate) struct NativeEndpointTransport {
+    sender: mpsc::SyncSender<WriterCommand>,
+    queued_bytes: Arc<AtomicUsize>,
+    stopped: Arc<AtomicBool>,
+    error: Arc<Mutex<Option<io::Error>>>,
+}
+
+impl NativeEndpointTransport {
+    pub(crate) fn with_lifetime(
+        mut stream: LocalStream,
+        lifetime: impl Send + 'static,
+    ) -> io::Result<Self> {
+        stream.set_nonblocking(true)?;
+        let (sender, receiver) = mpsc::sync_channel::<WriterCommand>(MAX_QUEUED_MESSAGES);
+        let queued_bytes = Arc::new(AtomicUsize::new(0));
+        let stopped = Arc::new(AtomicBool::new(false));
+        let error = Arc::new(Mutex::new(None));
+        let worker_bytes = queued_bytes.clone();
+        let worker_stop = stopped.clone();
+        let worker_error = error.clone();
+        std::thread::Builder::new()
+            .name("endpoint-writer".into())
+            .spawn(move || {
+                let _lifetime = lifetime;
+                while let Ok(command) = receiver.recv() {
+                    if worker_stop.load(Ordering::Acquire) {
+                        break;
+                    }
+                    let frame = match command {
+                        WriterCommand::Frame(frame) => frame,
+                        WriterCommand::Flush(done) => {
+                            let _ = done.send(());
+                            continue;
+                        }
+                    };
+                    let result = write_frame(&mut stream, &frame, &worker_stop);
+                    worker_bytes.fetch_sub(frame.len(), Ordering::AcqRel);
+                    if let Err(error) = result {
+                        if let Ok(mut slot) = worker_error.lock() {
+                            *slot = Some(error);
+                        }
+                        worker_stop.store(true, Ordering::Release);
+                        break;
+                    }
+                }
+            })?;
+        Ok(Self {
+            sender,
+            queued_bytes,
+            stopped,
+            error,
+        })
+    }
+
+    pub(crate) fn stop_handle(&self) -> Arc<AtomicBool> {
+        self.stopped.clone()
+    }
+}
+
+impl EndpointTransport for NativeEndpointTransport {
+    fn send(&mut self, message: &ClientMessage) -> io::Result<()> {
+        if self.stopped.load(Ordering::Acquire) {
+            return Err(io::Error::new(
+                io::ErrorKind::BrokenPipe,
+                "endpoint writer stopped",
+            ));
+        }
+        let mut frame = Vec::new();
+        crate::protocol::write_message(&mut frame, message)
+            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error.to_string()))?;
+        let len = frame.len();
+        if self
+            .queued_bytes
+            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |bytes| {
+                bytes
+                    .checked_add(len)
+                    .filter(|total| *total <= MAX_QUEUED_BYTES)
+            })
+            .is_err()
+        {
+            return Err(queue_full());
+        }
+        self.sender
+            .try_send(WriterCommand::Frame(frame))
+            .map_err(|error| {
+                self.queued_bytes.fetch_sub(len, Ordering::AcqRel);
+                match error {
+                    mpsc::TrySendError::Full(_) => queue_full(),
+                    mpsc::TrySendError::Disconnected(_) => {
+                        io::Error::new(io::ErrorKind::BrokenPipe, "endpoint writer stopped")
+                    }
+                }
+            })
+    }
+
+    fn disconnect(&mut self) {
+        self.stopped.store(true, Ordering::Release);
+    }
+
+    fn flush(&mut self, deadline: Instant) -> io::Result<()> {
+        let (done, completion) = mpsc::channel();
+        self.sender
+            .try_send(WriterCommand::Flush(done))
+            .map_err(|_| queue_full())?;
+        completion
+            .recv_timeout(deadline.saturating_duration_since(Instant::now()))
+            .map_err(|error| match error {
+                mpsc::RecvTimeoutError::Timeout => {
+                    io::Error::new(io::ErrorKind::TimedOut, "endpoint flush timed out")
+                }
+                mpsc::RecvTimeoutError::Disconnected => {
+                    io::Error::new(io::ErrorKind::BrokenPipe, "endpoint writer stopped")
+                }
+            })
+    }
+
+    fn take_error(&mut self) -> Option<io::Error> {
+        self.error.lock().ok()?.take()
+    }
+}
+
+impl Drop for NativeEndpointTransport {
+    fn drop(&mut self) {
+        self.stopped.store(true, Ordering::Release);
+    }
+}
+
+fn queue_full() -> io::Error {
+    // A message may already be partially written. Retrying or dropping only this message would
+    // lose input ordering; revoke the connection and recover through the normal lifecycle.
+    io::Error::new(
+        io::ErrorKind::ConnectionAborted,
+        "endpoint output queue is full",
+    )
+}
+
+fn write_frame(
+    writer: &mut impl io::Write,
+    mut frame: &[u8],
+    stopped: &AtomicBool,
+) -> io::Result<()> {
+    let deadline = Instant::now() + WRITE_TIMEOUT;
+    while !frame.is_empty() && !stopped.load(Ordering::Acquire) {
+        match writer.write(frame) {
+            Ok(0) => {}
+            Ok(written) => {
+                frame = &frame[written..];
+                continue;
+            }
+            Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
+            Err(error) if error.kind() == io::ErrorKind::WouldBlock => {}
+            Err(error) => return Err(error),
+        }
+        if Instant::now() >= deadline {
+            return Err(io::Error::new(
+                io::ErrorKind::TimedOut,
+                "endpoint write timed out",
+            ));
+        }
+        std::thread::sleep(IO_POLL_INTERVAL);
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn streams() -> (LocalStream, LocalStream, std::path::PathBuf) {
+        use interprocess::local_socket::traits::Listener as _;
+        let path = std::env::temp_dir().join(format!(
+            "herdr-writer-{}-{}.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 accepting = std::thread::spawn(move || listener.accept().unwrap());
+        let client = crate::ipc::connect_local_stream(&path).unwrap();
+        (client, accepting.join().unwrap(), path)
+    }
+
+    #[test]
+    fn native_endpoint_writer_delivers_ordered_protocol_frames() {
+        let (stream, mut peer, path) = streams();
+        let mut transport = NativeEndpointTransport::with_lifetime(stream, ()).unwrap();
+        let (done, received) = mpsc::channel();
+        let reader = std::thread::spawn(move || {
+            let first: ClientMessage =
+                crate::protocol::read_message(&mut peer, crate::protocol::MAX_FRAME_SIZE).unwrap();
+            let second: ClientMessage =
+                crate::protocol::read_message(&mut peer, crate::protocol::MAX_FRAME_SIZE).unwrap();
+            done.send((first, second)).unwrap();
+        });
+        transport
+            .send(&ClientMessage::ClientShellFocus { focused: true })
+            .unwrap();
+        transport
+            .send(&ClientMessage::ClientShellFocus { focused: false })
+            .unwrap();
+        let messages = received.recv_timeout(Duration::from_secs(3)).unwrap();
+        assert_eq!(
+            messages,
+            (
+                ClientMessage::ClientShellFocus { focused: true },
+                ClientMessage::ClientShellFocus { focused: false }
+            )
+        );
+        reader.join().unwrap();
+        drop(transport);
+        let _ = std::fs::remove_file(path);
+    }
+
+    #[test]
+    fn registry_exit_flushes_queued_input_and_a_complete_detach() {
+        let (stream, mut peer, path) = streams();
+        let transport = NativeEndpointTransport::with_lifetime(stream, ()).unwrap();
+        let (done, received) = mpsc::channel();
+        let reader = std::thread::spawn(move || {
+            let result = (|| {
+                let first: ClientMessage =
+                    crate::protocol::read_message(&mut peer, crate::protocol::MAX_FRAME_SIZE)?;
+                let second: ClientMessage =
+                    crate::protocol::read_message(&mut peer, crate::protocol::MAX_FRAME_SIZE)?;
+                Ok::<_, crate::protocol::FramingError>((first, second))
+            })();
+            done.send(result).unwrap();
+        });
+        let mut registry = super::super::EndpointRegistry::new(
+            transport,
+            1,
+            super::super::EndpointNegotiation::default(),
+        );
+        let input = ClientMessage::Input {
+            data: b"queued input".to_vec(),
+        };
+        assert_eq!(
+            registry.send(&input),
+            super::super::EndpointSendOutcome::Sent
+        );
+        drop(registry);
+        let (first, second) = received
+            .recv_timeout(Duration::from_secs(3))
+            .unwrap()
+            .expect("clean exit must flush complete frames before closing");
+        assert_eq!(first, input);
+        assert_eq!(second, ClientMessage::Detach);
+        reader.join().unwrap();
+        let _ = std::fs::remove_file(path);
+    }
+
+    #[test]
+    fn native_endpoint_flush_drains_large_frames_before_detach() {
+        let (stream, mut peer, path) = streams();
+        let mut transport = NativeEndpointTransport::with_lifetime(stream, ()).unwrap();
+        let (done, received) = mpsc::channel();
+        let reader = std::thread::spawn(move || {
+            let first: ClientMessage =
+                crate::protocol::read_message(&mut peer, crate::protocol::MAX_FRAME_SIZE).unwrap();
+            let second: ClientMessage =
+                crate::protocol::read_message(&mut peer, crate::protocol::MAX_FRAME_SIZE).unwrap();
+            done.send((first, second)).unwrap();
+        });
+        let input = ClientMessage::Input {
+            data: vec![b'x'; 1024 * 1024],
+        };
+        transport.send(&input).unwrap();
+        transport.send(&ClientMessage::Detach).unwrap();
+        // Large-frame correctness must not depend on the registry's short exit grace period.
+        transport
+            .flush(Instant::now() + Duration::from_secs(10))
+            .unwrap();
+        drop(transport);
+        let (first, second) = received.recv_timeout(Duration::from_secs(10)).unwrap();
+        assert_eq!(first, input);
+        assert_eq!(second, ClientMessage::Detach);
+        reader.join().unwrap();
+        let _ = std::fs::remove_file(path);
+    }
+
+    #[test]
+    fn native_endpoint_teardown_does_not_wait_for_a_stalled_peer() {
+        struct Lifetime(mpsc::Sender<()>);
+        impl Drop for Lifetime {
+            fn drop(&mut self) {
+                let _ = self.0.send(());
+            }
+        }
+        let (stream, peer, path) = streams();
+        let (done, dropped) = mpsc::channel();
+        let mut transport = NativeEndpointTransport::with_lifetime(stream, Lifetime(done)).unwrap();
+        transport
+            .send(&ClientMessage::Input {
+                data: vec![b'x'; 2 * 1024 * 1024],
+            })
+            .unwrap();
+        drop(transport);
+        dropped
+            .recv_timeout(Duration::from_secs(3))
+            .expect("worker lifetime is released without peer reads");
+        drop(peer);
+        let _ = std::fs::remove_file(path);
+    }
+
+    #[test]
+    fn partial_writes_preserve_frame_bytes() {
+        struct PartialWriter(Vec<u8>);
+        impl io::Write for PartialWriter {
+            fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
+                let count = bytes.len().min(3);
+                self.0.extend_from_slice(&bytes[..count]);
+                Ok(count)
+            }
+            fn flush(&mut self) -> io::Result<()> {
+                Ok(())
+            }
+        }
+        let mut writer = PartialWriter(Vec::new());
+        write_frame(&mut writer, b"first frame", &AtomicBool::new(false)).unwrap();
+        write_frame(&mut writer, b"second frame", &AtomicBool::new(false)).unwrap();
+        assert_eq!(writer.0, b"first framesecond frame");
+    }
+
+    #[test]
+    fn a_stalled_write_is_cancellable_without_peer_progress() {
+        struct StalledWriter(mpsc::Sender<()>);
+        impl io::Write for StalledWriter {
+            fn write(&mut self, _: &[u8]) -> io::Result<usize> {
+                let _ = self.0.send(());
+                Err(io::ErrorKind::WouldBlock.into())
+            }
+            fn flush(&mut self) -> io::Result<()> {
+                Ok(())
+            }
+        }
+        let stop = Arc::new(AtomicBool::new(false));
+        let worker_stop = stop.clone();
+        let (attempted, attempts) = mpsc::channel();
+        let (done, completion) = mpsc::channel();
+        let worker = std::thread::spawn(move || {
+            let result = write_frame(&mut StalledWriter(attempted), b"input", &worker_stop);
+            done.send(result).unwrap();
+        });
+        attempts.recv_timeout(Duration::from_secs(2)).unwrap();
+        stop.store(true, Ordering::Release);
+        completion
+            .recv_timeout(Duration::from_secs(2))
+            .unwrap()
+            .unwrap();
+        worker.join().unwrap();
+    }
+
+    #[test]
+    fn a_full_queue_is_a_connection_failure_not_silent_input_loss() {
+        let (sender, _receiver) = mpsc::sync_channel(1);
+        let mut transport = NativeEndpointTransport {
+            sender,
+            queued_bytes: Arc::new(AtomicUsize::new(0)),
+            stopped: Arc::new(AtomicBool::new(false)),
+            error: Arc::new(Mutex::new(None)),
+        };
+        transport.send(&ClientMessage::Detach).unwrap();
+        assert_eq!(
+            transport.send(&ClientMessage::Detach).unwrap_err().kind(),
+            io::ErrorKind::ConnectionAborted
+        );
+        transport
+            .queued_bytes
+            .store(MAX_QUEUED_BYTES, Ordering::Release);
+        assert_eq!(
+            transport.send(&ClientMessage::Detach).unwrap_err().kind(),
+            io::ErrorKind::ConnectionAborted
+        );
+    }
+}