Implemented initial user interface

src/lib.rs

@@ -4,4 +4,6 @@ pub mod animals;
 pub mod auth;
 pub mod helper;
 pub mod ssh;
+pub mod terminal;
+pub mod tui;
 pub mod tunnel;

src/ssh.rs

@@ -1,54 +1,112 @@
-use std::{iter::once, net::SocketAddr, sync::Arc, time::Duration};
+use std::{io::Write, iter::once, net::SocketAddr, sync::Arc, time::Duration};
 
 use clap::Parser;
 use indexmap::IndexMap;
+use ratatui::{Terminal, TerminalOptions, Viewport, layout::Rect, prelude::CrosstermBackend};
 use russh::{
     ChannelId,
     keys::PrivateKey,
     server::{Auth, Msg, Server as _, Session},
 };
-use tokio::{
-    net::ToSocketAddrs,
-    sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel},
-};
+use tokio::net::ToSocketAddrs;
 use tracing::{debug, trace, warn};
 
-use crate::tunnel::{Tunnel, TunnelAccess, Tunnels};
+use crate::{
+    terminal::TerminalHandle,
+    tui::Renderer,
+    tunnel::{Tunnel, TunnelAccess, Tunnels},
+};
 
 pub struct Handler {
-    tx: UnboundedSender<Vec<u8>>,
-    rx: Option<UnboundedReceiver<Vec<u8>>>,
-
     all_tunnels: Tunnels,
     tunnels: IndexMap<String, Option<Tunnel>>,
 
-    access: Option<TunnelAccess>,
+    user: Option<String>,
+
+    terminal: Option<Terminal<CrosstermBackend<TerminalHandle>>>,
+    renderer: Renderer,
 }
 
 impl Handler {
-    fn send(&self, data: impl AsRef<str>) {
-        let _ = self.tx.send(data.as_ref().as_bytes().to_vec());
-    }
-
-    fn sendln(&self, data: impl AsRef<str>) {
-        self.send(format!("{}\n\r", data.as_ref()));
+    fn new(all_tunnels: Tunnels) -> Self {
+        Self {
+            all_tunnels,
+            tunnels: IndexMap::new(),
+            user: None,
+            terminal: Default::default(),
+            renderer: Default::default(),
+        }
     }
 
     async fn set_access(&mut self, access: TunnelAccess) {
-        self.access = Some(access.clone());
-
         for (_address, tunnel) in &self.tunnels {
             if let Some(tunnel) = tunnel {
                 tunnel.set_access(access.clone()).await;
             }
         }
     }
+
+    pub async fn resize(&mut self, width: u32, height: u32) -> std::io::Result<()> {
+        let rect = Rect {
+            x: 0,
+            y: 0,
+            width: width as u16,
+            height: height as u16,
+        };
+
+        if let Some(terminal) = &mut self.terminal {
+            terminal.resize(rect)?;
+        } else {
+            todo!()
+        }
+
+        self.redraw().await?;
+
+        Ok(())
+    }
+
+    pub fn close(&mut self) -> std::io::Result<()> {
+        if let Some(terminal) = self.terminal.take() {
+            drop(terminal);
+        }
+
+        Ok(())
+    }
+
+    pub async fn redraw(&mut self) -> std::io::Result<()> {
+        if let Some(terminal) = &mut self.terminal {
+            trace!("redraw");
+            self.renderer.update_table(&self.tunnels).await;
+            terminal.draw(|frame| {
+                self.renderer.render(frame);
+            })?;
+        } else {
+            todo!()
+        }
+
+        Ok(())
+    }
+
+    pub async fn handle_input(&mut self, input: char) -> std::io::Result<bool> {
+        match input {
+            'q' => {
+                self.close()?;
+                return Ok(false);
+            }
+            _ => {
+                return Ok(false);
+            }
+        };
+
+        Ok(true)
+    }
 }
 
 /// Quickly create http tunnels for development
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 struct Args {
+    /// Make all tunnels public by default instead of private
     #[arg(short, long)]
     public: bool,
 }
@@ -59,37 +117,16 @@ impl russh::server::Handler for Handler {
     async fn channel_open_session(
         &mut self,
         channel: russh::Channel<Msg>,
-        _session: &mut Session,
+        session: &mut Session,
     ) -> Result<bool, Self::Error> {
         trace!("channel_open_session");
 
-        let Some(mut rx) = self.rx.take() else {
-            return Err(russh::Error::Inconsistent);
+        let terminal_handle = TerminalHandle::start(session.handle(), channel.id()).await?;
+        let backend = CrosstermBackend::new(terminal_handle);
+        let options = TerminalOptions {
+            viewport: Viewport::Fixed(Rect::default()),
         };
-
-        tokio::spawn(async move {
-            loop {
-                let Some(message) = rx.recv().await else {
-                    break;
-                };
-
-                trace!("Sending message to client");
-
-                if channel.data(message.as_ref()).await.is_err() {
-                    break;
-                }
-            }
-        });
-
-        // NOTE: I believe this happens as the final step when opening a session.
-        // At this point all the tunnels should be populated
-        for (address, tunnel) in &self.tunnels {
-            if tunnel.is_some() {
-                self.sendln(format!("http://{address}"));
-            } else {
-                self.sendln(format!("Failed to open {address}, address already in use"));
-            }
-        }
+        self.terminal = Some(Terminal::with_options(backend, options)?);
 
         Ok(true)
     }
@@ -101,7 +138,7 @@ impl russh::server::Handler for Handler {
     ) -> Result<Auth, Self::Error> {
         debug!("Login from {user}");
 
-        self.set_access(TunnelAccess::Private(user.into())).await;
+        self.user = Some(user.into());
 
         // TODO: Get ssh keys associated with user from ldap
         Ok(Auth::Accept)
@@ -113,8 +150,14 @@ impl russh::server::Handler for Handler {
         data: &[u8],
         _session: &mut Session,
     ) -> Result<(), Self::Error> {
-        if data == [3] {
-            return Err(russh::Error::Disconnect);
+        let Some(input) = data.first().cloned() else {
+            return Ok(());
+        };
+
+        trace!(input, "data");
+
+        if self.handle_input(input as char).await? {
+            self.redraw().await?;
         }
 
         Ok(())
@@ -130,28 +173,32 @@ impl russh::server::Handler for Handler {
 
         trace!(?cmd, "exec_request");
 
-        let cmd = once("<ssh>").chain(cmd.split_whitespace());
+        let cmd = once("<ssh command> --").chain(cmd.split_whitespace());
         match Args::try_parse_from(cmd) {
             Ok(args) => {
                 debug!("{args:?}");
                 if args.public {
                     trace!("Making tunnels public");
                     self.set_access(TunnelAccess::Public).await;
+                    self.redraw().await?;
                 }
-
-                session.channel_success(channel)
             }
             Err(err) => {
-                trace!("Sending error/help message and disconnecting");
-                session.disconnect(
-                    russh::Disconnect::ByApplication,
-                    &format!("\n\r{err}"),
-                    "EN",
-                )?;
+                trace!("Sending help message and disconnecting");
 
-                session.channel_failure(channel)
+                if let Some(terminal) = &mut self.terminal {
+                    let writer = terminal.backend_mut().writer_mut();
+
+                    writer.leave_alternate_screen()?;
+                    writer.write_all(err.to_string().replace('\n', "\n\r").as_bytes())?;
+                    writer.flush()?;
+                }
+
+                self.close()?;
             }
         }
+
+        session.channel_success(channel)
     }
 
     async fn tcpip_forward(
@@ -162,19 +209,58 @@ impl russh::server::Handler for Handler {
     ) -> Result<bool, Self::Error> {
         trace!(address, port, "tcpip_forward");
 
-        let Some(access) = self.access.clone() else {
+        let Some(user) = self.user.clone() else {
             return Err(russh::Error::Inconsistent);
         };
 
-        let tunnel = Tunnel::new(session.handle(), address, *port, access);
-        let Some(address) = self.all_tunnels.add_tunnel(address, tunnel.clone()).await else {
-            self.tunnels.insert(address.into(), None);
-            return Ok(false);
-        };
+        let tunnel = Tunnel::new(
+            session.handle(),
+            address,
+            *port,
+            TunnelAccess::Private(user),
+        );
+        let (success, address) = self.all_tunnels.add_tunnel(address, tunnel.clone()).await;
 
-        self.tunnels.insert(address, Some(tunnel));
+        let tunnel = if success { Some(tunnel) } else { None };
+        self.tunnels.insert(address, tunnel);
 
-        Ok(true)
+        Ok(success)
+    }
+
+    async fn window_change_request(
+        &mut self,
+        _channel: ChannelId,
+        col_width: u32,
+        row_height: u32,
+        _pix_width: u32,
+        _pix_height: u32,
+        _session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        trace!(col_width, row_height, "window_change_request");
+
+        self.resize(col_width, row_height).await?;
+
+        Ok(())
+    }
+
+    async fn pty_request(
+        &mut self,
+        channel: ChannelId,
+        _term: &str,
+        col_width: u32,
+        row_height: u32,
+        _pix_width: u32,
+        _pix_height: u32,
+        _modes: &[(russh::Pty, u32)],
+        session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        trace!(col_width, row_height, "pty_request");
+
+        self.resize(col_width, row_height).await?;
+
+        session.channel_success(channel)?;
+
+        Ok(())
     }
 }
 
@@ -215,6 +301,7 @@ impl Server {
             preferred: russh::Preferred {
                 ..Default::default()
             },
+            nodelay: true,
             ..Default::default()
         };
         let config = Arc::new(config);
@@ -229,15 +316,7 @@ impl russh::server::Server for Server {
     type Handler = Handler;
 
     fn new_client(&mut self, _peer_addr: Option<SocketAddr>) -> Self::Handler {
-        let (tx, rx) = unbounded_channel::<Vec<u8>>();
-
-        Handler {
-            tx,
-            rx: Some(rx),
-            all_tunnels: self.tunnels.clone(),
-            tunnels: IndexMap::new(),
-            access: None,
-        }
+        Handler::new(self.tunnels.clone())
     }
 
     fn handle_session_error(&mut self, error: <Self::Handler as russh::server::Handler>::Error) {

src/terminal.rs

@@ -0,0 +1,68 @@
+use crossterm::{
+    execute,
+    terminal::{EnterAlternateScreen, LeaveAlternateScreen},
+};
+use russh::{ChannelId, server::Handle};
+use tokio::sync::mpsc::{UnboundedSender, unbounded_channel};
+use tracing::error;
+
+pub struct TerminalHandle {
+    sender: UnboundedSender<Vec<u8>>,
+    sink: Vec<u8>,
+}
+
+impl TerminalHandle {
+    pub async fn start(handle: Handle, channel_id: ChannelId) -> std::io::Result<Self> {
+        let (sender, mut receiver) = unbounded_channel::<Vec<u8>>();
+
+        tokio::spawn(async move {
+            while let Some(data) = receiver.recv().await {
+                let result = handle.data(channel_id, data.into()).await;
+
+                if let Err(e) = result {
+                    error!("Failed to send data: {e:?}");
+                };
+            }
+
+            if let Err(e) = handle.close(channel_id).await {
+                error!("Failed to close session: {e:?}");
+            }
+        });
+
+        let mut terminal_handle = Self {
+            sender,
+            sink: Vec::new(),
+        };
+
+        execute!(terminal_handle, EnterAlternateScreen)?;
+
+        Ok(terminal_handle)
+    }
+
+    pub fn leave_alternate_screen(&mut self) -> std::io::Result<()> {
+        execute!(self, LeaveAlternateScreen)
+    }
+}
+
+impl Drop for TerminalHandle {
+    fn drop(&mut self) {
+        self.leave_alternate_screen().ok();
+    }
+}
+
+impl std::io::Write for TerminalHandle {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.sink.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        let result = self.sender.send(self.sink.clone());
+        if let Err(e) = result {
+            return Err(std::io::Error::new(std::io::ErrorKind::BrokenPipe, e));
+        }
+
+        self.sink.clear();
+        Ok(())
+    }
+}

src/tui.rs

@@ -0,0 +1,119 @@
+use std::cmp;
+
+use futures::StreamExt;
+use indexmap::IndexMap;
+use ratatui::{
+    Frame,
+    layout::{Constraint, Flex, Rect},
+    style::{Color, Style, Stylize as _},
+    text::{Line, Span},
+    widgets::{Cell, HighlightSpacing, Row, Table, TableState},
+};
+
+use crate::tunnel::{self, Tunnel};
+
+pub struct Renderer {
+    table_state: TableState,
+    table_rows: Vec<Vec<Span<'static>>>,
+    table_header: Vec<Span<'static>>,
+    table_widths: Vec<Constraint>,
+}
+
+impl Default for Renderer {
+    fn default() -> Self {
+        let mut renderer = Self {
+            table_state: Default::default(),
+            table_rows: Default::default(),
+            table_header: tunnel::tui::header(),
+            table_widths: Default::default(),
+        };
+
+        renderer.update_widths();
+
+        renderer
+    }
+}
+
+impl Renderer {
+    // NOTE: This needs to be a separate function as the render functions can not be async
+    pub async fn update_table(&mut self, tunnels: &IndexMap<String, Option<Tunnel>>) {
+        self.table_rows = futures::stream::iter(tunnels.iter())
+            .then(tunnel::tui::to_row)
+            .collect::<Vec<_>>()
+            .await;
+
+        self.update_widths();
+    }
+
+    pub fn render(&mut self, frame: &mut Frame) {
+        self.render_title(frame, frame.area());
+
+        let area = frame.area().inner(ratatui::layout::Margin {
+            horizontal: 1,
+            vertical: 1,
+        });
+
+        self.render_table(frame, area);
+    }
+
+    pub fn render_title(&self, frame: &mut Frame, rect: Rect) {
+        let title = format!(
+            "{} ({})",
+            std::env!("CARGO_PKG_NAME"),
+            std::env!("CARGO_PKG_VERSION")
+        );
+        let title = Line::from(title).centered();
+        frame.render_widget(title, rect);
+    }
+
+    fn update_widths(&mut self) {
+        self.table_widths = std::iter::once(&self.table_header)
+            .chain(&self.table_rows)
+            .map(|row| row.iter().map(|cell| cell.width() as u16))
+            .fold(vec![0; self.table_header.len()], |acc, row| {
+                acc.into_iter()
+                    .zip(row)
+                    .map(|v| cmp::max(v.0, v.1))
+                    .collect()
+            })
+            .into_iter()
+            .map(|c| Constraint::Length(c + 1))
+            .collect();
+    }
+
+    pub fn render_table(&mut self, frame: &mut Frame<'_>, rect: Rect) {
+        let highlight_style = Style::default().bg(Color::Blue);
+        let header_style = Style::default().bold().reversed();
+        let row_style = Style::default();
+
+        let rows = self.table_rows.iter().map(|row| {
+            row.iter()
+                .cloned()
+                .map(Cell::from)
+                .collect::<Row>()
+                .style(row_style)
+                .height(1)
+        });
+
+        let header = self
+            .table_header
+            .iter()
+            .cloned()
+            .map(Cell::from)
+            .collect::<Row>()
+            .style(header_style)
+            .height(1);
+
+        let t = Table::default()
+            .header(header)
+            .rows(rows)
+            .flex(Flex::Start)
+            .column_spacing(3)
+            .widths(&self.table_widths)
+            .row_highlight_style(highlight_style)
+            .highlight_symbol(Line::from("> "))
+            .highlight_spacing(HighlightSpacing::Always);
+
+        frame.render_stateful_widget(t, rect, &mut self.table_state);
+    }
+}

src/tunnel.rs

@@ -24,6 +24,8 @@ use crate::{
     helper::response,
 };
 
+pub mod tui;
+
 #[derive(Debug, Clone)]
 pub enum TunnelAccess {
     Private(String),
@@ -76,7 +78,7 @@ impl Tunnels {
         }
     }
 
-    pub async fn add_tunnel(&mut self, address: &str, tunnel: Tunnel) -> Option<String> {
+    pub async fn add_tunnel(&mut self, address: &str, tunnel: Tunnel) -> (bool, String) {
         let mut all_tunnels = self.tunnels.write().await;
 
         let address = if address == "localhost" {
@@ -93,7 +95,7 @@ impl Tunnels {
         } else {
             let address = format!("{address}.{}", self.domain);
             if all_tunnels.contains_key(&address) {
-                return None;
+                return (false, address);
             }
             address
         };
@@ -101,7 +103,7 @@ impl Tunnels {
         trace!(tunnel = address, "Adding tunnel");
         all_tunnels.insert(address.clone(), tunnel);
 
-        Some(address)
+        (true, address)
     }
 
     pub async fn remove_tunnels(&mut self, tunnels: &IndexMap<String, Option<Tunnel>>) {

src/tunnel/tui.rs

@@ -0,0 +1,26 @@
+use std::ops::Deref;
+
+use ratatui::style::Stylize;
+use ratatui::text::Span;
+
+use super::{Tunnel, TunnelAccess};
+
+pub fn header() -> Vec<Span<'static>> {
+    vec!["Access".into(), "Port".into(), "Address".into()]
+}
+
+pub async fn to_row((address, tunnel): (&String, &Option<Tunnel>)) -> Vec<Span<'static>> {
+    let (access, port) = if let Some(tunnel) = tunnel {
+        let access = match tunnel.access.read().await.deref() {
+            TunnelAccess::Private(owner) => owner.clone().yellow(),
+            TunnelAccess::Public => "PUBLIC".green(),
+        };
+
+        (access, tunnel.port.to_string().into())
+    } else {
+        ("FAILED".red(), "".into())
+    };
+    let address = format!("http://{address}").into();
+
+    vec![access, port, address]
+}