Make a fine distinction between enabled and connected

This is part of remembering which ports should be enabled and which
not; track ports even if they've gone away or between connection and
disconnection, &c.

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<oneshot::Sender<()>>,
+    desc: Option<PortDesc>,
+}
+
+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<u16>, 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<u16>, 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<u16>) {
+        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<UIEvent>,
-    listeners: HashMap<u16, oneshot::Sender<()>>,
+    ports: HashMap<u16, Listener>,
-    socks_port: u16,
+    socks_port: Option<u16>,
     running: bool,
     show_logs: bool,
     selection: usize,
-    ports: Option<Vec<PortDesc>>,
     lines: VecDeque<String>,
     alternate_screen: bool,
     raw_mode: bool,
@@ -75,12 +139,11 @@ impl UI {
     pub fn new(events: mpsc::Receiver<UIEvent>) -> UI {
         UI {
             events,
-            listeners: HashMap::new(),
+            ports: HashMap::new(),
-            socks_port: 0,
+            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 {
+        let max_ports: usize = if self.show_logs {
-                (rows / 2) - 1
+            (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() {
+            let state = if listener.enabled { "+" } else { " " };
-                print!(
+            let desc: &str = match &listener.desc {
-                            "{} {:>enabled_width$} {:port_width$} {:description_width$}\r\n",
+                Some(port_desc) => &port_desc.desc,
-                            if index == self.selection {
+                None => "",
-                                "\u{2B46}"
+            };
-                            } else {
+
-                                " "
+            print!("{caret} {state:>enabled_width$} {port:port_width$} {desc:description_width$}\r\n");
-                            },
-                            if self.listeners.contains_key(&port.port) {
-                                "+"
-                            } else {
-                                " "
-                            },
-                            port.port,
-                            port.desc
-                        );
-            }
         }
 
         // 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> {
+    fn get_ui_ports(&self) -> Vec<u16> {
-        if let Some(p) = &self.ports {
+        let mut ports: Vec<u16> = self.ports.keys().copied().collect();
-            if self.selection < p.len() {
+        ports.sort();
-                return Some(&p[self.selection]);
+        ports
-            }
+    }
+
+    fn get_selected_port(&self) -> Option<u16> {
+        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(listener) = self.ports.get_mut(&port) {
-            if let Some(_) = self.listeners.remove(&port) {
+            listener.set_enabled(self.socks_port, !listener.enabled());
-                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}");
-                }
-            });
         }
     }
 
@@ -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 != self.ports.len() - 1 {
-                        if self.selection != p.len() - 1 {
+                        self.selection += 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<UIEvent>) {
         match event {
             Some(UIEvent::Disconnected) => {
-                self.socks_port = 0;
+                self.socks_port = None;
-                self.listeners = HashMap::new(); // Bye.
+                for port in self.ports.values_mut() {
+                    port.disconnect();
+                }
             }
             Some(UIEvent::Connected(sp)) => {
-                self.socks_port = sp;
+                info!("Socks port {sp}");
-                info!("Socks port {socks_port}", socks_port = self.socks_port);
+                self.socks_port = Some(sp);
-            }
+                for port in self.ports.values_mut() {
-            Some(UIEvent::Ports(mut p)) => {
+                    port.start(self.socks_port);
-                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;
                 }
+            }
+            Some(UIEvent::Ports(p)) => {
+                let mut leftover_ports: HashSet<u16> =
+                    HashSet::from_iter(self.ports.keys().copied());
 
-                let mut new_listeners = HashMap::new();
+                for port_desc in p.into_iter() {
-                for port in &p {
+                    leftover_ports.remove(&port_desc.port);
-                    let port = port.port;
+                    if let Some(listener) = self.ports.get_mut(&port_desc.port)
-                    if let Some(l) = self.listeners.remove(&port) {
+                    {
-                        if !l.is_closed() {
+                        listener.connect(self.socks_port, port_desc);
-                            // `l` here is, of course, the channel that we
+                    } else {
-                            // use to tell the listener task to stop (see the
+                        self.ports.insert(
-                            // spawn call below). If it isn't closed then
+                            port_desc.port,
-                            // that means a spawn task is still running so we
+                            Listener::from_desc(port_desc),
-                            // 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
+                for port in leftover_ports {
-                // listener that didn't get copied over to the
+                    let mut enabled = false;
-                // new listeners table.
+                    if let Some(listener) = self.ports.get_mut(&port) {
-                self.listeners = new_listeners;
+                        enabled = listener.enabled;
-                self.ports = Some(p);
+                        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 {