From 19ec3714a680d8e539685c1f975e3629b1c6e0e5 Mon Sep 17 00:00:00 2001 From: Dreaded_X Date: Wed, 16 Apr 2025 02:55:42 +0200 Subject: [PATCH] Reorganized files --- src/cli.rs | 23 -------------- src/{ => helper}/animals.rs | 0 src/{ => helper}/animals.txt | 0 src/helper/mod.rs | 5 ++++ src/{ => helper}/units.rs | 0 src/{ => io}/input.rs | 0 src/io/mod.rs | 7 +++++ src/{wrapper.rs => io/stats.rs} | 45 ++++++++++++++++++++++++---- src/{io.rs => io/terminal_handle.rs} | 0 src/lib.rs | 23 +++----------- src/main.rs | 7 ++++- src/{ => ssh}/handler.rs | 35 +++++++++++++++++----- src/{server.rs => ssh/mod.rs} | 7 ++++- src/{tui.rs => ssh/renderer.rs} | 0 src/stats.rs | 36 ---------------------- src/tunnel/mod.rs | 18 +++++------ src/tunnel/registry.rs | 12 ++++---- src/{ => web}/auth.rs | 0 src/{web.rs => web/mod.rs} | 14 +++++---- src/{helper.rs => web/response.rs} | 0 20 files changed, 118 insertions(+), 114 deletions(-) delete mode 100644 src/cli.rs rename src/{ => helper}/animals.rs (100%) rename src/{ => helper}/animals.txt (100%) create mode 100644 src/helper/mod.rs rename src/{ => helper}/units.rs (100%) rename src/{ => io}/input.rs (100%) create mode 100644 src/io/mod.rs rename src/{wrapper.rs => io/stats.rs} (71%) rename src/{io.rs => io/terminal_handle.rs} (100%) rename src/{ => ssh}/handler.rs (95%) rename src/{server.rs => ssh/mod.rs} (93%) rename src/{tui.rs => ssh/renderer.rs} (100%) delete mode 100644 src/stats.rs rename src/{ => web}/auth.rs (100%) rename src/{web.rs => web/mod.rs} (96%) rename src/{helper.rs => web/response.rs} (100%) diff --git a/src/cli.rs b/src/cli.rs deleted file mode 100644 index f8396dd..0000000 --- a/src/cli.rs +++ /dev/null @@ -1,23 +0,0 @@ -use clap::Parser; - -/// Quickly create http tunnels for development -#[derive(Parser, Debug)] -#[command(version, about, long_about = None)] -pub struct Args { - /// Make all tunnels public by default instead of private - #[arg(long, group = "access")] - public: bool, - - #[arg(long, group = "access")] - protected: bool, -} - -impl Args { - pub fn make_public(&self) -> bool { - self.public - } - - pub fn make_protected(&self) -> bool { - self.protected - } -} diff --git a/src/animals.rs b/src/helper/animals.rs similarity index 100% rename from src/animals.rs rename to src/helper/animals.rs diff --git a/src/animals.txt b/src/helper/animals.txt similarity index 100% rename from src/animals.txt rename to src/helper/animals.txt diff --git a/src/helper/mod.rs b/src/helper/mod.rs new file mode 100644 index 0000000..8a5715e --- /dev/null +++ b/src/helper/mod.rs @@ -0,0 +1,5 @@ +mod animals; +mod units; + +pub use animals::get_animal_name; +pub use units::Unit; diff --git a/src/units.rs b/src/helper/units.rs similarity index 100% rename from src/units.rs rename to src/helper/units.rs diff --git a/src/input.rs b/src/io/input.rs similarity index 100% rename from src/input.rs rename to src/io/input.rs diff --git a/src/io/mod.rs b/src/io/mod.rs new file mode 100644 index 0000000..5f8a80f --- /dev/null +++ b/src/io/mod.rs @@ -0,0 +1,7 @@ +mod input; +mod stats; +mod terminal_handle; + +pub use input::Input; +pub use stats::{Stats, TrackStats}; +pub use terminal_handle::TerminalHandle; diff --git a/src/wrapper.rs b/src/io/stats.rs similarity index 71% rename from src/wrapper.rs rename to src/io/stats.rs index 4e1a162..fb19e06 100644 --- a/src/wrapper.rs +++ b/src/io/stats.rs @@ -7,23 +7,58 @@ use std::{ use pin_project_lite::pin_project; use russh::{ChannelStream, server::Msg}; -use crate::stats::Stats; +use crate::helper::Unit; + +use std::sync::atomic::{AtomicUsize, Ordering}; + +#[derive(Debug, Default)] +pub struct Stats { + connections: AtomicUsize, + rx: AtomicUsize, + tx: AtomicUsize, +} + +impl Stats { + pub fn add_connection(&self) { + self.connections.fetch_add(1, Ordering::Relaxed); + } + + pub fn add_rx_bytes(&self, n: usize) { + self.rx.fetch_add(n, Ordering::Relaxed); + } + + pub fn add_tx_bytes(&self, n: usize) { + self.tx.fetch_add(n, Ordering::Relaxed); + } + + pub fn connections(&self) -> usize { + self.connections.load(Ordering::Relaxed) + } + + pub fn rx(&self) -> Unit { + Unit::new(self.rx.load(Ordering::Relaxed), "B") + } + + pub fn tx(&self) -> Unit { + Unit::new(self.tx.load(Ordering::Relaxed), "B") + } +} pin_project! { - pub struct Wrapper { + pub struct TrackStats { #[pin] inner: ChannelStream, stats: Arc, } } -impl Wrapper { +impl TrackStats { pub fn new(inner: ChannelStream, stats: Arc) -> Self { Self { inner, stats } } } -impl hyper::rt::Read for Wrapper { +impl hyper::rt::Read for TrackStats { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, @@ -47,7 +82,7 @@ impl hyper::rt::Read for Wrapper { } } -impl hyper::rt::Write for Wrapper { +impl hyper::rt::Write for TrackStats { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, diff --git a/src/io.rs b/src/io/terminal_handle.rs similarity index 100% rename from src/io.rs rename to src/io/terminal_handle.rs diff --git a/src/lib.rs b/src/lib.rs index 149c98d..b14abe7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,22 +1,7 @@ #![feature(let_chains)] -mod animals; -pub mod auth; -mod cli; -mod handler; mod helper; -mod input; mod io; -mod ldap; -mod server; -mod stats; -mod tui; -mod tunnel; -mod units; -mod web; -mod wrapper; - -pub use ldap::Ldap; -pub use server::Server; -pub use tunnel::Registry; -pub use tunnel::Tunnel; -pub use web::Service; +pub mod ldap; +pub mod ssh; +pub mod tunnel; +pub mod web; diff --git a/src/main.rs b/src/main.rs index d5759cb..e9a917b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,7 +8,12 @@ use rand::rngs::OsRng; use tokio::net::TcpListener; use tracing::{error, info, warn}; use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; -use tunnel_rs::{Ldap, Registry, Server, Service, auth::ForwardAuth}; +use tunnel_rs::{ + ldap::Ldap, + ssh::Server, + tunnel::Registry, + web::{ForwardAuth, Service}, +}; #[tokio::main] async fn main() -> color_eyre::Result<()> { diff --git a/src/handler.rs b/src/ssh/handler.rs similarity index 95% rename from src/handler.rs rename to src/ssh/handler.rs index 6a34f0d..36b270e 100644 --- a/src/handler.rs +++ b/src/ssh/handler.rs @@ -1,6 +1,6 @@ use std::{cmp::min, io::Write, iter::once}; -use clap::Parser as _; +use clap::Parser; use ratatui::{Terminal, TerminalOptions, Viewport, layout::Rect, prelude::CrosstermBackend}; use russh::{ ChannelId, @@ -10,14 +10,33 @@ use russh::{ use tracing::{debug, trace, warn}; use crate::{ - Ldap, cli, - input::Input, - io::TerminalHandle, - ldap::LdapError, - tui::Renderer, + io::{Input, TerminalHandle}, + ldap::{Ldap, LdapError}, tunnel::{Registry, Tunnel, TunnelAccess}, }; +/// Quickly create http tunnels for development +#[derive(Parser, Debug)] +#[command(version, about, long_about = None)] +pub struct Args { + /// Make all tunnels public by default instead of private + #[arg(long, group = "access")] + public: bool, + + #[arg(long, group = "access")] + protected: bool, +} + +impl Args { + pub fn make_public(&self) -> bool { + self.public + } + + pub fn make_protected(&self) -> bool { + self.protected + } +} + #[derive(Debug, thiserror::Error)] pub enum HandlerError { #[error(transparent)] @@ -38,7 +57,7 @@ pub struct Handler { pty_channel: Option, terminal: Option>>, - renderer: Renderer, + renderer: super::Renderer, selected: Option, rename_buffer: Option, @@ -315,7 +334,7 @@ impl russh::server::Handler for Handler { trace!(?cmd, "exec_request"); let cmd = once(" --").chain(cmd.split_whitespace()); - match cli::Args::try_parse_from(cmd) { + match Args::try_parse_from(cmd) { Ok(args) => { debug!("{args:?}"); if args.make_public() { diff --git a/src/server.rs b/src/ssh/mod.rs similarity index 93% rename from src/server.rs rename to src/ssh/mod.rs index c5deaca..58c7216 100644 --- a/src/server.rs +++ b/src/ssh/mod.rs @@ -1,10 +1,15 @@ +mod handler; +mod renderer; + use std::{net::SocketAddr, sync::Arc, time::Duration}; use russh::{MethodKind, keys::PrivateKey, server::Server as _}; use tokio::net::ToSocketAddrs; use tracing::{debug, warn}; -use crate::{Ldap, handler::Handler, tunnel::Registry}; +use crate::{ldap::Ldap, tunnel::Registry}; +use handler::Handler; +use renderer::Renderer; pub struct Server { ldap: Ldap, diff --git a/src/tui.rs b/src/ssh/renderer.rs similarity index 100% rename from src/tui.rs rename to src/ssh/renderer.rs diff --git a/src/stats.rs b/src/stats.rs deleted file mode 100644 index 3f7a265..0000000 --- a/src/stats.rs +++ /dev/null @@ -1,36 +0,0 @@ -use std::sync::atomic::{AtomicUsize, Ordering}; - -use crate::units::Unit; - -#[derive(Debug, Default)] -pub struct Stats { - connections: AtomicUsize, - rx: AtomicUsize, - tx: AtomicUsize, -} - -impl Stats { - pub fn add_connection(&self) { - self.connections.fetch_add(1, Ordering::Relaxed); - } - - pub fn add_rx_bytes(&self, n: usize) { - self.rx.fetch_add(n, Ordering::Relaxed); - } - - pub fn add_tx_bytes(&self, n: usize) { - self.tx.fetch_add(n, Ordering::Relaxed); - } - - pub fn connections(&self) -> usize { - self.connections.load(Ordering::Relaxed) - } - - pub fn rx(&self) -> Unit { - Unit::new(self.rx.load(Ordering::Relaxed), "B") - } - - pub fn tx(&self) -> Unit { - Unit::new(self.tx.load(Ordering::Relaxed), "B") - } -} diff --git a/src/tunnel/mod.rs b/src/tunnel/mod.rs index abdd835..477fdd1 100644 --- a/src/tunnel/mod.rs +++ b/src/tunnel/mod.rs @@ -8,19 +8,19 @@ use tracing::trace; use russh::server::Handle; use tokio::sync::{RwLock, RwLockReadGuard}; -use crate::{stats::Stats, wrapper::Wrapper}; - pub use registry::Registry; +use crate::io::{Stats, TrackStats}; + #[derive(Debug, Clone)] -pub enum TunnelAccess { +pub(crate) enum TunnelAccess { Private(String), Protected, Public, } #[derive(Debug, Clone)] -pub struct TunnelInner { +pub(crate) struct TunnelInner { handle: Handle, internal_address: String, port: u32, @@ -29,7 +29,7 @@ pub struct TunnelInner { } impl TunnelInner { - pub async fn open(&self) -> Result { + pub(crate) async fn open(&self) -> Result { trace!("Opening tunnel"); self.stats.add_connection(); let channel = self @@ -42,20 +42,20 @@ impl TunnelInner { ) .await?; - Ok(Wrapper::new(channel.into_stream(), self.stats.clone())) + Ok(TrackStats::new(channel.into_stream(), self.stats.clone())) } - pub async fn is_public(&self) -> bool { + pub(crate) async fn is_public(&self) -> bool { matches!(*self.access.read().await, TunnelAccess::Public) } - pub async fn get_access(&self) -> RwLockReadGuard<'_, TunnelAccess> { + pub(crate) async fn get_access(&self) -> RwLockReadGuard<'_, TunnelAccess> { self.access.read().await } } #[derive(Debug)] -pub struct Tunnel { +pub(crate) struct Tunnel { inner: TunnelInner, registry: Registry, diff --git a/src/tunnel/registry.rs b/src/tunnel/registry.rs index f7f8f01..f053bae 100644 --- a/src/tunnel/registry.rs +++ b/src/tunnel/registry.rs @@ -6,19 +6,19 @@ use std::{ use tokio::sync::RwLock; use tracing::trace; -use crate::{Tunnel, animals::get_animal_name}; +use crate::{helper::get_animal_name, tunnel::Tunnel}; use super::TunnelInner; #[derive(Debug)] -pub struct RegistryEntry { +pub(crate) struct RegistryEntry { registry: Registry, name: String, address: Option, } impl RegistryEntry { - pub fn new(registry: Registry) -> Self { + pub(crate) fn new(registry: Registry) -> Self { Self { registry, name: Default::default(), @@ -26,11 +26,11 @@ impl RegistryEntry { } } - pub fn get_address(&self) -> Option<&String> { + pub(crate) fn get_address(&self) -> Option<&String> { self.address.as_ref() } - pub fn get_name(&self) -> &str { + pub(crate) fn get_name(&self) -> &str { &self.name } } @@ -124,7 +124,7 @@ impl Registry { self.register(tunnel).await; } - pub async fn get(&self, address: &str) -> Option { + pub(crate) async fn get(&self, address: &str) -> Option { self.tunnels.read().await.get(address).cloned() } } diff --git a/src/auth.rs b/src/web/auth.rs similarity index 100% rename from src/auth.rs rename to src/web/auth.rs diff --git a/src/web.rs b/src/web/mod.rs similarity index 96% rename from src/web.rs rename to src/web/mod.rs index 2cbed38..6718c57 100644 --- a/src/web.rs +++ b/src/web/mod.rs @@ -1,4 +1,7 @@ -use crate::Registry; +mod auth; +mod response; + +use crate::tunnel::Registry; use std::{ops::Deref, pin::Pin}; use bytes::Bytes; @@ -11,11 +14,10 @@ use hyper::{ }; use tracing::{debug, error, trace, warn}; -use crate::{ - auth::{AuthStatus, ForwardAuth}, - helper::response, - tunnel::TunnelAccess, -}; +use crate::tunnel::TunnelAccess; +use auth::AuthStatus; +pub use auth::ForwardAuth; +use response::response; #[derive(Debug, Clone)] pub struct Service { diff --git a/src/helper.rs b/src/web/response.rs similarity index 100% rename from src/helper.rs rename to src/web/response.rs