Enable/disable ports

src/lib.rs

@@ -3,7 +3,6 @@ use connection::ConnectionTable;
 use log::LevelFilter;
 use log::{error, info, warn};
 use message::{Message, MessageReader, MessageWriter};
-use std::collections::HashMap;
 use std::net::{Ipv4Addr, SocketAddrV4};
 use tokio::io::{
     AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader,
@@ -234,12 +233,9 @@ async fn client_listen(
 
 async fn client_read<T: AsyncRead + Unpin>(
     reader: &mut MessageReader<T>,
-    writer: mpsc::Sender<Message>,
     connections: ConnectionTable,
     events: mpsc::Sender<ui::UIEvent>,
 ) -> Result<()> {
-    let mut listeners: HashMap<u16, oneshot::Sender<()>> = HashMap::new();
-
     info!("Running");
     loop {
         let message = reader.read().await?;
@@ -267,41 +263,6 @@ async fn client_read<T: AsyncRead + Unpin>(
                 });
             }
             Ports(ports) => {
-                let mut new_listeners = HashMap::new();
-                for port in &ports {
-                    let port = port.port;
-                    if let Some(l) = 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);
-                        }
-                    }
-
-                    if !new_listeners.contains_key(&port) {
-                        let (l, stop) = oneshot::channel();
-                        new_listeners.insert(port, l);
-
-                        let (writer, connections) = (writer.clone(), connections.clone());
-                        tokio::spawn(async move {
-                            let result = tokio::select! {
-                                r = client_listen(port, writer, connections) => r,
-                                _ = stop => Ok(()),
-                            };
-                            if let Err(e) = result {
-                                error!("Error listening on port {port}: {e:?}");
-                            } else {
-                                info!("Stopped listening on port {port}");
-                            }
-                        });
-                    }
-                }
-
-                listeners = new_listeners;
                 if let Err(_) = events.send(ui::UIEvent::Ports(ports)).await {
                     // TODO: Log
                 }
@@ -351,10 +312,16 @@ async fn client_main<Reader: AsyncRead + Unpin, Writer: AsyncWrite + Unpin>(
 
     // And now really get into it...
     let (msg_sender, mut msg_receiver) = mpsc::channel(32);
-    let refresher = msg_sender.clone(); // Special for loop.
+
+    _ = events
+        .send(ui::UIEvent::Connected(
+            msg_sender.clone(),
+            connections.clone(),
+        ))
+        .await;
 
     let writing = pump_write(&mut msg_receiver, writer);
-    let reading = client_read(reader, msg_sender, connections, events);
+    let reading = client_read(reader, connections, events);
     tokio::pin!(reading);
     tokio::pin!(writing);
 
@@ -364,7 +331,7 @@ async fn client_main<Reader: AsyncRead + Unpin, Writer: AsyncWrite + Unpin>(
             result = async {
                 loop {
                     use tokio::time::{sleep, Duration};
-                    if let Err(e) = refresher.send(Message::Refresh).await {
+                    if let Err(e) = msg_sender.send(Message::Refresh).await {
                         break Err::<(), _>(e);
                     }
                     sleep(Duration::from_millis(500)).await;
@@ -458,8 +425,6 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender<ui::UIEvent>) {
             continue;
         }
 
-        _ = events.send(ui::UIEvent::Connected).await;
-
         let mut writer = MessageWriter::new(BufWriter::new(writer));
         let mut reader = MessageReader::new(reader);
 
@@ -477,13 +442,15 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender<ui::UIEvent>) {
 }
 
 pub async fn run_client(remote: &str) {
-    let (event_sender, mut event_receiver) = mpsc::channel(1024);
+    let (event_sender, event_receiver) = mpsc::channel(1024);
     _ = log::set_boxed_logger(ui::Logger::new(event_sender.clone()));
     log::set_max_level(LevelFilter::Info);
 
+    let mut ui = ui::UI::new(event_receiver);
+
     // Start the reconnect loop.
     tokio::select! {
-        _ = ui::run_ui(&mut event_receiver) => (),
+        _ = ui.run() => (),
         _ = client_connect_loop(remote, event_sender) => ()
     }
 }

src/ui.rs

@@ -1,4 +1,8 @@
-use crate::message::PortDesc;
+use crate::client_listen;
+use crate::{
+    message::{Message, PortDesc},
+    ConnectionTable,
+};
 use anyhow::Result;
 use crossterm::{
     cursor::MoveTo,
@@ -10,15 +14,17 @@ use crossterm::{
         EnterAlternateScreen, LeaveAlternateScreen,
     },
 };
-use log::{Level, Metadata, Record};
+use log::{error, info, Level, Metadata, Record};
 use open;
 use std::collections::vec_deque::VecDeque;
+use std::collections::HashMap;
 use std::io::{stdout, Write};
 use tokio::sync::mpsc;
+use tokio::sync::oneshot;
 use tokio_stream::StreamExt;
 
 pub enum UIEvent {
-    Connected,
+    Connected(mpsc::Sender<Message>, ConnectionTable),
     Disconnected,
     ServerLine(String),
     LogLine(log::Level, String),
@@ -53,15 +59,32 @@ impl log::Log for Logger {
     fn flush(&self) {}
 }
 
-pub async fn run_ui(events: &mut mpsc::Receiver<UIEvent>) -> Result<()> {
+pub struct UI {
+    events: mpsc::Receiver<UIEvent>,
+    writer: Option<mpsc::Sender<Message>>,
+    connections: Option<ConnectionTable>,
+    listeners: HashMap<u16, oneshot::Sender<()>>,
+}
+
+impl UI {
+    pub fn new(events: mpsc::Receiver<UIEvent>) -> UI {
+        UI {
+            events,
+            writer: None,
+            connections: None,
+            listeners: HashMap::new(),
+        }
+    }
+
+    pub async fn run(&mut self) -> Result<()> {
         enable_raw_mode()?;
-    let result = run_ui_core(events).await;
+        let result = self.run_core().await;
         execute!(stdout(), EnableLineWrap, LeaveAlternateScreen)?;
         disable_raw_mode()?;
         result
-}
+    }
 
-async fn run_ui_core(events: &mut mpsc::Receiver<UIEvent>) -> Result<()> {
+    async fn run_core(&mut self) -> Result<()> {
         let mut stdout = stdout();
 
         execute!(stdout, EnterAlternateScreen, DisableLineWrap)?;
@@ -83,6 +106,14 @@ async fn run_ui_core(events: &mut mpsc::Receiver<UIEvent>) -> Result<()> {
                                 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 {
@@ -112,16 +143,44 @@ async fn run_ui_core(events: &mut mpsc::Receiver<UIEvent>) -> Result<()> {
                         None => (),         // ....no events? what?
                     }
                 }
-            ev = events.recv() => {
+                ev = self.events.recv() => {
                     match ev {
                         Some(UIEvent::Disconnected) => {
+                            self.writer = None;
+                            self.connections = None;
                             connected = false;
                         }
-                    Some(UIEvent::Connected) => {
+                        Some(UIEvent::Connected(w,t)) => {
+                            self.writer = Some(w);
+                            self.connections = Some(t);
                             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;
+                            }
+
+                            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)) => {
@@ -149,14 +208,17 @@ async fn run_ui_core(events: &mut mpsc::Receiver<UIEvent>) -> Result<()> {
                 // 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 + port_width + padding);
+                let description_width =
+                    columns - (padding + padding + enabled_width + padding + port_width + padding);
 
                 print!(
                     "{}",
                     format!(
-                    "  {port:>port_width$} {description:<description_width$}\r\n",
+                        "  {enabled:>enabled_width$} {port:>port_width$} {description:<description_width$}\r\n",
+                        enabled = "E",
                         port = "port",
                         description = "description"
                     )
@@ -166,8 +228,13 @@ async fn run_ui_core(events: &mut mpsc::Receiver<UIEvent>) -> Result<()> {
                     let max_ports: usize = if show_logs { (rows / 2) - 1 } else { rows - 2 }.into();
                     for (index, port) in ports.into_iter().take(max_ports).enumerate() {
                         print!(
-                        "{} {:port_width$} {:description_width$}\r\n",
+                            "{} {:>enabled_width$} {:port_width$} {:description_width$}\r\n",
                             if index == selection { "\u{2B46}" } else { " " },
+                            if self.listeners.contains_key(&port.port) {
+                                "+"
+                            } else {
+                                " "
+                            },
                             port.port,
                             port.desc
                         );
@@ -205,12 +272,37 @@ async fn run_ui_core(events: &mut mpsc::Receiver<UIEvent>) -> Result<()> {
                 "{}",
                 format!(
                     "{:columns$}",
-                " q - quit  |  l - toggle log  |  \u{2191}/\u{2193} - select port  | <enter> - browse"
+                    " q - quit  |  l - toggle log  |  \u{2191}/\u{2193} - select port  | e - enable/disable | <enter> - browse"
-            )
+                ).negative()
-            .negative()
             );
             stdout.flush()?;
         }
 
         Ok(())
+    }
+
+    fn enable_disable_port(&mut self, port: u16) {
+        if let (Some(writer), Some(connections)) = (&self.writer, &self.connections) {
+            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 (writer, connections) = (writer.clone(), connections.clone());
+            tokio::spawn(async move {
+                let result = tokio::select! {
+                    r = client_listen(port, writer, connections) => r,
+                    _ = stop => Ok(()),
+                };
+                if let Err(e) = result {
+                    error!("Error listening on port {port}: {e:?}");
+                } else {
+                    info!("Stopped listening on port {port}");
+                }
+            });
+        }
+    }
 }