diff --git a/src/lib.rs b/src/lib.rs index 0b274e2..e0b30c7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ use anyhow::{bail, Result}; +use bytes::BytesMut; use log::LevelFilter; use log::{debug, error, info, warn}; use message::{Message, MessageReader, MessageWriter}; @@ -84,12 +85,17 @@ async fn client_sync( let mut stderr = tokio::io::stderr(); let mut stdout = tokio::io::stdout(); + let mut buf = BytesMut::with_capacity(1024); + let mut seen = 0; - tokio::select! { - result = tokio::io::copy(client_stderr, &mut stderr) => match result { - Ok(_) => Ok(()), // ? - Err(e) => Err(e), - }, + let result = tokio::select! { + result = async { + loop { + buf.clear(); + client_stderr.read_buf(&mut buf).await?; + stderr.write_all(&buf[..]).await?; + } + } => result, result = async { while seen < 8 { let byte = reader.read_u8().await?; @@ -103,7 +109,17 @@ async fn client_sync( Ok::<_, tokio::io::Error>(()) } => result, + }; + + if let Err(_) = result { + // Something went wrong, let's just make sure we flush the client's + // stderr before we return. + _ = stderr.write_all(&buf[..]).await; + _ = tokio::io::copy(client_stderr, &mut stderr).await; + _ = stderr.flush().await; } + + result } /// Handle an incoming client connection, by forwarding it to the SOCKS5 @@ -356,12 +372,10 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender) { let (mut child, socks_port) = spawn_ssh(remote).await.expect("failed to spawn"); - let mut stderr = BufReader::new( - child - .stderr - .take() - .expect("child did not have a handle to stderr"), - ); + let mut stderr = child + .stderr + .take() + .expect("child did not have a handle to stderr"); let writer = child .stdin @@ -395,6 +409,7 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender) { continue; } + let mut stderr = BufReader::new(stderr); let mut writer = MessageWriter::new(BufWriter::new(writer)); let mut reader = MessageReader::new(reader); diff --git a/src/ui.rs b/src/ui.rs index bf86cbb..f56a1e4 100644 --- a/src/ui.rs +++ b/src/ui.rs @@ -2,7 +2,7 @@ use crate::client_listen; use crate::message::PortDesc; use anyhow::Result; use crossterm::{ - cursor::MoveTo, + cursor::{MoveTo, RestorePosition, SavePosition}, event::{Event, EventStream, KeyCode, KeyEvent, KeyModifiers}, execute, queue, style::{Color, PrintStyledContent, Stylize}, @@ -88,14 +88,6 @@ impl UI { } pub async fn run(&mut self) -> Result<()> { - self.enter_alternate_screen()?; - let result = self.run_core().await; - _ = self.disable_raw_mode(); - _ = self.leave_alternate_screen(); - result - } - - async fn run_core(&mut self) -> Result<()> { let mut console_events = EventStream::new(); self.running = true; @@ -113,6 +105,7 @@ impl UI { } fn render_disconnected(&mut self) -> Result<()> { + self.enter_alternate_screen()?; self.disable_raw_mode()?; let mut stdout = stdout(); @@ -121,18 +114,20 @@ impl UI { execute!( stdout, - Clear(ClearType::All), + SavePosition, MoveTo(0, 0), PrintStyledContent( format!("{:^columns$}\r\n", "Not Connected") .with(Color::Black) .on(Color::Red) ), + RestorePosition, )?; Ok(()) } fn render_connected(&mut self) -> Result<()> { + self.enter_alternate_screen()?; self.enable_raw_mode()?; let mut stdout = stdout(); @@ -404,3 +399,10 @@ impl UI { } } } + +impl Drop for UI { + fn drop(&mut self) { + _ = self.disable_raw_mode(); + _ = self.leave_alternate_screen(); + } +}