Tunnels are now always stored in the handler

If tunnel.domain is set that means the tunnel is actually open.
This should make it possible to retry a failed tunnel in the future.
This commit is contained in:
Dreaded_X 2025-04-11 17:03:39 +02:00
parent da9dc7700c
commit 4fa885843f
Signed by: Dreaded_X
GPG Key ID: FA5F485356B0D2D4
6 changed files with 88 additions and 73 deletions

1
Cargo.lock generated
View File

@ -3221,7 +3221,6 @@ dependencies = [
"http-body-util",
"hyper",
"hyper-util",
"indexmap",
"rand 0.8.5",
"ratatui",
"reqwest",

View File

@ -15,7 +15,6 @@ 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"] }

View File

@ -1,7 +1,6 @@
use std::{io::Write, iter::once};
use clap::Parser as _;
use indexmap::IndexMap;
use ratatui::{Terminal, TerminalOptions, Viewport, layout::Rect, prelude::CrosstermBackend};
use russh::{
ChannelId,
@ -19,7 +18,7 @@ use crate::{
pub struct Handler {
all_tunnels: Tunnels,
tunnels: IndexMap<String, Option<Tunnel>>,
tunnels: Vec<Tunnel>,
user: Option<String>,
pty_channel: Option<ChannelId>,
@ -33,7 +32,7 @@ impl Handler {
pub fn new(all_tunnels: Tunnels) -> Self {
Self {
all_tunnels,
tunnels: IndexMap::new(),
tunnels: Default::default(),
user: None,
pty_channel: None,
terminal: None,
@ -43,12 +42,10 @@ impl Handler {
}
async fn set_access_all(&mut self, access: TunnelAccess) {
for (_address, tunnel) in &self.tunnels {
if let Some(tunnel) = tunnel {
for tunnel in &self.tunnels {
tunnel.set_access(access.clone()).await;
}
}
}
async fn resize(&mut self, width: u32, height: u32) -> std::io::Result<()> {
if let Some(terminal) = &mut self.terminal {
@ -92,7 +89,7 @@ impl Handler {
async fn set_access_selection(&mut self, access: TunnelAccess) {
if let Some(selected) = self.selected {
if let Some((_, Some(tunnel))) = self.tunnels.get_index_mut(selected) {
if let Some(tunnel) = self.tunnels.get_mut(selected) {
tunnel.set_access(access).await;
}
} else {
@ -267,18 +264,16 @@ impl russh::server::Handler for Handler {
return Err(russh::Error::Inconsistent);
};
let tunnel = Tunnel::new(
session.handle(),
address,
*port,
TunnelAccess::Private(user),
);
let (success, address) = self.all_tunnels.add_tunnel(address, tunnel.clone()).await;
let tunnel = self
.all_tunnels
.add_tunnel(session.handle(), address, *port, user)
.await;
let tunnel = if success { Some(tunnel) } else { None };
self.tunnels.insert(address, tunnel);
self.tunnels.push(tunnel);
Ok(success)
// Technically forwarding has failed if tunnel.domain = None, however by lying to the ssh
// client we can retry in the future
Ok(true)
}
async fn window_change_request(

View File

@ -1,7 +1,6 @@
use std::cmp;
use futures::StreamExt;
use indexmap::IndexMap;
use ratatui::{
Frame,
layout::{Constraint, Flex, Layout, Rect},
@ -25,12 +24,8 @@ fn command<'c>(key: &'c str, text: &'c str) -> Vec<Span<'c>> {
impl Renderer {
// NOTE: This needs to be a separate function as the render functions can not be async
pub async fn update(
&mut self,
tunnels: &IndexMap<String, Option<Tunnel>>,
index: Option<usize>,
) {
self.table_rows = futures::stream::iter(tunnels.iter())
pub async fn update(&mut self, tunnels: &[Tunnel], index: Option<usize>) {
self.table_rows = futures::stream::iter(tunnels)
.then(tunnel::tui::to_row)
.collect::<Vec<_>>()
.await;

View File

@ -8,8 +8,12 @@ use hyper::{
service::Service,
};
use hyper_util::rt::TokioIo;
use indexmap::IndexMap;
use std::{collections::HashMap, ops::Deref, pin::Pin, sync::Arc};
use std::{
collections::{HashMap, hash_map::Entry},
ops::Deref,
pin::Pin,
sync::Arc,
};
use tracing::{debug, error, trace, warn};
use russh::{
@ -37,20 +41,12 @@ pub enum TunnelAccess {
pub struct Tunnel {
handle: Handle,
name: String,
domain: Option<String>,
port: u32,
access: Arc<RwLock<TunnelAccess>>,
}
impl Tunnel {
pub fn new(handle: Handle, name: impl Into<String>, port: u32, access: TunnelAccess) -> Self {
Self {
handle,
name: name.into(),
port,
access: Arc::new(RwLock::new(access)),
}
}
pub async fn open_tunnel(&self) -> Result<Channel<Msg>, russh::Error> {
trace!(tunnel = self.name, "Opening tunnel");
self.handle
@ -65,6 +61,12 @@ impl Tunnel {
pub async fn is_public(&self) -> bool {
matches!(*self.access.read().await, TunnelAccess::Public)
}
pub fn get_address(&self) -> Option<String> {
self.domain
.clone()
.map(|domain| format!("{}.{domain}", self.name))
}
}
#[derive(Debug, Clone)]
@ -83,40 +85,57 @@ impl Tunnels {
}
}
pub async fn add_tunnel(&mut self, address: &str, tunnel: Tunnel) -> (bool, String) {
let mut all_tunnels = self.tunnels.write().await;
pub async fn add_tunnel(
&mut self,
handle: Handle,
name: impl Into<String>,
port: u32,
user: impl Into<String>,
) -> Tunnel {
let mut tunnel = Tunnel {
handle,
name: name.into(),
domain: Some(self.domain.clone()),
port,
access: Arc::new(RwLock::new(TunnelAccess::Private(user.into()))),
};
let address = if address == "localhost" {
if tunnel.name == "localhost" {
// NOTE: It is technically possible to become stuck in this loop.
// However, that really only becomes a concern if a (very) high
// number of tunnels is open at the same time.
loop {
let address = get_animal_name();
let address = format!("{address}.{}", self.domain);
if !all_tunnels.contains_key(&address) {
break address;
tunnel.name = get_animal_name().into();
if !self
.tunnels
.read()
.await
.contains_key(&tunnel.get_address().expect("domain is set"))
{
break;
}
trace!(tunnel = tunnel.name, "Already in use, picking new name");
}
} else {
let address = format!("{address}.{}", self.domain);
if all_tunnels.contains_key(&address) {
return (false, address);
}
address
};
let address = tunnel.get_address().expect("domain is set");
trace!(tunnel = address, "Adding tunnel");
all_tunnels.insert(address.clone(), tunnel);
(true, address)
if let Entry::Vacant(e) = self.tunnels.write().await.entry(address) {
trace!(tunnel = tunnel.name, "Adding tunnel");
e.insert(tunnel.clone());
} else {
trace!("Address already in use");
tunnel.domain = None
}
pub async fn remove_tunnels(&mut self, tunnels: &IndexMap<String, Option<Tunnel>>) {
tunnel
}
pub async fn remove_tunnels(&mut self, tunnels: &[Tunnel]) {
let mut all_tunnels = self.tunnels.write().await;
for (address, tunnel) in tunnels {
if tunnel.is_some() {
trace!(address, "Removing tunnel");
all_tunnels.remove(address);
for tunnel in tunnels {
if let Some(address) = tunnel.get_address() {
trace!(tunnel.name, "Removing tunnel");
all_tunnels.remove(&address);
}
}
}

View File

@ -6,22 +6,30 @@ use ratatui::text::Span;
use super::{Tunnel, TunnelAccess};
pub fn header() -> Vec<Span<'static>> {
vec!["Access".into(), "Port".into(), "Address".into()]
vec![
"Name".into(),
"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 {
pub async fn to_row(tunnel: &Tunnel) -> Vec<Span<'static>> {
let access = match tunnel.access.read().await.deref() {
TunnelAccess::Private(owner) => owner.clone().yellow(),
TunnelAccess::Protected => "PROTECTED".blue(),
TunnelAccess::Public => "PUBLIC".green(),
};
(access, tunnel.port.to_string().into())
} else {
("FAILED".red(), "".into())
};
let address = format!("http://{address}").into();
let address = tunnel
.get_address()
.map(|address| format!("http://{address}").into())
.unwrap_or("FAILED".red());
vec![access, port, address]
vec![
tunnel.name.clone().into(),
access,
tunnel.port.to_string().into(),
address,
]
}