Small tweaks and fixes

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<Read: AsyncRead + Unpin>(
-    reader: &mut Read,
+async fn client_sync<S: AsyncRead + Unpin, T: AsyncRead + Unpin>(
+    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<ui::UIEvent>) {
                 .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) => {

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<Vec<PortDesc>>,
     lines: VecDeque<String>,
     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;