fix: prevent stdout/stderr race (#19)

Author: imos

src/proxy.rs

@@ -138,11 +138,13 @@ pub async fn run(args: ProxyArgs) -> Result<()> {
         local_port, pod_name, remote_port
     );
 
-    let stream = TcpStream::connect(("127.0.0.1", local_port))
+    let pump_result = match TcpStream::connect(("127.0.0.1", local_port))
         .await
-        .context("failed to connect to forwarded sshd port")?;
-
-    let pump_result = proxy_io::pump(stream).await;
+        .context("failed to connect to forwarded sshd port")
+    {
+        Ok(stream) => proxy_io::pump(stream).await,
+        Err(err) => Err(err),
+    };
     let stop_result = forward.stop().await;
 
     pump_result?;

src/proxy_io.rs

@@ -1,31 +1,94 @@
 use anyhow::Result;
-use tokio::io::{self, AsyncWriteExt};
+use log::debug;
+use tokio::io::{self, AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio::net::TcpStream;
 
 pub async fn pump(stream: TcpStream) -> Result<()> {
-    let (mut reader, mut writer) = stream.into_split();
-    let mut stdin = io::stdin();
-    let mut stdout = io::stdout();
+    let (reader, writer) = stream.into_split();
+    pump_io(io::stdin(), io::stdout(), reader, writer).await
+}
 
-    let to_remote = tokio::spawn(async move {
-        let copied = tokio::io::copy(&mut stdin, &mut writer).await?;
-        writer.shutdown().await?;
+async fn pump_io<I, O, R, W>(
+    mut input: I,
+    mut output: O,
+    mut remote_reader: R,
+    mut remote_writer: W,
+) -> Result<()>
+where
+    I: AsyncRead + Unpin,
+    O: AsyncWrite + Unpin,
+    R: AsyncRead + Unpin,
+    W: AsyncWrite + Unpin,
+{
+    let to_remote = async {
+        let copied = tokio::io::copy(&mut input, &mut remote_writer).await?;
+        remote_writer.shutdown().await?;
         Ok::<_, anyhow::Error>(copied)
-    });
+    };
 
-    let from_remote = tokio::spawn(async move {
-        let copied = tokio::io::copy(&mut reader, &mut stdout).await?;
-        stdout.flush().await?;
+    let from_remote = async {
+        let copied = tokio::io::copy(&mut remote_reader, &mut output).await?;
+        output.flush().await?;
         Ok::<_, anyhow::Error>(copied)
-    });
-
-    let (a, b) = tokio::join!(to_remote, from_remote);
-    let to_bytes = a??;
-    let from_bytes = b??;
-    // Debug logging kept compact
-    eprintln!(
-        "[sshpod][proxy_io] bytes_to_remote={} bytes_from_remote={}",
-        to_bytes, from_bytes
-    );
+    };
+
+    tokio::pin!(to_remote);
+    tokio::pin!(from_remote);
+
+    tokio::select! {
+        result = &mut to_remote => {
+            let bytes = result?;
+            debug!("[sshpod][proxy_io] bytes_to_remote={}", bytes);
+        }
+        result = &mut from_remote => {
+            let bytes = result?;
+            debug!("[sshpod][proxy_io] bytes_from_remote={}", bytes);
+        }
+    }
+
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::pump_io;
+    use tokio::io::{self, AsyncReadExt};
+    use tokio::net::{TcpListener, TcpStream};
+    use tokio::time::{timeout, Duration};
+
+    #[tokio::test]
+    async fn remote_eof_finishes_even_when_local_input_stays_open() {
+        let (_local_writer, local_reader) = io::duplex(64);
+
+        let result = timeout(
+            Duration::from_millis(100),
+            pump_io(local_reader, io::sink(), io::empty(), io::sink()),
+        )
+        .await;
+
+        assert!(result.is_ok(), "pump_io should not wait for local input");
+        result.unwrap().unwrap();
+    }
+
+    #[tokio::test]
+    async fn local_eof_shuts_down_tcp_write_half() {
+        let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
+        let addr = listener.local_addr().unwrap();
+
+        let client = tokio::spawn(TcpStream::connect(addr));
+        let (mut server, _) = listener.accept().await.unwrap();
+        let client = client.await.unwrap().unwrap();
+        let (remote_reader, remote_writer) = client.into_split();
+
+        pump_io(io::empty(), io::sink(), remote_reader, remote_writer)
+            .await
+            .unwrap();
+
+        let mut buf = [0; 1];
+        let read = timeout(Duration::from_millis(100), server.read(&mut buf))
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(read, 0);
+    }
+}