Fix up port state interactions in the face of configuration

src/client/config.rs

@@ -1,4 +1,5 @@
 use anyhow::{bail, Result};
+use std::collections::hash_map;
 use std::collections::HashMap;
 use toml::Value;
 
@@ -25,6 +26,10 @@ impl ServerConfig {
         self.ports.insert(port, config);
     }
 
+    pub fn iter(&self) -> hash_map::Iter<u16, PortConfig> {
+        self.ports.iter()
+    }
+
     pub fn contains_key(&self, port: u16) -> bool {
         self.ports.contains_key(&port)
     }
@@ -69,7 +74,7 @@ pub fn load_config() -> Result<Config> {
         },
     };
 
-    Ok(parse_config(&contents.parse::<Value>()?)?)
+    parse_config(&contents.parse::<Value>()?)
 }
 
 fn default() -> Config {
@@ -84,10 +89,7 @@ fn parse_config(value: &Value) -> Result<Config> {
                 Some(Value::Boolean(v)) => *v,
                 Some(v) => bail!("expected a true or false, got {:?}", v),
             };
-            Config {
+            Config { auto, servers: get_servers(table, auto)? }
-                auto,
-                servers: get_servers(&table, auto)?,
-            }
         }),
         _ => bail!("top level must be a table"),
     }

src/client/mod.rs

@@ -432,8 +432,13 @@ pub async fn run_client(remote: &str, sudo: bool, log_filter: &str) {
     _ = log::set_boxed_logger(ui::Logger::new(event_sender.clone()));
     log::set_max_level(LevelFilter::Info);
 
+    let server = if let Some((_user, server)) = remote.split_once("@") {
+        server
+    } else {
+        remote
+    };
     let config = match config::load_config() {
-        Ok(config) => config.get(remote),
+        Ok(config) => config.get(server),
         Err(e) => {
             eprintln!("Error loading configuration: {:?}", e);
             return;

src/client/ui.rs

@@ -1,4 +1,7 @@
-use super::{client_listen, config::ServerConfig};
+use super::{
+    client_listen,
+    config::{PortConfig, ServerConfig},
+};
 use crate::message::PortDesc;
 use anyhow::Result;
 use copypasta::{ClipboardContext, ClipboardProvider};
@@ -75,6 +78,7 @@ pub enum State {
     Enabled,
     Broken,
     Disabled,
+    Configured,
 }
 
 impl State {
@@ -86,35 +90,44 @@ impl State {
 #[derive(Debug)]
 struct Listener {
     state: std::sync::Arc<std::sync::Mutex<State>>,
+    config: Option<PortConfig>,
     stop: Option<oneshot::Sender<()>>,
     desc: Option<PortDesc>,
 }
 
 impl Listener {
-    pub fn from_desc(
+    pub fn from_desc(socks_port: Option<u16>, desc: PortDesc) -> Listener {
-        socks_port: Option<u16>,
-        desc: PortDesc,
-        enabled: bool,
-    ) -> Listener {
         let mut listener = Listener {
-            state: if enabled {
+            state: State::Enabled.boxed(),
-                State::Enabled.boxed()
+            config: None,
-            } else {
-                State::Disabled.boxed()
-            },
             stop: None,
             desc: Some(desc),
         };
-        if enabled {
+        listener.start(socks_port);
-            listener.start(socks_port);
-        }
         listener
     }
 
+    pub fn from_config(config: PortConfig) -> Self {
+        Listener {
+            state: State::Configured.boxed(),
+            config: Some(config),
+            stop: None,
+            desc: None,
+        }
+    }
+
     pub fn enabled(&self) -> bool {
         self.state() == State::Enabled
     }
 
+    pub fn description(&self) -> &str {
+        self.config
+            .as_ref()
+            .and_then(|c| c.description.as_deref())
+            .or_else(|| self.desc.as_ref().map(|d| d.desc.as_str()))
+            .unwrap_or("")
+    }
+
     fn state(&self) -> State {
         *self.state.lock().unwrap()
     }
@@ -130,6 +143,16 @@ impl Listener {
     }
 
     pub fn connect(&mut self, socks_port: Option<u16>, desc: PortDesc) {
+        // If we're just sitting idle and the port comes in from the remote
+        // server then we should become enabled. Otherwise we should become
+        // real, but disabled.
+        if self.state() == State::Configured {
+            if self.config.as_ref().unwrap().enabled {
+                self.state = State::Enabled.boxed();
+            } else {
+                self.state = State::Disabled.boxed();
+            }
+        }
         self.desc = Some(desc);
         self.start(socks_port);
     }
@@ -137,6 +160,13 @@ impl Listener {
     pub fn disconnect(&mut self) {
         self.desc = None;
         self.stop = None;
+
+        // When we get disconnected, but we're present in the configuration,
+        // we go back to being merely 'Configured'. If the port shows up
+        // again, our auto-enable behavior will depend on the configuration.
+        if self.config.is_some() {
+            self.state = State::Configured.boxed();
+        }
     }
 
     pub fn start(&mut self, socks_port: Option<u16>) {
@@ -185,9 +215,14 @@ pub struct UI {
 
 impl UI {
     pub fn new(events: mpsc::Receiver<UIEvent>, config: ServerConfig) -> UI {
+        let mut ports = HashMap::new();
+        for (port, config) in config.iter() {
+            ports.insert(*port, Listener::from_config(config.clone()));
+        }
+
         UI {
             events,
-            ports: HashMap::new(),
+            ports,
             socks_port: None,
             running: true,
             show_logs: false,
@@ -288,13 +323,15 @@ impl UI {
             let (symbol, style) = match listener.state() {
                 State::Enabled => (" ✓ ", enabled_port_style),
                 State::Broken => (" ✗ ", broken_port_style),
-                State::Disabled => ("", disabled_port_style),
+                State::Disabled | State::Configured => {
+                    ("", disabled_port_style)
+                }
             };
             rows.push(
                 Row::new(vec![
                     symbol,
                     &*port_strings[index],
-                    listener.desc.as_ref().map(|pd| &*pd.desc).unwrap_or(""),
+                    listener.description(),
                 ])
                 .style(style),
             );
@@ -390,11 +427,7 @@ impl UI {
     }
 
     fn enable_disable_port(&mut self, port: u16) {
-        let state = self.ports.get(&port).map(Listener::state);
+        if let Some(listener) = self.ports.get_mut(&port) {
-        if state == Some(State::Broken) {
-            // try turning it off and on again, it will at least get logs visible
-            self.ports.remove(&port);
-        } else if let Some(listener) = self.ports.get_mut(&port) {
             listener.set_enabled(self.socks_port, !listener.enabled());
         }
     }
@@ -571,24 +604,16 @@ impl UI {
                 // Grab the selected port
                 let selected_port = self.get_selected_port();
 
-                for mut port_desc in p.into_iter() {
+                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 {
-                        let config = self.config.get(port_desc.port);
+                        assert!(!self.config.contains_key(port_desc.port));
-                        info!("Port config {port_desc:?} -> {config:?}");
-                        port_desc.desc =
-                            config.description.unwrap_or(port_desc.desc);
-
                         self.ports.insert(
                             port_desc.port,
-                            Listener::from_desc(
+                            Listener::from_desc(self.socks_port, port_desc),
-                                self.socks_port,
-                                port_desc,
-                                config.enabled,
-                            ),
                         );
                     }
                 }
@@ -961,8 +986,8 @@ mod tests {
             desc: "my-service".to_string(),
         }])));
 
-        let description = ui.ports.get(&8080).unwrap().desc.as_ref().unwrap();
+        let listener = ui.ports.get(&8080).unwrap();
-        assert_eq!(description.desc, "override");
+        assert_eq!(listener.description(), "override");
 
         drop(sender);
     }
@@ -992,4 +1017,104 @@ mod tests {
 
         drop(sender);
     }
+
+    #[test]
+    fn port_config_missing_but_still_there() {
+        let (sender, receiver) = mpsc::channel(64);
+        let mut config = ServerConfig::default();
+        config.insert(
+            8080,
+            PortConfig {
+                enabled: false,
+                description: Some("override".to_string()),
+            },
+        );
+
+        let mut ui = UI::new(receiver, config);
+
+        // There are no ports...
+        ui.handle_internal_event(Some(UIEvent::Ports(vec![])));
+
+        // But there should still be ports, man.
+        let listener = ui.ports.get(&8080).unwrap();
+        assert_eq!(listener.state(), State::Configured);
+        assert_eq!(listener.description(), "override");
+
+        drop(sender);
+    }
+
+    #[test]
+    fn port_config_state_interactions() {
+        let (sender, receiver) = mpsc::channel(64);
+        let mut config = ServerConfig::default();
+        config.insert(8080, PortConfig { enabled: false, description: None });
+        config.insert(8081, PortConfig { enabled: true, description: None });
+
+        let mut ui = UI::new(receiver, config);
+
+        // No ports have been received, make sure everything's "configured"
+        assert_eq!(ui.ports.get(&8080).unwrap().state(), State::Configured);
+        assert_eq!(ui.ports.get(&8081).unwrap().state(), State::Configured);
+
+        // 8080 shows up.... not configured as enabled so it becomes "disabled"
+        ui.handle_internal_event(Some(UIEvent::Ports(vec![PortDesc {
+            port: 8080,
+            desc: "python3".to_string(),
+        }])));
+        assert_eq!(ui.ports.get(&8080).unwrap().state(), State::Disabled);
+        assert_eq!(ui.ports.get(&8081).unwrap().state(), State::Configured);
+
+        // 8081 shows up.... configured as enabled so it becomes "enabled"
+        ui.handle_internal_event(Some(UIEvent::Ports(vec![
+            PortDesc { port: 8080, desc: "python3".to_string() },
+            PortDesc { port: 8081, desc: "python3".to_string() },
+        ])));
+        assert_eq!(ui.ports.get(&8080).unwrap().state(), State::Disabled);
+        assert_eq!(ui.ports.get(&8081).unwrap().state(), State::Enabled);
+
+        // 8082 shows up.... it should be enabled by default!
+        ui.handle_internal_event(Some(UIEvent::Ports(vec![
+            PortDesc { port: 8080, desc: "python3".to_string() },
+            PortDesc { port: 8081, desc: "python3".to_string() },
+            PortDesc { port: 8082, desc: "python3".to_string() },
+        ])));
+        assert_eq!(ui.ports.get(&8080).unwrap().state(), State::Disabled);
+        assert_eq!(ui.ports.get(&8081).unwrap().state(), State::Enabled);
+        assert_eq!(ui.ports.get(&8082).unwrap().state(), State::Enabled);
+
+        // 8081 goes away.... back to configured.
+        ui.handle_internal_event(Some(UIEvent::Ports(vec![
+            PortDesc { port: 8080, desc: "python3".to_string() },
+            PortDesc { port: 8082, desc: "python3".to_string() },
+        ])));
+        assert_eq!(ui.ports.get(&8080).unwrap().state(), State::Disabled);
+        assert_eq!(ui.ports.get(&8081).unwrap().state(), State::Configured);
+        assert_eq!(ui.ports.get(&8082).unwrap().state(), State::Enabled);
+
+        // All gone, state resets itself.
+        ui.handle_internal_event(Some(UIEvent::Ports(vec![])));
+        assert_eq!(ui.ports.get(&8080).unwrap().state(), State::Configured);
+        assert_eq!(ui.ports.get(&8081).unwrap().state(), State::Configured);
+        assert!(!ui.ports.contains_key(&8082));
+
+        drop(sender);
+    }
+
+    #[test]
+    fn port_defaults_respected() {
+        let (sender, receiver) = mpsc::channel(64);
+        let config = ServerConfig::default();
+        let mut ui = UI::new(receiver, config);
+
+        ui.handle_internal_event(Some(UIEvent::Ports(vec![PortDesc {
+            port: 8080,
+            desc: "python3".to_string(),
+        }])));
+
+        let listener = ui.ports.get(&8080).unwrap();
+        assert_eq!(listener.state(), State::Enabled);
+        assert_eq!(listener.description(), "python3");
+
+        drop(sender);
+    }
 }