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

1 files changed, 146 insertions, 0 deletions

diff --git a/src/client/transport.rs b/src/client/transport.rs
new file mode 100644
index 00000000..4807c31d
--- /dev/null
+++ b/src/client/transport.rs
@@ -0,0 +1,146 @@
+use super::*;
+
+pub(super) fn start_endpoint_transport(
+    stream: LocalStream,
+    lifetime: impl Send + 'static,
+    event_tx: &tokio::sync::mpsc::Sender<ClientLoopEvent>,
+    endpoint_id: endpoint::ClientEndpointId,
+    generation: u64,
+    max_frame_size: usize,
+) -> Result<endpoint::NativeEndpointTransport, ClientError> {
+    let reader = stream.try_clone().map_err(ClientError::ConnectionFailed)?;
+    let transport = endpoint::NativeEndpointTransport::with_lifetime(stream, lifetime)
+        .map_err(ClientError::ConnectionFailed)?;
+    let stopped = transport.stop_handle();
+    let event_tx = event_tx.clone();
+    std::thread::Builder::new()
+        .name("endpoint-reader".into())
+        .spawn(move || {
+            server_reader_thread(
+                reader,
+                event_tx,
+                &stopped,
+                max_frame_size,
+                endpoint_id,
+                generation,
+            );
+        })
+        .map_err(ClientError::ConnectionFailed)?;
+    Ok(transport)
+}
+
+/// Reads complete frames while retaining partial-read progress across nonblocking polls.
+pub(super) fn server_reader_thread(
+    mut stream: LocalStream,
+    event_tx: tokio::sync::mpsc::Sender<ClientLoopEvent>,
+    should_quit: &Arc<AtomicBool>,
+    max_frame_size: usize,
+    endpoint_id: endpoint::ClientEndpointId,
+    generation: u64,
+) {
+    if stream.set_nonblocking(true).is_err() {
+        let _ = event_tx.blocking_send(ClientLoopEvent::ServerDisconnected {
+            endpoint_id,
+            generation,
+        });
+        return;
+    }
+
+    let mut stream = EndpointReader {
+        stream: &mut stream,
+        stopped: should_quit,
+    };
+    loop {
+        if should_quit.load(Ordering::Acquire) {
+            break;
+        }
+
+        match protocol::read_message(&mut stream, max_frame_size) {
+            Ok(msg) => {
+                if event_tx
+                    .blocking_send(ClientLoopEvent::ServerMessage {
+                        endpoint_id: endpoint_id.clone(),
+                        generation,
+                        message: Box::new(msg),
+                    })
+                    .is_err()
+                {
+                    break;
+                }
+            }
+            Err(protocol::FramingError::UnexpectedEof) => {
+                let _ = event_tx.blocking_send(ClientLoopEvent::ServerDisconnected {
+                    endpoint_id: endpoint_id.clone(),
+                    generation,
+                });
+                break;
+            }
+            Err(protocol::FramingError::Io(err)) if err.kind() == io::ErrorKind::WouldBlock => {
+                std::thread::sleep(Duration::from_millis(1));
+            }
+            Err(err) => {
+                warn!(err = %err, "server read error");
+                let _ = event_tx.blocking_send(ClientLoopEvent::ServerDisconnected {
+                    endpoint_id: endpoint_id.clone(),
+                    generation,
+                });
+                break;
+            }
+        }
+    }
+}
+
+struct EndpointReader<'a> {
+    stream: &'a mut LocalStream,
+    stopped: &'a AtomicBool,
+}
+
+impl io::Read for EndpointReader<'_> {
+    fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
+        loop {
+            if self.stopped.load(Ordering::Acquire) {
+                return Ok(0);
+            }
+            match crate::ipc::poll_local_stream_read_count(self.stream, buffer)? {
+                crate::ipc::LocalStreamReadCount::Data(count) => return Ok(count),
+                crate::ipc::LocalStreamReadCount::Closed => return Ok(0),
+                crate::ipc::LocalStreamReadCount::Pending => {
+                    crate::platform::wait_client_stream_readable(self.stream)?;
+                }
+            }
+        }
+    }
+}
+
+pub(in crate::client) fn write_to_local_server(
+    stream: &mut LocalStream,
+    msg: &ClientMessage,
+) -> io::Result<()> {
+    protocol::write_message(stream, msg).map_err(|error| io::Error::other(error.to_string()))
+}
+
+pub(super) trait ClientMessageSink {
+    fn send_client_message(&mut self, message: &ClientMessage) -> io::Result<()>;
+}
+
+impl ClientMessageSink for LocalStream {
+    fn send_client_message(&mut self, message: &ClientMessage) -> io::Result<()> {
+        write_to_local_server(self, message)
+    }
+}
+
+impl ClientMessageSink for endpoint::EndpointRegistry {
+    fn send_client_message(&mut self, message: &ClientMessage) -> io::Result<()> {
+        // The lifecycle loop consumes failures for every endpoint, including Local. A send
+        // failure must not bypass that transition or tear down unrelated connections.
+        self.send(message);
+        Ok(())
+    }
+}
+
+pub(super) fn write_to_server(
+    stream: &mut impl ClientMessageSink,
+    msg: &ClientMessage,
+) -> io::Result<()> {
+    stream.send_client_message(msg)
+}