From 750713b6b0c03fdf4b3c503d70ebee4ce4d14ae5 Mon Sep 17 00:00:00 2001 From: Dreaded_X Date: Thu, 10 Apr 2025 03:29:04 +0200 Subject: [PATCH] Implemented initial user interface --- .typos.toml | 2 + Cargo.lock | 351 +++++++++++++++++++++++++++++++++++++++++++++- Cargo.toml | 4 + src/lib.rs | 2 + src/ssh.rs | 229 ++++++++++++++++++++---------- src/terminal.rs | 68 +++++++++ src/tui.rs | 119 ++++++++++++++++ src/tunnel.rs | 8 +- src/tunnel/tui.rs | 26 ++++ 9 files changed, 729 insertions(+), 80 deletions(-) create mode 100644 .typos.toml create mode 100644 src/terminal.rs create mode 100644 src/tui.rs create mode 100644 src/tunnel/tui.rs diff --git a/.typos.toml b/.typos.toml new file mode 100644 index 0000000..372dc95 --- /dev/null +++ b/.typos.toml @@ -0,0 +1,2 @@ +[default.extend-words] +ratatui = "ratatui" diff --git a/Cargo.lock b/Cargo.lock index fd58291..ae81346 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -67,6 +67,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "android-tzdata" version = "0.1.1" @@ -261,6 +267,21 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +[[package]] +name = "cassowary" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df8670b8c7b9dae1793364eafadf7239c40d669904660c5960d74cfd80b46a53" + +[[package]] +name = "castaway" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0abae9be0aaf9ea96a3b1b8b1b55c602ca751eba1b1500220cea4ecbafe7c0d5" +dependencies = [ + "rustversion", +] + [[package]] name = "cbc" version = "0.1.2" @@ -414,12 +435,35 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990" +[[package]] +name = "compact_str" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b79c4069c6cad78e2e0cdfcbd26275770669fb39fd308a752dc110e83b9af32" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "static_assertions", +] + [[package]] name = "const-oid" version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" +[[package]] +name = "convert_case" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb402b8d4c85569410425650ce3eddc7d698ed96d39a73f941b08fb63082f1e7" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -454,6 +498,49 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossterm" +version = "0.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6" +dependencies = [ + "bitflags", + "crossterm_winapi", + "mio", + "parking_lot", + "rustix 0.38.44", + "signal-hook", + "signal-hook-mio", + "winapi", +] + +[[package]] +name = "crossterm" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" +dependencies = [ + "bitflags", + "crossterm_winapi", + "derive_more", + "document-features", + "mio", + "parking_lot", + "rustix 1.0.5", + "signal-hook", + "signal-hook-mio", + "winapi", +] + +[[package]] +name = "crossterm_winapi" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acdd7c62a3665c7f6830a51635d9ac9b23ed385797f70a83bb8bafe9c572ab2b" +dependencies = [ + "winapi", +] + [[package]] name = "crypto-bigint" version = "0.5.5" @@ -513,6 +600,41 @@ dependencies = [ "syn", ] +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn", +] + [[package]] name = "data-encoding" version = "2.8.0" @@ -541,6 +663,27 @@ dependencies = [ "zeroize", ] +[[package]] +name = "derive_more" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "093242cf7570c207c83073cf82f79706fe7b8317e98620a47d5be7c3d8497678" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bda628edc44c4bb645fbe0f758797143e4e07926f7ebf4e9bdfbd3d2ce621df3" +dependencies = [ + "convert_case", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "digest" version = "0.10.7" @@ -564,6 +707,15 @@ dependencies = [ "syn", ] +[[package]] +name = "document-features" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95249b50c6c185bee49034bcb378a49dc2b5dff0be90ff6616d31d64febab05d" +dependencies = [ + "litrs", +] + [[package]] name = "dotenvy" version = "0.15.7" @@ -609,6 +761,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + [[package]] name = "elliptic-curve" version = "0.13.8" @@ -715,6 +873,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "foreign-types" version = "0.3.2" @@ -917,6 +1081,11 @@ name = "hashbrown" version = "0.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] [[package]] name = "heck" @@ -1233,6 +1402,12 @@ dependencies = [ "syn", ] +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "idna" version = "1.0.3" @@ -1270,6 +1445,12 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "indoc" +version = "2.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd" + [[package]] name = "inout" version = "0.1.4" @@ -1280,6 +1461,19 @@ dependencies = [ "generic-array", ] +[[package]] +name = "instability" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf9fed6d91cfb734e7476a06bde8300a1b94e217e1b523b6f0cd1a01998c71d" +dependencies = [ + "darling", + "indoc", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "internal-russh-forked-ssh-key" version = "0.6.10+upstream-0.6.7" @@ -1331,6 +1525,15 @@ version = "1.70.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" @@ -1368,6 +1571,12 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + [[package]] name = "linux-raw-sys" version = "0.9.3" @@ -1380,6 +1589,12 @@ version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "23fb14cb19457329c82206317a5663005a4d404783dc74f4252769b0d5f42856" +[[package]] +name = "litrs" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ce301924b7887e9d637144fdade93f9dfff9b60981d4ac161db09720d39aa5" + [[package]] name = "lock_api" version = "0.4.12" @@ -1396,6 +1611,15 @@ version = "0.4.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" +[[package]] +name = "lru" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +dependencies = [ + "hashbrown", +] + [[package]] name = "matchers" version = "0.1.0" @@ -1448,6 +1672,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2886843bf800fba2e3377cff24abf6379b4c4d5c6681eaf9ea5b0d15090450bd" dependencies = [ "libc", + "log", "wasi 0.11.0+wasi-snapshot-preview1", "windows-sys 0.52.0", ] @@ -1714,6 +1939,12 @@ dependencies = [ "subtle", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pbkdf2" version = "0.12.2" @@ -1974,6 +2205,27 @@ dependencies = [ "getrandom 0.3.2", ] +[[package]] +name = "ratatui" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eabd94c2f37801c20583fc49dd5cd6b0ba68c716787c2dd6ed18571e1e63117b" +dependencies = [ + "bitflags", + "cassowary", + "compact_str", + "crossterm 0.28.1", + "indoc", + "instability", + "itertools", + "lru", + "paste", + "strum", + "unicode-segmentation", + "unicode-truncate", + "unicode-width 0.2.0", +] + [[package]] name = "redox_syscall" version = "0.5.10" @@ -2231,6 +2483,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys 0.4.15", + "windows-sys 0.59.0", +] + [[package]] name = "rustix" version = "1.0.5" @@ -2240,7 +2505,7 @@ dependencies = [ "bitflags", "errno", "libc", - "linux-raw-sys", + "linux-raw-sys 0.9.3", "windows-sys 0.59.0", ] @@ -2467,6 +2732,27 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "signal-hook" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8621587d4798caf8eb44879d42e56b9a93ea5dcd315a6487c357130095b62801" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-mio" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34db1a06d485c9142248b7a054f034b349b212551f3dfd19c94d45a754a217cd" +dependencies = [ + "libc", + "mio", + "signal-hook", +] + [[package]] name = "signal-hook-registry" version = "1.4.2" @@ -2562,12 +2848,40 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "strsim" version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn", +] + [[package]] name = "subtle" version = "2.6.1" @@ -2635,7 +2949,7 @@ dependencies = [ "fastrand", "getrandom 0.3.2", "once_cell", - "rustix", + "rustix 1.0.5", "windows-sys 0.59.0", ] @@ -2901,18 +3215,22 @@ dependencies = [ "clap", "clio", "color-eyre", + "crossterm 0.29.0", "dotenvy", + "futures", "http-body-util", "hyper", "hyper-util", "indexmap", "rand 0.8.5", + "ratatui", "reqwest", "russh", "thiserror 2.0.12", "tokio", "tracing", "tracing-subscriber", + "unicode-width 0.2.0", ] [[package]] @@ -2927,6 +3245,35 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" +[[package]] +name = "unicode-segmentation" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + +[[package]] +name = "unicode-truncate" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3644627a5af5fa321c95b9b235a72fd24cd29c648c2c379431e6628655627bf" +dependencies = [ + "itertools", + "unicode-segmentation", + "unicode-width 0.1.14", +] + +[[package]] +name = "unicode-width" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" + +[[package]] +name = "unicode-width" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd" + [[package]] name = "universal-hash" version = "0.5.1" diff --git a/Cargo.toml b/Cargo.toml index a0c2162..d5a1f1b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,15 +9,19 @@ bytes = "1.10.1" clap = { version = "4.5.35", features = ["derive"] } clio = { version = "0.3.5", features = ["clap-parse"] } color-eyre = "0.6.3" +crossterm = "0.29.0" dotenvy = "0.15.7" +futures = "0.3.31" http-body-util = { version = "0.1.3", features = ["full"] } hyper = { version = "1.6.0", features = ["full"] } hyper-util = { version = "0.1.11", features = ["full"] } indexmap = "2.9.0" rand = "0.8.5" +ratatui = { version = "0.29.0", features = ["unstable-backend-writer"] } reqwest = { version = "0.12.15", features = ["rustls-tls"] } russh = "0.51.1" thiserror = "2.0.12" tokio = { version = "1.44.1", features = ["full"] } tracing = "0.1.41" tracing-subscriber = { version = "0.3.19", features = ["json", "env-filter"] } +unicode-width = "0.2.0" diff --git a/src/lib.rs b/src/lib.rs index 76eb4ee..3427ef4 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/ssh.rs b/src/ssh.rs index eeadcdc..868707c 100644 --- a/src/ssh.rs +++ b/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>, - rx: Option>>, - all_tunnels: Tunnels, tunnels: IndexMap>, - access: Option, + user: Option, + + terminal: Option>>, + renderer: Renderer, } impl Handler { - fn send(&self, data: impl AsRef) { - let _ = self.tx.send(data.as_ref().as_bytes().to_vec()); - } - - fn sendln(&self, data: impl AsRef) { - 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 { + 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, - _session: &mut Session, + session: &mut Session, ) -> Result { 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 { 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("").chain(cmd.split_whitespace()); + let cmd = once(" --").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 { 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) -> Self::Handler { - let (tx, rx) = unbounded_channel::>(); - - 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: ::Error) { diff --git a/src/terminal.rs b/src/terminal.rs new file mode 100644 index 0000000..0bc83b0 --- /dev/null +++ b/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>, + sink: Vec, +} + +impl TerminalHandle { + pub async fn start(handle: Handle, channel_id: ChannelId) -> std::io::Result { + let (sender, mut receiver) = unbounded_channel::>(); + + 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 { + 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(()) + } +} diff --git a/src/tui.rs b/src/tui.rs new file mode 100644 index 0000000..bb9af63 --- /dev/null +++ b/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>>, + table_header: Vec>, + table_widths: Vec, +} + +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>) { + self.table_rows = futures::stream::iter(tunnels.iter()) + .then(tunnel::tui::to_row) + .collect::>() + .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::() + .style(row_style) + .height(1) + }); + + let header = self + .table_header + .iter() + .cloned() + .map(Cell::from) + .collect::() + .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); + } +} diff --git a/src/tunnel.rs b/src/tunnel.rs index f19b7c7..320c5bb 100644 --- a/src/tunnel.rs +++ b/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 { + 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>) { diff --git a/src/tunnel/tui.rs b/src/tunnel/tui.rs new file mode 100644 index 0000000..c001c5e --- /dev/null +++ b/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> { + vec!["Access".into(), "Port".into(), "Address".into()] +} + +pub async fn to_row((address, tunnel): (&String, &Option)) -> Vec> { + 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] +}