From 6cc17e7fca0cb46dc1cd24a6b2684e01c54afc1d Mon Sep 17 00:00:00 2001 From: John Doty Date: Sun, 16 Oct 2022 08:55:30 -0700 Subject: [PATCH] Try to make the UI better when unconnected --- rustfmt.toml | 2 + src/lib.rs | 66 ++++++-- src/ui.rs | 440 ++++++++++++++++++++++++++++++--------------------- 3 files changed, 321 insertions(+), 187 deletions(-) create mode 100644 rustfmt.toml diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..3214201 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,2 @@ +max_width = 80 +struct_lit_width = 40 diff --git a/src/lib.rs b/src/lib.rs index f954731..bf788db 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,8 +4,8 @@ use log::{debug, error, info, warn}; use message::{Message, MessageReader, MessageWriter}; use std::net::{Ipv4Addr, SocketAddrV4}; use tokio::io::{ - AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, - BufWriter, + AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, + AsyncWriteExt, BufReader, BufWriter, }; use tokio::net::{TcpListener, TcpStream}; use tokio::process; @@ -119,7 +119,9 @@ async fn server_main( } } -async fn client_sync(reader: &mut Read) -> Result<(), tokio::io::Error> { +async fn client_sync( + reader: &mut Read, +) -> Result<(), tokio::io::Error> { info!("Waiting for synchronization marker..."); // Run these two loops in parallel; the copy of stdin should stop when @@ -149,7 +151,11 @@ async fn client_sync(reader: &mut Read) -> Result<(), t /// This contains a very simplified implementation of a SOCKS5 connector, /// enough to work with the SSH I have. I would have liked it to be SOCKS4, /// which is a much simpler protocol, but somehow it didn't work. -async fn client_handle_connection(socks_port: u16, port: u16, socket: TcpStream) -> Result<()> { +async fn client_handle_connection( + socks_port: u16, + port: u16, + socket: TcpStream, +) -> Result<()> { debug!("Handling connection!"); let dest_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, socks_port); @@ -257,14 +263,18 @@ async fn client_handle_connection(socks_port: u16, port: u16, socket: TcpStream) async fn client_listen(port: u16, socks_port: u16) -> Result<()> { loop { - let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)).await?; + let listener = + TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)) + .await?; loop { // The second item contains the IP and port of the new // connection, but we don't care. let (socket, _) = listener.accept().await?; tokio::spawn(async move { - if let Err(e) = client_handle_connection(socks_port, port, socket).await { + if let Err(e) = + client_handle_connection(socks_port, port, socket).await + { error!("Error handling connection: {:?}", e); } else { debug!("Done???"); @@ -402,7 +412,9 @@ pub async fn run_server() { } } -async fn spawn_ssh(server: &str) -> Result<(tokio::process::Child, u16), std::io::Error> { +async fn spawn_ssh( + server: &str, +) -> Result<(tokio::process::Child, u16), std::io::Error> { let socks_port = { let listener = TcpListener::bind("127.0.0.1:0").await?; listener.local_addr()?.port() @@ -423,11 +435,30 @@ async fn spawn_ssh(server: &str) -> Result<(tokio::process::Child, u16), std::io Ok((child, socks_port)) } +#[cfg(target_family = "windows")] +fn is_sigint(status: std::process::ExitStatus) -> bool { + match status.code() { + Some(255) => true, + _ => false, + } +} + +#[cfg(target_family = "unix")] +fn is_sigint(status: std::process::ExitStatus) -> bool { + use std::os::unix::process::ExitStatusExt; + match status.signal() { + Some(2) => true, + Some(_) => false, + None => false, + } +} + async fn client_connect_loop(remote: &str, events: mpsc::Sender) { loop { _ = events.send(ui::UIEvent::Disconnected).await; - let (mut child, socks_port) = spawn_ssh(remote).await.expect("failed to spawn"); + let (mut child, socks_port) = + spawn_ssh(remote).await.expect("failed to spawn"); let mut stderr = BufReader::new( child @@ -450,6 +481,20 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender) { if let Err(e) = client_sync(&mut reader).await { error!("Error synchronizing: {:?}", e); + match child.wait().await { + Ok(status) => { + if is_sigint(status) { + return; + } else { + match status.code() { + Some(127) => eprintln!("Cannot find `fwd` remotely, make sure it is installed"), + _ => (), + }; + } + } + Err(_) => (), + }; + tokio::time::sleep(tokio::time::Duration::from_secs(3)).await; continue; } @@ -462,7 +507,10 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender) { client_pipe_stderr(&mut stderr, sec).await; }); - if let Err(e) = client_main(socks_port, &mut reader, &mut writer, events.clone()).await { + if let Err(e) = + client_main(socks_port, &mut reader, &mut writer, events.clone()) + .await + { error!("Server disconnected with error: {:?}", e); } else { warn!("Disconnected from server, reconnecting..."); diff --git a/src/ui.rs b/src/ui.rs index cbabf1b..aca5917 100644 --- a/src/ui.rs +++ b/src/ui.rs @@ -2,13 +2,14 @@ use crate::client_listen; use crate::message::PortDesc; use anyhow::Result; use crossterm::{ - cursor::MoveTo, - event::{Event, EventStream, KeyCode, KeyEvent}, + cursor::{MoveTo, RestorePosition, SavePosition}, + event::{Event, EventStream, KeyCode, KeyEvent, KeyModifiers}, execute, queue, style::{Color, PrintStyledContent, Stylize}, terminal::{ - disable_raw_mode, enable_raw_mode, size, Clear, ClearType, DisableLineWrap, EnableLineWrap, - EnterAlternateScreen, LeaveAlternateScreen, + disable_raw_mode, enable_raw_mode, size, Clear, ClearType, + DisableLineWrap, EnableLineWrap, EnterAlternateScreen, + LeaveAlternateScreen, }, }; use log::{error, info, Level, Metadata, Record}; @@ -56,10 +57,17 @@ impl log::Log for Logger { fn flush(&self) {} } +#[derive(Debug)] pub struct UI { events: mpsc::Receiver, listeners: HashMap>, socks_port: u16, + running: bool, + show_logs: bool, + selection: usize, + ports: Option>, + lines: VecDeque, + alternate_screen: bool, } impl UI { @@ -68,162 +76,107 @@ impl UI { events, listeners: HashMap::new(), socks_port: 0, + running: true, + show_logs: false, + selection: 0, + ports: None, + lines: VecDeque::with_capacity(1024), + alternate_screen: false, } } pub async fn run(&mut self) -> Result<()> { - enable_raw_mode()?; let result = self.run_core().await; - execute!(stdout(), EnableLineWrap, LeaveAlternateScreen)?; - disable_raw_mode()?; + _ = self.leave_alternate_screen(); result } async fn run_core(&mut self) -> Result<()> { - let mut stdout = stdout(); - - execute!(stdout, EnterAlternateScreen, DisableLineWrap)?; let mut console_events = EventStream::new(); - let mut connected = false; - let mut selection = 0; - let mut show_logs = false; - let mut lines: VecDeque = VecDeque::with_capacity(1024); - let mut ports: Option> = None; - loop { - tokio::select! { - ev = console_events.next() => { - match ev { - Some(Ok(Event::Key(ev))) => { - match ev { - KeyEvent {code:KeyCode::Esc, ..} - | KeyEvent {code:KeyCode::Char('q'), ..} => { break; }, - KeyEvent {code:KeyCode::Char('l'), ..} => { - show_logs = !show_logs; - } - KeyEvent {code:KeyCode::Char('e'), ..} => { - if let Some(ports) = &ports { - if selection < ports.len() { - let p = ports[selection].port; - self.enable_disable_port(p); - } - } - } - KeyEvent { code:KeyCode::Up, ..} - | KeyEvent { code:KeyCode::Char('j'), ..} => { - if selection > 0 { - selection -= 1; - } - } - KeyEvent { code:KeyCode::Down, ..} - | KeyEvent { code:KeyCode::Char('k'), ..} => { - if let Some(p) = &ports { - if selection != p.len() - 1 { - selection += 1; - } - } - } - KeyEvent { code:KeyCode::Enter, ..} => { - if let Some(p) = &ports { - if selection < p.len() { - _ = open::that(format!("http://127.0.0.1:{}/", p[selection].port)); - } - } - } - _ => () - } - }, - Some(Ok(_)) => (), // Don't care about this event... - Some(Err(_)) => (), // Hmmmmmm.....? - None => (), // ....no events? what? - } - } - ev = self.events.recv() => { - match ev { - Some(UIEvent::Disconnected) => { - self.socks_port = 0; - connected = false; - } - Some(UIEvent::Connected(sp)) => { - self.socks_port = sp; - info!("Socks port {socks_port}", socks_port=self.socks_port); - connected = true; - } - Some(UIEvent::Ports(mut p)) => { - p.sort_by(|a, b| a.port.partial_cmp(&b.port).unwrap()); - if selection >= p.len() { - selection = p.len()-1; - } + self.running = true; + while self.running { + self.handle_events(&mut console_events).await; - let mut new_listeners = HashMap::new(); - for port in &p { - let port = port.port; - if let Some(l) = self.listeners.remove(&port) { - if !l.is_closed() { - // `l` here is, of course, the channel that we - // use to tell the listener task to stop (see the - // spawn call below). If it isn't closed then - // that means a spawn task is still running so we - // should just let it keep running and re-use the - // existing listener. - new_listeners.insert(port, l); - } - } - } - - // This has the side effect of closing any - // listener that didn't get copied over to the - // new listeners table. - self.listeners = new_listeners; - ports = Some(p); - } - Some(UIEvent::ServerLine(line)) => { - while lines.len() >= 1024 { - lines.pop_front(); - } - lines.push_back(format!("[SERVER] {line}")); - } - Some(UIEvent::LogLine(_level, line)) => { - while lines.len() >= 1024 { - lines.pop_front(); - } - lines.push_back(format!("[CLIENT] {line}")); - } - None => break, - } - } + if self.connected() { + self.render_connected()?; + } else { + self.render_disconnected()?; } + } - let (columns, rows) = size()?; - let columns: usize = columns.into(); + Ok(()) + } - queue!(stdout, Clear(ClearType::All), MoveTo(0, 0))?; - if connected { - // List of open ports - // How wide are all the things? - let padding = 1; - let enabled_width = 1; // Just 1 character - let port_width = 5; // 5 characters for 16-bit number + fn render_disconnected(&mut self) -> Result<()> { + self.leave_alternate_screen()?; + let mut stdout = stdout(); - let description_width = - columns - (padding + padding + enabled_width + padding + port_width + padding); + let (columns, _) = size()?; + let columns: usize = columns.into(); + execute!( + stdout, + SavePosition, + MoveTo(0, 0), + PrintStyledContent( + format!("{:^columns$}", "Not Connected") + .with(Color::Black) + .on(Color::Red) + ), + RestorePosition, + )?; + Ok(()) + } + + fn render_connected(&mut self) -> Result<()> { + self.enter_alternate_screen()?; + let mut stdout = stdout(); + + let (columns, rows) = size()?; + let columns: usize = columns.into(); + + queue!(stdout, Clear(ClearType::All), MoveTo(0, 0))?; + + // List of open ports + // How wide are all the things? + let padding = 1; + let enabled_width = 1; // Just 1 character + let port_width = 5; // 5 characters for 16-bit number + + let description_width = columns + - (padding + + padding + + enabled_width + + padding + + port_width + + padding); + + print!( + "{}", + format!( + " {enabled:>enabled_width$} {port:>port_width$} {description:enabled_width$} {port:>port_width$} {description:enabled_width$} {:port_width$} {:description_width$}\r\n", - if index == selection { "\u{2B46}" } else { " " }, + if index == self.selection { + "\u{2B46}" + } else { + " " + }, if self.listeners.contains_key(&port.port) { "+" } else { @@ -232,49 +185,50 @@ impl UI { port.port, port.desc ); - } - } - } else { - queue!( - stdout, - PrintStyledContent( - format!("{:^columns$}", "Not Connected") - .with(Color::Black) - .on(Color::Red) - ) - )?; } - - // Log - if show_logs { - let hr: usize = ((rows / 2) - 2).into(); - let start: usize = if lines.len() > hr { - lines.len() - hr - } else { - 0 - }; - - queue!(stdout, MoveTo(0, rows / 2))?; - print!("{}", format!("{:columns$}", " Log").negative()); - for line in lines.range(start..) { - print!("{}\r\n", line); - } - } - - queue!(stdout, MoveTo(0, rows - 1))?; - print!( - "{}", - format!( - "{:columns$}", - " q - quit | l - toggle log | \u{2191}/\u{2193} - select port | e - enable/disable | - browse" - ).negative() - ); - stdout.flush()?; } + // Log + if self.show_logs { + let hr: usize = ((rows / 2) - 2).into(); + let start: usize = if self.lines.len() > hr { + self.lines.len() - hr + } else { + 0 + }; + + queue!(stdout, MoveTo(0, rows / 2))?; + print!("{}", format!("{:columns$}", " Log").negative()); + for line in self.lines.range(start..) { + print!("{}\r\n", line); + } + } + + queue!(stdout, MoveTo(0, rows - 1))?; + print!( + "{}", + format!( + "{:columns$}", + " q - quit | l - toggle log | \u{2191}/\u{2193} - select port | e - enable/disable | - browse" + ).negative() + ); + stdout.flush()?; Ok(()) } + fn connected(&self) -> bool { + self.socks_port != 0 + } + + fn get_selected_port(&self) -> Option<&PortDesc> { + if let Some(p) = &self.ports { + if self.selection < p.len() { + return Some(&p[self.selection]); + } + } + None + } + fn enable_disable_port(&mut self, port: u16) { if self.socks_port != 0 { if let Some(_) = self.listeners.remove(&port) { @@ -299,4 +253,134 @@ impl UI { }); } } + + fn enter_alternate_screen(&mut self) -> Result<()> { + if !self.alternate_screen { + enable_raw_mode()?; + execute!(stdout(), EnterAlternateScreen, DisableLineWrap)?; + self.alternate_screen = true; + } + Ok(()) + } + + fn leave_alternate_screen(&mut self) -> Result<()> { + if self.alternate_screen { + execute!(stdout(), LeaveAlternateScreen, EnableLineWrap)?; + disable_raw_mode()?; + self.alternate_screen = false; + } + Ok(()) + } + + async fn handle_events(&mut self, console_events: &mut EventStream) { + tokio::select! { + ev = console_events.next() => self.handle_console_event(ev), + ev = self.events.recv() => self.handle_internal_event(ev), + } + } + + fn handle_console_event( + &mut self, + ev: Option>, + ) { + match ev { + Some(Ok(Event::Key(ev))) => match ev { + KeyEvent { code: KeyCode::Char('c'), .. } => { + if ev.modifiers.intersects(KeyModifiers::CONTROL) { + self.running = false; + } + } + KeyEvent { code: KeyCode::Esc, .. } + | KeyEvent { code: KeyCode::Char('q'), .. } => { + self.running = false; + } + KeyEvent { code: KeyCode::Char('l'), .. } => { + self.show_logs = !self.show_logs; + } + KeyEvent { code: KeyCode::Char('e'), .. } => { + if let Some(p) = self.get_selected_port() { + self.enable_disable_port(p.port); + } + } + KeyEvent { code: KeyCode::Up, .. } + | KeyEvent { code: KeyCode::Char('j'), .. } => { + if self.selection > 0 { + self.selection -= 1; + } + } + KeyEvent { code: KeyCode::Down, .. } + | KeyEvent { code: KeyCode::Char('k'), .. } => { + if let Some(p) = &self.ports { + if self.selection != p.len() - 1 { + self.selection += 1; + } + } + } + KeyEvent { code: KeyCode::Enter, .. } => { + if let Some(p) = self.get_selected_port() { + _ = open::that(format!("http://127.0.0.1:{}/", p.port)); + } + } + _ => (), + }, + Some(Ok(_)) => (), // Don't care about this event... + Some(Err(_)) => (), // Hmmmmmm.....? + None => (), // ....no events? what? + } + } + + fn handle_internal_event(&mut self, event: Option) { + match event { + Some(UIEvent::Disconnected) => { + self.socks_port = 0; + } + Some(UIEvent::Connected(sp)) => { + self.socks_port = sp; + info!("Socks port {socks_port}", socks_port = self.socks_port); + } + Some(UIEvent::Ports(mut p)) => { + p.sort_by(|a, b| a.port.partial_cmp(&b.port).unwrap()); + if self.selection >= p.len() { + self.selection = p.len() - 1; + } + + let mut new_listeners = HashMap::new(); + for port in &p { + let port = port.port; + if let Some(l) = self.listeners.remove(&port) { + if !l.is_closed() { + // `l` here is, of course, the channel that we + // use to tell the listener task to stop (see the + // spawn call below). If it isn't closed then + // that means a spawn task is still running so we + // should just let it keep running and re-use the + // existing listener. + new_listeners.insert(port, l); + } + } + } + + // This has the side effect of closing any + // listener that didn't get copied over to the + // new listeners table. + self.listeners = new_listeners; + self.ports = Some(p); + } + Some(UIEvent::ServerLine(line)) => { + while self.lines.len() >= 1024 { + self.lines.pop_front(); + } + self.lines.push_back(format!("[SERVER] {line}")); + } + Some(UIEvent::LogLine(_level, line)) => { + while self.lines.len() >= 1024 { + self.lines.pop_front(); + } + self.lines.push_back(format!("[CLIENT] {line}")); + } + None => { + self.running = false; + } + } + } }