Refactor shell I/O handling with unified stdin reader

ares-shell/src/dumb.rs

@@ -1,102 +1,131 @@
-use std::io::{Error, Read, Write, stdin};
-use std::sync::{Arc, Mutex};
+use std::io::{Error, Write};
 use std::thread;
-use std::thread::JoinHandle;
 use std::time::Duration;
 
-use crossbeam_channel::{Sender, select, unbounded};
+use crossbeam_channel::{select, tick};
 use libssh_rs::Channel;
 use libssh_rs::Error::TryAgain;
 
+use crate::io::{io_error, spawn_stdin_reader};
+
+const POLL_INTERVAL: Duration = Duration::from_millis(10);
+
 pub(crate) fn shell(ch: Channel) -> Result<i32, Error> {
-    let (tx, rx) = unbounded::<Vec<u8>>();
-    let events = EventThread::new(tx);
-    let mut buf = [0; 1024];
-    let mut stderr = false;
+    let stdin_rx = spawn_stdin_reader();
+    let ticker = tick(POLL_INTERVAL);
+
+    let mut buf = [0u8; 8192];
+    let mut input_open = true;
+    let mut pending_cr = false;
+
     loop {
-        if ch.is_eof() {
+        drain(&ch, &mut buf, false, &mut std::io::stdout().lock())?;
+        drain(&ch, &mut buf, true, &mut std::io::stderr().lock())?;
+
+        if ch.is_eof() || ch.is_closed() {
             break;
         }
-        select! {
-            recv(rx) -> item => match item {
-                Ok(data) => {
-                    ch.stdin().write_all(&data[..])?;
-                    ch.stdin().flush()?;
-                }
-                Err(_) => {
-                    break;
-                }
-            },
-            default => {
-                match ch.read_nonblocking(&mut buf, stderr) {
-                    Err(TryAgain) | Ok(0) => {
-                        if ch.is_closed() {
-                            break;
-                        }
-                        thread::sleep(Duration::from_millis(1))
-                    }
-                    Ok(size) => {
-                        if stderr {
-                            let mut stderr = std::io::stderr();
-                            stderr.write_all(&buf[..size])?;
-                            stderr.flush()?;
-                        } else {
-                            let mut stdout = std::io::stdout();
-                            stdout.write_all(&buf[..size])?;
-                            stdout.flush()?;
-                        }
+
+        if input_open {
+            select! {
+                recv(stdin_rx) -> msg => match msg {
+                    Ok(bytes) => {
+                        // Without a remote PTY there is no line discipline to
+                        // translate carriage returns, so normalize them here.
+                        let bytes = crlf_to_lf(&bytes, &mut pending_cr);
+                        let mut stdin = ch.stdin();
+                        stdin.write_all(&bytes)?;
+                        stdin.flush()?;
                     }
-                    Err(e) => {
-                        return Err(Error::new(std::io::ErrorKind::Other, e.to_string()));
+                    Err(_) => {
+                        input_open = false;
+                        let _ = ch.send_eof();
                     }
-                }
-                stderr = !stderr;
+                },
+                recv(ticker) -> _ => {}
             }
+        } else {
+            thread::sleep(POLL_INTERVAL);
         }
     }
-    drop(events);
-    Ok(ch.get_exit_status().unwrap_or(-1) as i32)
+
+    Ok(ch.get_exit_status().unwrap_or(-1))
 }
 
-struct EventThread {
-    handle: Mutex<Option<JoinHandle<()>>>,
-    terminated: Arc<Mutex<bool>>,
+/// Converts CR and CRLF line endings to LF.
+///
+/// In dumb mode there is no remote pseudo-terminal, so nothing on the far end
+/// turns a carriage return into a newline. A raw `\r` would otherwise reach the
+/// shell as part of the command (e.g. `ls\r`), which fails with
+/// "<cmd>: not found". `pending_cr` carries the state of a `\r` seen at the end
+/// of the previous chunk so a `\r\n` split across reads still collapses to one
+/// `\n`.
+fn crlf_to_lf(bytes: &[u8], pending_cr: &mut bool) -> Vec<u8> {
+    let mut out = Vec::with_capacity(bytes.len());
+    for &b in bytes {
+        match b {
+            b'\r' => {
+                out.push(b'\n');
+                *pending_cr = true;
+            }
+            // The `\n` of a `\r\n` pair: the newline was already emitted.
+            b'\n' if *pending_cr => *pending_cr = false,
+            _ => {
+                out.push(b);
+                *pending_cr = false;
+            }
+        }
+    }
+    out
 }
 
-impl EventThread {
-    fn new(tx: Sender<Vec<u8>>) -> Self {
-        let terminated = Arc::new(Mutex::new(false));
-        let thread_terminated = Arc::downgrade(&terminated);
-        Self {
-            terminated,
-            handle: Mutex::new(Some(thread::spawn(move || {
-                loop {
-                    if let Some(terminated) = thread_terminated.upgrade() {
-                        if *terminated.lock().unwrap() {
-                            break;
-                        }
-                    } else {
-                        break;
-                    }
-                    let mut buf = [0; 1024];
-                    match stdin().read(&mut buf) {
-                        Ok(size) => {
-                            if !tx.send(buf[..size].to_vec()).is_ok() {
-                                break;
-                            }
-                        }
-                        Err(_) => {
-                            break;
-                        }
-                    }
-                }
-            }))),
+/// Writes all currently-available data from one remote stream without blocking.
+fn drain<W: Write>(
+    ch: &Channel,
+    buf: &mut [u8],
+    is_stderr: bool,
+    out: &mut W,
+) -> Result<(), Error> {
+    let mut wrote = false;
+    loop {
+        match ch.read_nonblocking(buf, is_stderr) {
+            Ok(0) | Err(TryAgain) => break,
+            Ok(size) => {
+                out.write_all(&buf[..size])?;
+                wrote = true;
+            }
+            Err(e) => return Err(io_error(e)),
         }
     }
+    if wrote {
+        out.flush()?;
+    }
+    Ok(())
 }
 
-impl Drop for EventThread {
-    fn drop(&mut self) {
-        *self.terminated.lock().unwrap() = true;
+#[cfg(test)]
+mod tests {
+    use super::crlf_to_lf;
+
+    fn convert(chunks: &[&[u8]]) -> Vec<u8> {
+        let mut pending_cr = false;
+        let mut out = Vec::new();
+        for chunk in chunks {
+            out.extend(crlf_to_lf(chunk, &mut pending_cr));
+        }
+        out
+    }
+
+    #[test]
+    fn collapses_crlf_and_bare_cr() {
+        assert_eq!(convert(&[b"ls\r\n"]), b"ls\n");
+        assert_eq!(convert(&[b"ls\n"]), b"ls\n");
+        assert_eq!(convert(&[b"ls\r"]), b"ls\n");
+        assert_eq!(convert(&[b"a\r\rb"]), b"a\n\nb");
+    }
+
+    #[test]
+    fn handles_crlf_split_across_chunks() {
+        assert_eq!(convert(&[b"ls\r", b"\npwd\r\n"]), b"ls\npwd\n");
     }
 }

ares-shell/src/io.rs

@@ -0,0 +1,34 @@
+use std::io::{Read, stdin};
+use std::thread;
+
+use crossbeam_channel::{Receiver, unbounded};
+
+/// Reads raw bytes from stdin on a background thread and forwards them
+/// verbatim, exactly like `ssh`. Doing the read on its own thread lets the
+/// main loop block on either input or remote output without busy-polling.
+///
+/// In PTY mode the terminal is in raw mode, so control keys, UTF-8 input and
+/// escape sequences all arrive already correctly encoded and are passed
+/// through untouched. The channel closes when stdin reaches EOF.
+pub(crate) fn spawn_stdin_reader() -> Receiver<Vec<u8>> {
+    let (tx, rx) = unbounded::<Vec<u8>>();
+    thread::spawn(move || {
+        let mut stdin = stdin().lock();
+        let mut buf = [0u8; 8192];
+        loop {
+            match stdin.read(&mut buf) {
+                Ok(0) | Err(_) => break,
+                Ok(size) => {
+                    if tx.send(buf[..size].to_vec()).is_err() {
+                        break;
+                    }
+                }
+            }
+        }
+    });
+    rx
+}
+
+pub(crate) fn io_error(e: libssh_rs::Error) -> std::io::Error {
+    std::io::Error::new(std::io::ErrorKind::Other, e.to_string())
+}

ares-shell/src/main.rs

@@ -8,6 +8,7 @@ use crossterm::terminal;
 use crossterm::tty::IsTty;
 
 mod dumb;
+mod io;
 mod pty;
 
 #[derive(Parser, Debug)]
@@ -43,7 +44,8 @@ fn main() {
     let mut has_pty = false;
     if !cli.no_pty && (cli.pty || stdout().is_tty()) {
         let (width, height) = terminal::size().unwrap_or((80, 24));
-        if let Err(e) = ch.request_pty("xterm", width as u32, height as u32) {
+        let term = std::env::var("TERM").unwrap_or_else(|_| String::from("xterm"));
+        if let Err(e) = ch.request_pty(&term, u32::from(width), u32::from(height)) {
             eprintln!("Can't request pty: {:?}", e);
             if cli.pty {
                 exit(255);
@@ -66,6 +68,7 @@ fn main() {
         Ok(code) => exit(code),
         Err(e) => {
             eprintln!("Error: {:?}", e);
+            exit(1);
         }
     }
 }