Split client and server into modules

src/client/mod.rs

@@ -0,0 +1,384 @@
+use crate::message::{Message, MessageReader, MessageWriter};
+use anyhow::{bail, Result};
+use bytes::BytesMut;
+use log::LevelFilter;
+use log::{debug, error, info, warn};
+use std::net::{Ipv4Addr, SocketAddrV4};
+use tokio::io::{
+    AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite,
+    AsyncWriteExt, BufReader, BufWriter,
+};
+use tokio::net::{TcpListener, TcpStream};
+use tokio::process;
+use tokio::sync::mpsc;
+
+mod ui;
+
+/// Wait for the server to be ready; we know the server is there and
+/// listening when we see the special sync marker, which is 8 NUL bytes in a
+/// row.
+async fn client_sync<S: AsyncRead + Unpin, T: AsyncRead + Unpin>(
+    reader: &mut S,
+    client_stderr: &mut T,
+) -> Result<(), tokio::io::Error> {
+    info!("Waiting for synchronization marker...");
+
+    let mut stderr = tokio::io::stderr();
+    let mut stdout = tokio::io::stdout();
+    let mut buf = BytesMut::with_capacity(1024);
+
+    let mut seen = 0;
+    let result = tokio::select! {
+        result = async {
+            loop {
+                buf.clear();
+                client_stderr.read_buf(&mut buf).await?;
+                stderr.write_all(&buf[..]).await?;
+            }
+        } => result,
+        result = async {
+            while seen < 8 {
+                let byte = reader.read_u8().await?;
+                if byte == 0 {
+                    seen += 1;
+                } else {
+                    stdout.write_u8(byte).await?;
+                    seen = 0;
+                }
+            }
+
+            Ok::<_, tokio::io::Error>(())
+        } => result,
+    };
+
+    if let Err(_) = result {
+        // Something went wrong, let's just make sure we flush the client's
+        // stderr before we return.
+        _ = stderr.write_all(&buf[..]).await;
+        _ = tokio::io::copy(client_stderr, &mut stderr).await;
+        _ = stderr.flush().await;
+    }
+
+    result
+}
+
+/// Handle an incoming client connection, by forwarding it to the SOCKS5
+/// server at the specified port. This is the core of the entire thing.
+///
+/// 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<()> {
+    debug!("Handling connection!");
+
+    let dest_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, socks_port);
+    let mut dest_socket = TcpStream::connect(dest_addr).await?;
+
+    debug!("Connected, sending handshake request");
+    let packet: [u8; 3] = [
+        0x05, // v5
+        0x01, // 1 auth method
+        0x00, // my one auth method is no auth
+    ];
+    dest_socket.write_all(&packet[..]).await?;
+    debug!("Initial handshake sent. Awaiting handshake response");
+
+    let mut response: [u8; 2] = [0; 2];
+    dest_socket.read_exact(&mut response).await?;
+    if response[0] != 0x05 {
+        bail!("SOCKS incorrect response version {}", response[0]);
+    }
+    if response[1] == 0xFF {
+        bail!("SOCKS server says no acceptable auth");
+    }
+    if response[1] != 0x00 {
+        bail!("SOCKS server chose something wild? {}", response[1]);
+    }
+
+    debug!("Handshake response received, sending connect request");
+    let packet: [u8; 10] = [
+        0x05,                                       // version again :P
+        0x01,                                       // connect
+        0x00,                                       // reserved!
+        0x01,                                       // ipv4
+        127,                                        // lo..
+        0,                                          // ..cal..
+        0,                                          // ..ho..
+        1,                                          // ..st
+        ((port & 0xFF00) >> 8).try_into().unwrap(), // port (high)
+        ((port & 0x00FF) >> 0).try_into().unwrap(), // port (low)
+    ];
+    dest_socket.write_all(&packet[..]).await?;
+
+    debug!("Connect request sent, awaiting response");
+    let mut response: [u8; 4] = [0; 4];
+    dest_socket.read_exact(&mut response).await?;
+    if response[0] != 0x05 {
+        bail!("SOCKS5 incorrect response version again? {}", response[0]);
+    }
+    if response[1] != 0x00 {
+        bail!("SOCKS5 reports a connect error {}", response[1]);
+    }
+    // Now we 100% do not care about the following information but we must
+    // discard it so we can get to the good stuff. response[3] is the type of
+    // address...
+    if response[3] == 0x01 {
+        // IPv4 - 4 bytes.
+        let mut response: [u8; 4] = [0; 4];
+        dest_socket.read_exact(&mut response).await?;
+    } else if response[3] == 0x03 {
+        // Domain Name
+        let len = dest_socket.read_u8().await?;
+        for _ in 0..len {
+            dest_socket.read_u8().await?; // So slow!
+        }
+    } else if response[3] == 0x04 {
+        // IPv6 - 8 bytes
+        let mut response: [u8; 8] = [0; 8];
+        dest_socket.read_exact(&mut response).await?;
+    } else {
+        bail!(
+            "SOCKS5 sent me an address I don't understand {}",
+            response[3]
+        );
+    }
+    // Finally the port number. Again, garbage, but it's in the packet we
+    // need to skip.
+    let mut response: [u8; 2] = [0; 2];
+    dest_socket.read_exact(&mut response).await?;
+
+    info!("Connection established on port {}", port);
+
+    let mut socket = socket;
+    tokio::io::copy_bidirectional(&mut socket, &mut dest_socket).await?;
+    Ok(())
+}
+
+/// Listen on a port that we are currently forwarding, and use the SOCKS5
+/// proxy on the specified port to handle the connections.
+async fn client_listen(port: u16, socks_port: u16) -> Result<()> {
+    loop {
+        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
+                {
+                    error!("Error handling connection: {:?}", e);
+                } else {
+                    debug!("Done???");
+                }
+            });
+        }
+    }
+}
+
+async fn client_handle_messages<T: AsyncRead + Unpin>(
+    reader: &mut MessageReader<T>,
+    events: mpsc::Sender<ui::UIEvent>,
+) -> Result<()> {
+    loop {
+        use Message::*;
+        match reader.read().await? {
+            Ping => (),
+            Ports(ports) => {
+                if let Err(_) = events.send(ui::UIEvent::Ports(ports)).await {
+                    // TODO: Log
+                }
+            }
+            message => panic!("Unsupported: {:?}", message),
+        };
+    }
+}
+
+async fn client_pipe_stderr<Debug: AsyncBufRead + Unpin>(
+    debug: &mut Debug,
+    events: mpsc::Sender<ui::UIEvent>,
+) {
+    loop {
+        let mut line = String::new();
+        match debug.read_line(&mut line).await {
+            Err(e) => {
+                error!("Error reading stderr from server: {:?}", e);
+                break;
+            }
+            Ok(0) => {
+                warn!("stderr stream closed");
+                break;
+            }
+            _ => {
+                _ = events.send(ui::UIEvent::ServerLine(line)).await;
+            }
+        }
+    }
+}
+
+async fn client_main<Reader: AsyncRead + Unpin, Writer: AsyncWrite + Unpin>(
+    socks_port: u16,
+    reader: &mut MessageReader<Reader>,
+    writer: &mut MessageWriter<Writer>,
+    events: mpsc::Sender<ui::UIEvent>,
+) -> Result<()> {
+    // Wait for the server's announcement.
+    if let Message::Hello(major, minor, _) = reader.read().await? {
+        if major != 0 || minor > 1 {
+            bail!("Unsupported remote protocol version {}.{}", major, minor);
+        }
+    } else {
+        bail!("Expected a hello message from the remote server");
+    }
+
+    // And now really get into it...
+    _ = events.send(ui::UIEvent::Connected(socks_port)).await;
+
+    tokio::select! {
+        result = async {
+            loop {
+                use tokio::time::{sleep, Duration};
+                if let Err(e) = writer.write(Message::Refresh).await {
+                    break Err::<(), _>(e);
+                }
+                sleep(Duration::from_millis(500)).await;
+            }
+        } => {
+            if let Err(e) = result {
+                return Err(e.into());
+            }
+        },
+        result = client_handle_messages(reader, events) => {
+            if let Err(e) = result {
+                return Err(e.into());
+            }
+        },
+    }
+    Ok(())
+}
+
+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()
+    };
+
+    let mut cmd = process::Command::new("ssh");
+    cmd.arg("-T")
+        .arg("-D")
+        .arg(socks_port.to_string())
+        .arg(server)
+        .arg("fwd")
+        .arg("--server");
+
+    cmd.stdout(std::process::Stdio::piped());
+    cmd.stdin(std::process::Stdio::piped());
+    cmd.stderr(std::process::Stdio::piped());
+    let child = cmd.spawn()?;
+    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<ui::UIEvent>) {
+    loop {
+        _ = events.send(ui::UIEvent::Disconnected).await;
+
+        let (mut child, socks_port) =
+            spawn_ssh(remote).await.expect("failed to spawn");
+
+        let mut stderr = child
+            .stderr
+            .take()
+            .expect("child did not have a handle to stderr");
+
+        let writer = child
+            .stdin
+            .take()
+            .expect("child did not have a handle to stdin");
+
+        let mut reader = BufReader::new(
+            child
+                .stdout
+                .take()
+                .expect("child did not have a handle to stdout"),
+        );
+
+        if let Err(e) = client_sync(&mut reader, &mut stderr).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;
+        }
+
+        let mut stderr = BufReader::new(stderr);
+        let mut writer = MessageWriter::new(BufWriter::new(writer));
+        let mut reader = MessageReader::new(reader);
+
+        let sec = events.clone();
+        tokio::spawn(async move {
+            client_pipe_stderr(&mut stderr, sec).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...");
+        }
+    }
+}
+
+pub async fn run_client(remote: &str) {
+    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() => (),
+        _ = client_connect_loop(remote, event_sender) => ()
+    }
+}

src/client/ui.rs

@@ -0,0 +1,470 @@
+use super::client_listen;
+use crate::message::PortDesc;
+use anyhow::Result;
+use crossterm::{
+    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,
+    },
+};
+use log::{error, info, Level, Metadata, Record};
+use open;
+use std::collections::vec_deque::VecDeque;
+use std::collections::{HashMap, HashSet};
+use std::io::{stdout, Write};
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio_stream::StreamExt;
+
+pub enum UIEvent {
+    Connected(u16),
+    Disconnected,
+    ServerLine(String),
+    LogLine(log::Level, String),
+    Ports(Vec<PortDesc>),
+}
+
+#[derive(Debug, Clone)]
+pub struct Logger {
+    line_sender: mpsc::Sender<UIEvent>,
+}
+
+impl Logger {
+    pub fn new(line_sender: mpsc::Sender<UIEvent>) -> Box<Logger> {
+        Box::new(Logger { line_sender })
+    }
+}
+
+impl log::Log for Logger {
+    fn enabled(&self, metadata: &Metadata) -> bool {
+        metadata.level() <= Level::Info
+    }
+
+    fn log(&self, record: &Record) {
+        if self.enabled(record.metadata()) {
+            let line = format!("{}", record.args());
+            _ = self
+                .line_sender
+                .try_send(UIEvent::LogLine(record.level(), line));
+        }
+    }
+
+    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>,
+    ports: HashMap<u16, Listener>,
+    socks_port: Option<u16>,
+    running: bool,
+    show_logs: bool,
+    selection: usize,
+    lines: VecDeque<String>,
+    alternate_screen: bool,
+    raw_mode: bool,
+}
+
+impl UI {
+    pub fn new(events: mpsc::Receiver<UIEvent>) -> UI {
+        UI {
+            events,
+            ports: HashMap::new(),
+            socks_port: None,
+            running: true,
+            show_logs: false,
+            selection: 0,
+            lines: VecDeque::with_capacity(1024),
+            alternate_screen: false,
+            raw_mode: false,
+        }
+    }
+
+    pub async fn run(&mut self) -> Result<()> {
+        let mut console_events = EventStream::new();
+
+        self.running = true;
+        while self.running {
+            self.handle_events(&mut console_events).await;
+
+            if self.connected() {
+                self.render_connected()?;
+            } else {
+                self.render_disconnected()?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn render_disconnected(&mut self) -> Result<()> {
+        self.enter_alternate_screen()?;
+        self.disable_raw_mode()?;
+        let mut stdout = stdout();
+
+        let (columns, _) = size()?;
+        let columns: usize = columns.into();
+
+        execute!(
+            stdout,
+            SavePosition,
+            MoveTo(0, 0),
+            PrintStyledContent(
+                format!("{:^columns$}\r\n", "Not Connected")
+                    .with(Color::Black)
+                    .on(Color::Red)
+            ),
+            RestorePosition,
+        )?;
+        Ok(())
+    }
+
+    fn render_connected(&mut self) -> Result<()> {
+        self.enter_alternate_screen()?;
+        self.enable_raw_mode()?;
+        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:<description_width$}\r\n",
+                enabled = "E",
+                port = "port",
+                description = "description"
+            )
+                .negative()
+        );
+
+        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 {
+                " "
+            };
+
+            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
+        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 | <enter> - browse"
+            ).negative()
+        );
+        stdout.flush()?;
+        Ok(())
+    }
+
+    fn connected(&self) -> bool {
+        match self.socks_port {
+            Some(_) => true,
+            None => false,
+        }
+    }
+
+    fn get_ui_ports(&self) -> Vec<u16> {
+        let mut ports: Vec<u16> = self.ports.keys().copied().collect();
+        ports.sort();
+        ports
+    }
+
+    fn get_selected_port(&self) -> Option<u16> {
+        if self.selection < self.ports.len() {
+            Some(self.get_ui_ports()[self.selection])
+        } else {
+            None
+        }
+    }
+
+    fn enable_disable_port(&mut self, port: u16) {
+        if let Some(listener) = self.ports.get_mut(&port) {
+            listener.set_enabled(self.socks_port, !listener.enabled());
+        }
+    }
+
+    fn enable_raw_mode(&mut self) -> Result<()> {
+        if !self.raw_mode {
+            enable_raw_mode()?;
+            self.raw_mode = true;
+        }
+        Ok(())
+    }
+
+    fn disable_raw_mode(&mut self) -> Result<()> {
+        if self.raw_mode {
+            disable_raw_mode()?;
+            self.raw_mode = false;
+        }
+        Ok(())
+    }
+
+    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<Result<Event, std::io::Error>>,
+    ) {
+        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);
+                    }
+                }
+                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 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));
+                    }
+                }
+                _ => (),
+            },
+            Some(Ok(_)) => (), // Don't care about this event...
+            Some(Err(_)) => (), // Hmmmmmm.....?
+            None => (),        // ....no events? what?
+        }
+    }
+
+    fn handle_internal_event(&mut self, event: Option<UIEvent>) {
+        match event {
+            Some(UIEvent::Disconnected) => {
+                self.socks_port = None;
+                for port in self.ports.values_mut() {
+                    port.disconnect();
+                }
+            }
+            Some(UIEvent::Connected(sp)) => {
+                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<u16> =
+                    HashSet::from_iter(self.ports.keys().copied());
+
+                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),
+                        );
+                    }
+                }
+
+                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 {
+                    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;
+            }
+        }
+    }
+}
+
+impl Drop for UI {
+    fn drop(&mut self) {
+        _ = self.disable_raw_mode();
+        _ = self.leave_alternate_screen();
+    }
+}