diff --git a/src/lib.rs b/src/lib.rs index 2cbeb99..0b274e2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -76,16 +76,20 @@ pub async fn run_server() { /// Wait for the server to be ready; we know the server is there and /// listening when we see the special sync marker, which is 8 NUL bytes in a /// row. -/// -/// TODO: We should be pumping stderr too. -async fn client_sync( - reader: &mut Read, +async fn client_sync( + reader: &mut S, + client_stderr: &mut T, ) -> Result<(), tokio::io::Error> { info!("Waiting for synchronization marker..."); + let mut stderr = tokio::io::stderr(); let mut stdout = tokio::io::stdout(); let mut seen = 0; tokio::select! { + result = tokio::io::copy(client_stderr, &mut stderr) => match result { + Ok(_) => Ok(()), // ? + Err(e) => Err(e), + }, result = async { while seen < 8 { let byte = reader.read_u8().await?; @@ -193,30 +197,8 @@ async fn client_handle_connection( info!("Connection established on port {}", port); - let (client_read_half, client_write_half) = socket.into_split(); - let (server_read_half, server_write_half) = dest_socket.into_split(); - let client_to_server = tokio::spawn(async move { - let mut client_read_half = client_read_half; - let mut server_write_half = server_write_half; - tokio::io::copy(&mut client_read_half, &mut server_write_half).await - }); - let server_to_client = tokio::spawn(async move { - let mut server_read_half = server_read_half; - let mut client_write_half = client_write_half; - tokio::io::copy(&mut server_read_half, &mut client_write_half).await - }); - - let client_err = client_to_server.await; - debug!("Done client -> server"); - let svr_err = server_to_client.await; - debug!("Done server -> client"); - - if let Ok(Err(e)) = client_err { - return Err(e.into()); - } else if let Ok(Err(e)) = svr_err { - return Err(e.into()); - } - + let mut socket = socket; + tokio::io::copy_bidirectional(&mut socket, &mut dest_socket).await?; Ok(()) } @@ -393,7 +375,7 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender) { .expect("child did not have a handle to stdout"), ); - if let Err(e) = client_sync(&mut reader).await { + if let Err(e) = client_sync(&mut reader, &mut stderr).await { error!("Error synchronizing: {:?}", e); match child.wait().await { Ok(status) => { diff --git a/src/ui.rs b/src/ui.rs index e374cb6..bf86cbb 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, RestorePosition, SavePosition}, + cursor::MoveTo, event::{Event, EventStream, KeyCode, KeyEvent, KeyModifiers}, execute, queue, style::{Color, PrintStyledContent, Stylize}, @@ -68,6 +68,7 @@ pub struct UI { ports: Option>, lines: VecDeque, alternate_screen: bool, + raw_mode: bool, } impl UI { @@ -82,11 +83,14 @@ impl UI { ports: None, lines: VecDeque::with_capacity(1024), alternate_screen: false, + raw_mode: false, } } 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 } @@ -109,7 +113,7 @@ impl UI { } fn render_disconnected(&mut self) -> Result<()> { - self.leave_alternate_screen()?; + self.disable_raw_mode()?; let mut stdout = stdout(); let (columns, _) = size()?; @@ -117,20 +121,19 @@ impl UI { execute!( stdout, - SavePosition, + Clear(ClearType::All), MoveTo(0, 0), PrintStyledContent( - format!("{:^columns$}", "Not Connected") + 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(); let (columns, rows) = size()?; @@ -254,6 +257,22 @@ impl UI { } } + fn enable_raw_mode(&mut self) -> Result<()> { + if !self.raw_mode { + enable_raw_mode()?; + self.raw_mode = true; + } + Ok(()) + } + + fn disable_raw_mode(&mut self) -> Result<()> { + if self.raw_mode { + disable_raw_mode()?; + self.raw_mode = false; + } + Ok(()) + } + fn enter_alternate_screen(&mut self) -> Result<()> { if !self.alternate_screen { enable_raw_mode()?; @@ -333,6 +352,7 @@ impl UI { match event { Some(UIEvent::Disconnected) => { self.socks_port = 0; + self.listeners = HashMap::new(); // Bye. } Some(UIEvent::Connected(sp)) => { self.socks_port = sp;