diff --git a/ares-shell/src/dumb.rs b/ares-shell/src/dumb.rs index 9bf6b6b..0cd6abe 100644 --- a/ares-shell/src/dumb.rs +++ b/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 { - let (tx, rx) = unbounded::>(); - 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>>, - terminated: Arc>, +/// 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 +/// ": 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 { + 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>) -> 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( + 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 { + 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"); } } diff --git a/ares-shell/src/io.rs b/ares-shell/src/io.rs new file mode 100644 index 0000000..466fc43 --- /dev/null +++ b/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> { + let (tx, rx) = unbounded::>(); + 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()) +} diff --git a/ares-shell/src/main.rs b/ares-shell/src/main.rs index 43c16ed..beaf4ba 100644 --- a/ares-shell/src/main.rs +++ b/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); } } } diff --git a/ares-shell/src/pty.rs b/ares-shell/src/pty.rs index ff40c34..616a9eb 100644 --- a/ares-shell/src/pty.rs +++ b/ares-shell/src/pty.rs @@ -1,179 +1,104 @@ use std::io::Write; -use std::sync::{Arc, Mutex}; use std::thread; -use std::thread::JoinHandle; use std::time::Duration; -use crossbeam_channel::{Sender, select, unbounded}; -use crossterm::event::{Event, KeyCode, KeyEvent, KeyEventKind}; +use crossbeam_channel::{select, tick}; use crossterm::terminal; +use libssh_rs::Channel; use libssh_rs::Error::TryAgain; -use libssh_rs::{Channel, Error}; + +use crate::io::{io_error, spawn_stdin_reader}; + +/// How long the main loop waits for input before checking the remote for +/// output and the local terminal for resizes. Small enough to feel instant, +/// large enough to keep the process idle when nothing is happening. +const POLL_INTERVAL: Duration = Duration::from_millis(10); + +/// Restores the terminal to cooked mode when dropped, even on early return. +struct RawMode; + +impl RawMode { + fn enable() -> std::io::Result { + terminal::enable_raw_mode()?; + Ok(Self) + } +} + +impl Drop for RawMode { + fn drop(&mut self) { + let _ = terminal::disable_raw_mode(); + } +} pub(crate) fn shell(ch: Channel) -> Result { - terminal::enable_raw_mode()?; - let (tx, rx) = unbounded::(); - let events = EventThread::new(tx); - let mut buf = [0; 1024]; + let _raw = RawMode::enable()?; + let stdin_rx = spawn_stdin_reader(); + let ticker = tick(POLL_INTERVAL); + + let mut buf = [0u8; 8192]; + let mut last_size = terminal::size().unwrap_or((80, 24)); + let mut input_open = true; + loop { - if ch.is_eof() { + // Forward everything the remote has sent us. In a PTY, stderr is + // folded into stdout, so reading the stdout stream is enough. + drain(&ch, &mut buf)?; + + if ch.is_eof() || ch.is_closed() { break; } - select! { - recv(rx) -> ev => match ev { - Ok(Event::Key(key)) => { - if key.kind == KeyEventKind::Release { - continue; - } - send_key(&mut ch.stdin(), &key)?; - } - Ok(Event::Resize(width, height)) => { - ch.change_pty_size(width as u32, height as u32)?; - } - Ok(_) => { - } - Err(_) => { - break; - } - }, - default => { - match ch.read_nonblocking(&mut buf, false) { - Err(TryAgain) | Ok(0) => { - if ch.is_closed() { - break; - } - thread::sleep(Duration::from_millis(1)) - } - Ok(size) => { - let mut stdout = std::io::stdout(); - stdout.write_all(&buf[..size])?; - stdout.flush()?; - } - Err(e) => { - return Err(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())); - } - } - } - } - } - drop(events); - Ok(ch.get_exit_status().unwrap_or(-1) as i32) -} -fn send_key(stdin: &mut Stdin, key: &KeyEvent) -> Result<(), Error> { - match key.code { - KeyCode::Backspace => { - stdin.write_all(&[0x08, 0x20, 0x08])?; - } - KeyCode::Enter => { - stdin.write_all(&[0x0d])?; - } - KeyCode::Left => { - stdin.write_all(&[0x1b, 0x5b, 0x44])?; - } - KeyCode::Right => { - stdin.write_all(&[0x1b, 0x5b, 0x43])?; - } - KeyCode::Up => { - stdin.write_all(&[0x1b, 0x5b, 0x41])?; - } - KeyCode::Down => { - stdin.write_all(&[0x1b, 0x5b, 0x42])?; - } - KeyCode::Home => { - stdin.write_all(&[0x1b, 0x5b, 0x48])?; - } - KeyCode::End => { - stdin.write_all(&[0x1b, 0x5b, 0x46])?; - } - KeyCode::PageUp => { - stdin.write_all(&[0x1b, 0x5b, 0x35, 0x7e])?; - } - KeyCode::PageDown => { - stdin.write_all(&[0x1b, 0x5b, 0x36, 0x7e])?; - } - KeyCode::Tab => { - stdin.write_all(&[0x09])?; - } - KeyCode::BackTab => { - stdin.write_all(&[0x1b, 0x5b, 0x5a])?; - } - KeyCode::Delete => { - stdin.write_all(&[0x1b, 0x5b, 0x33, 0x7e])?; - } - KeyCode::Insert => { - stdin.write_all(&[0x1b, 0x5b, 0x32, 0x7e])?; - } - KeyCode::F(n) => { - stdin.write_all(&[0x1b, 0x5b, 0x4f, 0x30 + n])?; - } - KeyCode::Char(c) => { - if key - .modifiers - .contains(crossterm::event::KeyModifiers::CONTROL) - { - stdin.write_all(&[c as u8 - 0x61])?; - } else if key.modifiers.contains(crossterm::event::KeyModifiers::ALT) { - stdin.write_all(&[0x1b, c as u8])?; - } else { - stdin.write_all(&[c as u8])?; + // Mirror local terminal resizes to the remote PTY. + match terminal::size() { + Ok(size) if size != last_size => { + ch.change_pty_size(u32::from(size.0), u32::from(size.1)) + .map_err(io_error)?; + last_size = size; } + _ => {} } - KeyCode::Null => { - stdin.write_all(&[0x0])?; - } - KeyCode::Esc => { - stdin.write_all(&[0x1b])?; - } - _ => {} - } - Ok(()) -} - -struct EventThread { - handle: Mutex>>, - terminated: Arc>, -} -impl EventThread { - fn new(tx: Sender) -> 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 Ok(has_event) = crossterm::event::poll(Duration::from_millis(20)) else { - break; - }; - if !has_event { - continue; + if input_open { + select! { + recv(stdin_rx) -> msg => match msg { + Ok(bytes) => { + let mut stdin = ch.stdin(); + stdin.write_all(&bytes)?; + stdin.flush()?; } - let Ok(event) = crossterm::event::read() else { - break; - }; - if !tx.send(event).is_ok() { - break; + // Local stdin reached EOF: let the remote know we are done. + Err(_) => { + input_open = false; + let _ = ch.send_eof(); } - } - }))), + }, + recv(ticker) -> _ => {} + } + } else { + // No more input to forward; just keep draining remote output. + thread::sleep(POLL_INTERVAL); } } + + Ok(ch.get_exit_status().unwrap_or(-1)) } -impl Drop for EventThread { - fn drop(&mut self) { - terminal::disable_raw_mode().unwrap(); - *self.terminated.lock().unwrap() = true; - if let Some(hnd) = self.handle.lock().unwrap().take() { - hnd.join().unwrap(); +/// Writes all currently-available remote output to stdout without blocking. +fn drain(ch: &Channel, buf: &mut [u8]) -> Result<(), std::io::Error> { + let mut stdout = std::io::stdout().lock(); + let mut wrote = false; + loop { + match ch.read_nonblocking(buf, false) { + Ok(0) | Err(TryAgain) => break, + Ok(size) => { + stdout.write_all(&buf[..size])?; + wrote = true; + } + Err(e) => return Err(io_error(e)), } } + if wrote { + stdout.flush()?; + } + Ok(()) }