diff --git a/src/ui.rs b/src/ui.rs index f56a1e4..199b5c6 100644 --- a/src/ui.rs +++ b/src/ui.rs @@ -15,7 +15,7 @@ use crossterm::{ use log::{error, info, Level, Metadata, Record}; use open; use std::collections::vec_deque::VecDeque; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::io::{stdout, Write}; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -57,15 +57,79 @@ impl log::Log for Logger { fn flush(&self) {} } +#[derive(Debug)] +struct Listener { + enabled: bool, + stop: Option>, + desc: Option, +} + +impl Listener { + pub fn from_desc(desc: PortDesc) -> Listener { + Listener { + enabled: false, + stop: None, + desc: Some(desc), + } + } + + pub fn enabled(&self) -> bool { + self.enabled + } + + pub fn set_enabled(&mut self, socks_port: Option, enabled: bool) { + if enabled { + self.enabled = true; + self.start(socks_port); + } else { + self.enabled = false; + self.stop = None; + } + } + + pub fn connect(&mut self, socks_port: Option, desc: PortDesc) { + self.desc = Some(desc); + self.start(socks_port); + } + + pub fn disconnect(&mut self) { + self.desc = None; + self.stop = None; + } + + pub fn start(&mut self, socks_port: Option) { + if self.enabled { + if let (Some(desc), Some(socks_port), None) = + (&self.desc, socks_port, &self.stop) + { + let (l, stop) = oneshot::channel(); + let port = desc.port; + tokio::spawn(async move { + let result = tokio::select! { + r = client_listen(port, socks_port) => r, + _ = stop => Ok(()), + }; + if let Err(e) = result { + error!("Error listening on port {port}: {e:?}"); + } else { + info!("Stopped listening on port {port}"); + } + }); + + self.stop = Some(l); + } + } + } +} + #[derive(Debug)] pub struct UI { events: mpsc::Receiver, - listeners: HashMap>, - socks_port: u16, + ports: HashMap, + socks_port: Option, running: bool, show_logs: bool, selection: usize, - ports: Option>, lines: VecDeque, alternate_screen: bool, raw_mode: bool, @@ -75,12 +139,11 @@ impl UI { pub fn new(events: mpsc::Receiver) -> UI { UI { events, - listeners: HashMap::new(), - socks_port: 0, + ports: HashMap::new(), + socks_port: None, running: true, show_logs: false, selection: 0, - ports: None, lines: VecDeque::with_capacity(1024), alternate_screen: false, raw_mode: false, @@ -160,30 +223,31 @@ impl UI { ) .negative() ); - if let Some(ports) = &self.ports { - let max_ports: usize = if self.show_logs { - (rows / 2) - 1 + + let max_ports: usize = if self.show_logs { + (rows / 2) - 1 + } else { + rows - 2 + } + .into(); + + let ports = self.get_ui_ports(); + for (index, port) in ports.into_iter().take(max_ports).enumerate() { + let listener = self.ports.get(&port).unwrap(); + + let caret = if index == self.selection { + "\u{2B46}" } else { - rows - 2 - } - .into(); - for (index, port) in ports.into_iter().take(max_ports).enumerate() { - print!( - "{} {:>enabled_width$} {:port_width$} {:description_width$}\r\n", - if index == self.selection { - "\u{2B46}" - } else { - " " - }, - if self.listeners.contains_key(&port.port) { - "+" - } else { - " " - }, - port.port, - port.desc - ); - } + " " + }; + + let state = if listener.enabled { "+" } else { " " }; + let desc: &str = match &listener.desc { + Some(port_desc) => &port_desc.desc, + None => "", + }; + + print!("{caret} {state:>enabled_width$} {port:port_width$} {desc:description_width$}\r\n"); } // Log @@ -215,40 +279,29 @@ impl UI { } fn connected(&self) -> bool { - self.socks_port != 0 + match self.socks_port { + Some(_) => true, + None => false, + } } - fn get_selected_port(&self) -> Option<&PortDesc> { - if let Some(p) = &self.ports { - if self.selection < p.len() { - return Some(&p[self.selection]); - } + fn get_ui_ports(&self) -> Vec { + let mut ports: Vec = self.ports.keys().copied().collect(); + ports.sort(); + ports + } + + fn get_selected_port(&self) -> Option { + if self.selection < self.ports.len() { + Some(self.get_ui_ports()[self.selection]) + } else { + None } - None } fn enable_disable_port(&mut self, port: u16) { - if self.socks_port != 0 { - if let Some(_) = self.listeners.remove(&port) { - return; // We disabled the listener. - } - - // We need to enable the listener. - let (l, stop) = oneshot::channel(); - self.listeners.insert(port, l); - - let socks_port = self.socks_port; - tokio::spawn(async move { - let result = tokio::select! { - r = client_listen(port, socks_port) => r, - _ = stop => Ok(()), - }; - if let Err(e) = result { - error!("Error listening on port {port}: {e:?}"); - } else { - info!("Stopped listening on port {port}"); - } - }); + if let Some(listener) = self.ports.get_mut(&port) { + listener.set_enabled(self.socks_port, !listener.enabled()); } } @@ -313,7 +366,7 @@ impl UI { } KeyEvent { code: KeyCode::Char('e'), .. } => { if let Some(p) = self.get_selected_port() { - self.enable_disable_port(p.port); + self.enable_disable_port(p); } } KeyEvent { code: KeyCode::Up, .. } @@ -324,15 +377,13 @@ impl UI { } KeyEvent { code: KeyCode::Down, .. } | KeyEvent { code: KeyCode::Char('k'), .. } => { - if let Some(p) = &self.ports { - if self.selection != p.len() - 1 { - self.selection += 1; - } + if self.selection != self.ports.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)); + _ = open::that(format!("http://127.0.0.1:{}/", p)); } } _ => (), @@ -346,40 +397,51 @@ impl UI { fn handle_internal_event(&mut self, event: Option) { match event { Some(UIEvent::Disconnected) => { - self.socks_port = 0; - self.listeners = HashMap::new(); // Bye. + self.socks_port = None; + for port in self.ports.values_mut() { + port.disconnect(); + } } 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() && p.len() > 0 { - self.selection = p.len() - 1; + info!("Socks port {sp}"); + self.socks_port = Some(sp); + for port in self.ports.values_mut() { + port.start(self.socks_port); } + } + Some(UIEvent::Ports(p)) => { + let mut leftover_ports: HashSet = + HashSet::from_iter(self.ports.keys().copied()); - 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); - } + for port_desc in p.into_iter() { + leftover_ports.remove(&port_desc.port); + if let Some(listener) = self.ports.get_mut(&port_desc.port) + { + listener.connect(self.socks_port, port_desc); + } else { + self.ports.insert( + port_desc.port, + Listener::from_desc(port_desc), + ); } } - // 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); + for port in leftover_ports { + let mut enabled = false; + if let Some(listener) = self.ports.get_mut(&port) { + enabled = listener.enabled; + listener.disconnect(); + } + + if !enabled { + // Just... forget it? Or leave it around forever? + self.ports.remove(&port); + } + } + + if self.selection >= self.ports.len() && self.ports.len() > 0 { + self.selection = self.ports.len() - 1; + } } Some(UIEvent::ServerLine(line)) => { while self.lines.len() >= 1024 {