diff --git a/client/src/control/mod.rs b/client/src/control/mod.rs index 22c881e..ba910a9 100644 --- a/client/src/control/mod.rs +++ b/client/src/control/mod.rs @@ -1,3 +1,3 @@ mod session; -pub use session::{ActiveSession, Control, SessionEnd}; +pub use session::{ActiveSession, Control, LoginRejected, SessionEnd}; diff --git a/client/src/control/session.rs b/client/src/control/session.rs index f990923..ec1c9c9 100644 --- a/client/src/control/session.rs +++ b/client/src/control/session.rs @@ -1,6 +1,5 @@ use crate::connector::{build_connector, Connector}; use crate::reload::{ReloadLevel, ReloadOutcome, TunnelChanges}; -use crate::session_id; use crate::tunnel::TunnelManager; use anyhow::{anyhow, Result}; use orbien_core::auth; @@ -9,7 +8,7 @@ use orbien_core::msg::{self, CloseTunnel, Login, Message, NewDataConn, NewTunnel use orbien_core::transport::DynStream; use orbien_core::VERSION; use std::collections::{HashMap, HashSet}; -use std::path::Path; +use std::fmt; use std::sync::atomic::{AtomicI64, Ordering}; use std::sync::{Arc, Mutex as StdMutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -19,6 +18,19 @@ use tokio::task::JoinSet; use tokio::time::{interval, sleep}; use tokio_util::sync::CancellationToken; +#[derive(Debug, Clone)] +pub struct LoginRejected { + pub reason: String, +} + +impl fmt::Display for LoginRejected { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "login rejected: {}", self.reason) + } +} + +impl std::error::Error for LoginRejected {} + #[derive(Debug)] pub enum SessionEnd { Disconnected { session_id: String }, @@ -73,7 +85,6 @@ impl Control { pub async fn open_session( cfg: Arc>, previous_session_id: String, - config_path: &Path, parent_cancel: CancellationToken, on_connected: impl FnOnce(), on_tunnel_remote: OnTunnelRemote, @@ -98,6 +109,7 @@ impl Control { os: std::env::consts::OS.into(), arch: std::env::consts::ARCH.into(), user: cfg_snapshot.user.clone(), + agent_id: cfg_snapshot.agent_id.clone(), auth_digest, timestamp, session_id: previous_session_id, @@ -108,6 +120,7 @@ impl Control { os = %login.os, arch = %login.arch, user = %login.user, + agent_id = %login.agent_id, "login identity" ); @@ -123,13 +136,10 @@ impl Control { }; if !resp.error.is_empty() { - return Err(anyhow!("login failed: {}", resp.error)); + return Err(LoginRejected { reason: resp.error }.into()); } tracing::info!(session_id = %resp.session_id, "login ok"); - if let Err(e) = session_id::save(config_path, &resp.session_id) { - tracing::warn!(error = %e, "failed to persist session_id"); - } let (reader, writer) = tokio::io::split(stream); let ctl = Arc::new(Control { @@ -393,7 +403,7 @@ impl Control { match msg { Message::KickOut(k) => { - tracing::warn!(reason = %k.reason, "kicked by server — will exit"); + tracing::warn!(reason = %k.reason, "kicked by server"); return Ok(ReaderEnd::Kicked(k.reason)); } Message::ReqDataConn(_) => { diff --git a/client/src/handle.rs b/client/src/handle.rs index 3a06e9b..5a8d06b 100644 --- a/client/src/handle.rs +++ b/client/src/handle.rs @@ -313,7 +313,7 @@ impl ClientHandle { Arc::new(move || h.clear_tunnel_remotes()) }; - let result = Service::new(cfg, config_path.clone()) + let result = Service::new(cfg) .run( cancel.clone(), &mut reload_rx, @@ -400,7 +400,7 @@ impl ClientHandle { Ok(Ok(())) => {} Ok(Err(e)) => tracing::warn!(error = %e, "client task join error"), Err(_) => { - tracing::warn!("client stop timed out after 5s — aborting task"); + tracing::warn!("client stop timed out after 5s, aborting task"); abort.abort(); self.set_status(ClientStatus::Stopped); } diff --git a/client/src/lib.rs b/client/src/lib.rs index e21635b..8532bff 100644 --- a/client/src/lib.rs +++ b/client/src/lib.rs @@ -5,7 +5,6 @@ pub mod local_control; mod plugin; mod reload; mod service; -mod session_id; mod tunnel; pub use handle::{ClientHandle, ClientStatus, StartOptions}; diff --git a/client/src/service.rs b/client/src/service.rs index b854c85..001e15c 100644 --- a/client/src/service.rs +++ b/client/src/service.rs @@ -1,9 +1,8 @@ -use crate::control::{ActiveSession, Control, SessionEnd}; +use crate::control::{ActiveSession, Control, LoginRejected, SessionEnd}; use crate::handle::ClientStatus; use crate::reload::{ empty_outcome, outcome_from_plan, outcome_level, plan_reload, ReloadOutcome, ReloadPlan, }; -use crate::session_id; use anyhow::Result; use orbien_core::config::ClientConfig; use std::path::{Path, PathBuf}; @@ -29,15 +28,11 @@ struct PendingReloadReply { pub struct Service { cfg: ClientConfig, - config_path: PathBuf, } impl Service { - pub fn new(cfg: ClientConfig, config_path: impl Into) -> Self { - Self { - cfg, - config_path: config_path.into(), - } + pub fn new(cfg: ClientConfig) -> Self { + Self { cfg } } pub async fn run( @@ -55,12 +50,8 @@ impl Service { let mut guard = cfg.write().await; *guard = self.cfg; } - let config_path = self.config_path; - let mut session_id = session_id::load(&config_path); - if !session_id.is_empty() { - tracing::info!(%session_id, "restored persisted session_id"); - } + let mut session_id = String::new(); let mut first_attempt = true; let mut backoff_secs = RECONNECT_BASE_SECS; @@ -99,7 +90,6 @@ impl Service { res = Control::open_session( Arc::clone(&cfg), session_id.clone(), - &config_path, connect_cancel, || { on_status(ClientStatus::Running); @@ -120,7 +110,15 @@ impl Service { ); return Ok(()); } + if let Some(rej) = e.downcast_ref::() { + fail_pending_reload(&mut pending_reload_reply, &rej.to_string()); + tracing::error!(reason = %rej.reason, "login rejected, stopping"); + on_log(format!("ERROR {rej}")); + on_status(ClientStatus::Stopped); + return Err(e); + } on_log(format!("ERROR failed to connect: {e}")); + tracing::warn!(error = %e, "connect failed, retrying"); on_status(ClientStatus::Reconnecting); first_attempt = false; let delay = backoff_secs; @@ -173,7 +171,7 @@ impl Service { tracing::warn!( session_id = %rid, %reason, - "kicked by server — stopping (no reconnect)" + "kicked by server, stopping" ); on_log(format!("WARN kicked by server: {reason}")); return Ok(()); diff --git a/client/src/session_id.rs b/client/src/session_id.rs deleted file mode 100644 index f6ce7bb..0000000 --- a/client/src/session_id.rs +++ /dev/null @@ -1,33 +0,0 @@ -use std::path::{Path, PathBuf}; - -pub fn path_for(config_path: &Path) -> PathBuf { - let mut p = config_path.to_path_buf(); - let ext = p - .extension() - .and_then(|e| e.to_str()) - .map(|e| format!("{e}.session_id")) - .unwrap_or_else(|| "session_id".into()); - p.set_extension(ext); - p -} - -fn read_valid_id(path: &Path) -> Option { - std::fs::read_to_string(path) - .ok() - .map(|s| s.trim().to_string()) - .filter(|s| { - !s.is_empty() && s.len() <= 64 && s.chars().all(|c| c.is_ascii_hexdigit() || c == '-') - }) -} - -pub fn load(config_path: &Path) -> String { - read_valid_id(&path_for(config_path)).unwrap_or_default() -} - -pub fn save(config_path: &Path, session_id: &str) -> std::io::Result<()> { - let path = path_for(config_path); - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent)?; - } - std::fs::write(path, session_id) -} diff --git a/core/src/config/client.rs b/core/src/config/client.rs index 647bc1c..3ec7dca 100644 --- a/core/src/config/client.rs +++ b/core/src/config/client.rs @@ -11,6 +11,9 @@ pub struct ClientConfig { #[serde(default)] pub user: String, + #[serde(default, rename = "agentId", alias = "agent_id")] + pub agent_id: String, + #[serde(default)] pub auth: AuthConfig, @@ -561,6 +564,7 @@ impl ClientConfig { pub fn connection_settings_eq(&self, other: &Self) -> bool { self.server == other.server && self.user == other.user + && self.agent_id == other.agent_id && self.auth == other.auth && self.transport == other.transport && self.udp_packet_size == other.udp_packet_size diff --git a/core/src/msg/types.rs b/core/src/msg/types.rs index 63f099b..1e2f368 100644 --- a/core/src/msg/types.rs +++ b/core/src/msg/types.rs @@ -27,6 +27,8 @@ pub struct Login { #[serde(default)] pub user: String, #[serde(default)] + pub agent_id: String, + #[serde(default)] pub auth_digest: String, #[serde(default)] pub timestamp: i64, diff --git a/desktop/src/config_bridge.rs b/desktop/src/config_bridge.rs index 4ba68bc..1d40747 100644 --- a/desktop/src/config_bridge.rs +++ b/desktop/src/config_bridge.rs @@ -123,6 +123,7 @@ pub fn load_merge_server_fields( ClientConfig { server: String::new(), user: String::new(), + agent_id: String::new(), auth: Default::default(), transport: TransportConfig::default(), tunnels: Vec::new(), @@ -220,6 +221,7 @@ fn build_base( let mut cfg = ClientConfig { server, user: user.trim().into(), + agent_id: String::new(), auth: Default::default(), transport, tunnels, diff --git a/server-ui/src/types/api.ts b/server-ui/src/types/api.ts index 5a17001..ae41a4d 100644 --- a/server-ui/src/types/api.ts +++ b/server-ui/src/types/api.ts @@ -41,6 +41,7 @@ export interface SystemStatus { export interface ClientInfo { sessionId: string + agentId?: string user: string hostname: string os: string diff --git a/server-ui/src/views/ClientDetail.vue b/server-ui/src/views/ClientDetail.vue index 51379f7..652b25d 100644 --- a/server-ui/src/views/ClientDetail.vue +++ b/server-ui/src/views/ClientDetail.vue @@ -195,7 +195,7 @@ onUnmounted(() => { {{ t('nav.clients') }} - {{ client?.sessionId || sessionId }} + {{ client?.agentId || client?.sessionId || sessionId }}
{{ t('traffic.loading') }}
@@ -209,11 +209,12 @@ onUnmounted(() => {
-

{{ client.sessionId }}

+

{{ client.agentId || client.sessionId }}

v{{ client.version }} {{ client.user }}
+ session {{ client.sessionId }} {{ client.clientIP }}
diff --git a/server-ui/src/views/Clients.vue b/server-ui/src/views/Clients.vue index c1f78c2..7bd7dc2 100644 --- a/server-ui/src/views/Clients.vue +++ b/server-ui/src/views/Clients.vue @@ -137,7 +137,8 @@ async function onKick(sessionId: string, evt: Event) {
-

{{ c.sessionId }}

+

{{ c.agentId || c.sessionId }}

+ {{ c.sessionId }} {{ c.hostname }} {{ c.user }} v{{ c.version }} diff --git a/server/src/control/session/mod.rs b/server/src/control/session/mod.rs index 12dcc79..650ba48 100644 --- a/server/src/control/session/mod.rs +++ b/server/src/control/session/mod.rs @@ -3,7 +3,9 @@ mod register; use crate::access::AccessPolicy; use crate::metrics::{MemMetrics, ServerMetrics}; -use crate::tunnel::{HttpGw, HttpsGw, TunnelManager}; +use crate::tunnel::{ + DetachedTunnel, HttpGw, HttpsGw, PortTable, TunnelManager, TunnelOwner, TunnelRegistry, +}; use anyhow::Result; use orbien_core::config::ServerConfig; use orbien_core::msg::{self, KickOut, Message, Ping, Pong}; @@ -13,7 +15,7 @@ use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::io::AsyncWriteExt; use tokio::io::{ReadHalf, WriteHalf}; -use tokio::sync::{mpsc, Mutex, Notify}; +use tokio::sync::{mpsc, watch, Mutex, Notify}; use tokio::task::JoinSet; use tokio::time::sleep; @@ -22,7 +24,9 @@ type CtrlWrite = WriteHalf; pub struct Control { pub session_id: String, + pub generation: u64, pub user: String, + pub agent_id: String, pub hostname: String, pub os: String, pub arch: String, @@ -37,8 +41,13 @@ pub struct Control { data_notify: Notify, shutdown_notify: Notify, tunnels: Mutex, + tunnel_registry: Arc, + tcp_ports: Arc, + udp_ports: Arc, bg_tasks: Mutex>, closed: AtomicBool, + finished: watch::Sender, + activated: AtomicBool, pool_count: usize, http_gw: Option>, https_gw: Option>, @@ -50,6 +59,7 @@ pub struct Control { impl Control { pub fn new( session_id: String, + generation: u64, stream: DynStream, cfg: ServerConfig, pool_count: usize, @@ -57,18 +67,25 @@ impl Control { https_gw: Option>, access: Arc, user: String, + agent_id: String, hostname: String, os: String, arch: String, version: String, client_ip: String, metrics: Arc, + tunnel_registry: Arc, + tcp_ports: Arc, + udp_ports: Arc, ) -> Self { let (reader, writer) = tokio::io::split(stream); + let (finished, _) = watch::channel(false); let (data_tx, data_rx) = mpsc::channel(64); Self { session_id, + generation, user, + agent_id, hostname, os, arch, @@ -83,8 +100,13 @@ impl Control { data_notify: Notify::new(), shutdown_notify: Notify::new(), tunnels: Mutex::new(TunnelManager::new()), + tunnel_registry, + tcp_ports, + udp_ports, bg_tasks: Mutex::new(JoinSet::new()), closed: AtomicBool::new(false), + finished, + activated: AtomicBool::new(false), pool_count: pool_count.max(1), http_gw, https_gw, @@ -99,6 +121,21 @@ impl Control { } } + pub fn owner(&self) -> TunnelOwner { + TunnelOwner { + session_id: self.session_id.clone(), + generation: self.generation, + } + } + + pub fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } + + pub fn is_accepting_data(&self) -> bool { + self.activated.load(Ordering::Acquire) && !self.is_closed() + } + pub async fn tunnel_summaries(&self) -> Vec { self.tunnels.lock().await.summaries() } @@ -107,6 +144,73 @@ impl Control { self.tunnels.lock().await.len() } + pub async fn wait_finished(&self) { + let mut rx = self.finished.subscribe(); + loop { + if *rx.borrow_and_update() { + return; + } + if rx.changed().await.is_err() { + return; + } + } + } + + fn mark_finished(&self) { + let _ = self.finished.send(true); + } + + pub async fn send_login_ok(&self, version: &str) -> Result<()> { + let mut writer = self.writer.lock().await; + msg::write_msg( + &mut *writer, + &Message::LoginResp(orbien_core::msg::LoginResp { + version: version.into(), + session_id: self.session_id.clone(), + error: String::new(), + }), + ) + .await?; + drop(writer); + self.activated.store(true, Ordering::Release); + Ok(()) + } + + pub async fn send_login_err(&self, version: &str, error: &str) -> Result<()> { + let mut writer = self.writer.lock().await; + msg::write_msg( + &mut *writer, + &Message::LoginResp(orbien_core::msg::LoginResp { + version: version.into(), + session_id: String::new(), + error: error.into(), + }), + ) + .await?; + Ok(()) + } + + pub(super) fn release_global_slot(&self, name: &str, detached: &DetachedTunnel) { + if let Some(port) = detached.remote_port { + match detached.tunnel_type { + "tcp" => self.tcp_ports.release(port, name), + "udp" => self.udp_ports.release(port, name), + _ => {} + } + } + self.tunnel_registry.remove_if_owner(name, &self.owner()); + } + + pub(super) async fn detach_tunnel(&self, name: &str) -> Option<&'static str> { + let detached = { + let mut tm = self.tunnels.lock().await; + tm.remove(name).await + }?; + let ty = detached.tunnel_type; + self.release_global_slot(name, &detached); + Some(ty) + } + pub async fn run(self: Arc) -> Result<()> { for _ in 0..self.pool_count { if self.closed.load(Ordering::SeqCst) { @@ -133,6 +237,7 @@ impl Control { if last > 0 && now.saturating_sub(last) > timeout { tracing::warn!( session_id = %this.session_id, + generation = this.generation, timeout_secs = timeout, "heartbeat timeout" ); @@ -184,14 +289,16 @@ impl Control { if self.closed.swap(true, Ordering::SeqCst) { self.shutdown_notify.notify_waiters(); self.data_notify.notify_waiters(); + self.wait_finished().await; return; } self.shutdown_notify.notify_waiters(); self.data_notify.notify_waiters(); { let mut tm = self.tunnels.lock().await; - for (name, ty) in tm.close_all().await { - self.metrics.close_tunnel(&name, ty); + for (name, detached) in tm.close_all().await { + self.release_global_slot(&name, &detached); + self.metrics.close_tunnel(&name, detached.tunnel_type); } } { @@ -201,6 +308,7 @@ impl Control { let mut bg = self.bg_tasks.lock().await; bg.abort_all(); while bg.join_next().await.is_some() {} + self.mark_finished(); } pub async fn kick(&self, reason: impl Into) { @@ -215,7 +323,12 @@ impl Control { ) .await; } - tracing::info!(session_id = %self.session_id, %reason, "kicking client"); + tracing::info!( + session_id = %self.session_id, + generation = self.generation, + %reason, + "kicking client" + ); self.shutdown().await; } @@ -246,3 +359,17 @@ impl Control { Ok(()) } } + +impl Drop for Control { + fn drop(&mut self) { + if *self.finished.borrow() { + return; + } + if let Ok(mut tm) = self.tunnels.try_lock() { + for (name, detached) in tm.abandon_all() { + self.release_global_slot(&name, &detached); + } + } + self.mark_finished(); + } +} diff --git a/server/src/control/session/register.rs b/server/src/control/session/register.rs index 54414fa..e06e705 100644 --- a/server/src/control/session/register.rs +++ b/server/src/control/session/register.rs @@ -51,6 +51,9 @@ impl Control { } async fn register_tunnel(self: &Arc, np: &NewTunnel) -> Result { + if self.is_closed() { + return Err(anyhow!("control session is closed")); + } match np.protocol.as_str() { "tcp" => self.register_tcp_tunnel(np).await, "http" => self.register_http_tunnel(np).await, @@ -60,43 +63,66 @@ impl Control { } } + async fn prepare_name_slot(&self, name: &str) { + if let Some(old_ty) = self.detach_tunnel(name).await { + self.metrics.close_tunnel(name, old_ty); + } + } + async fn register_tcp_tunnel(self: &Arc, np: &NewTunnel) -> Result { if np.remote_port <= 0 || np.remote_port > 65535 { return Err(anyhow!("invalid remote_port")); } let limiter = Self::tunnel_transport(np)?; - let bind_addr = self.cfg.proxy_addr.clone(); let remote_port = np.remote_port as u16; let name = np.tunnel_name.clone(); - let control = Arc::clone(self); + let owner = self.owner(); - { - let mut tm = self.tunnels.lock().await; - if let Some(old_ty) = tm.remove(&name).await { - self.metrics.close_tunnel(&name, old_ty); - } + self.prepare_name_slot(&name).await; + self.tunnel_registry.try_insert(&name, owner.clone())?; + if let Err(e) = self.tcp_ports.claim(remote_port, &name) { + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(e); } - let tunnel = TcpTunnel::start( + let tunnel = match TcpTunnel::start( name.clone(), bind_addr, remote_port, - control, + Arc::clone(self), limiter, Arc::clone(&self.access), ) - .await?; - let remote_addr = format!(":{}", remote_port); + .await + { + Ok(t) => t, + Err(e) => { + self.tcp_ports.release(remote_port, &name); + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(e); + } + }; + let remote_addr = format!(":{remote_port}"); let local_addr = format_local_addr(&np.local_ip, np.local_port); let mut tm = self.tunnels.lock().await; - let _ = tm - .insert(name.clone(), RegisteredTunnel::Tcp(tunnel), local_addr) - .await; + if let Err(tunnel) = tm.insert(name.clone(), RegisteredTunnel::Tcp(tunnel), local_addr) { + drop(tm); + tunnel.close().await; + self.tcp_ports.release(remote_port, &name); + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(anyhow!("tunnel `{name}` already present in this session")); + } self.note_tunnel_registered(&name, "tcp"); - tracing::info!(tunnel = %np.tunnel_name, port = remote_port, "tcp tunnel registered"); + tracing::info!( + tunnel = %np.tunnel_name, + port = remote_port, + session_id = %self.session_id, + generation = self.generation, + "tcp tunnel registered" + ); Ok(remote_addr) } @@ -107,23 +133,27 @@ impl Control { .ok_or_else(|| anyhow!("http tunnel requires server httpGwPort > 0"))?; let limiter = Self::tunnel_transport(np)?; - let name = np.tunnel_name.clone(); - { - let mut tm = self.tunnels.lock().await; - if let Some(old_ty) = tm.remove(&name).await { - self.metrics.close_tunnel(&name, old_ty); - } - } + let owner = self.owner(); - let tunnel = HttpTunnel::register( + self.prepare_name_slot(&name).await; + self.tunnel_registry.try_insert(&name, owner.clone())?; + + let tunnel = match HttpTunnel::register( np, Arc::clone(self), Arc::clone(&gw), &self.cfg.root_domain, limiter, ) - .await?; + .await + { + Ok(t) => t, + Err(e) => { + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(e); + } + }; let remote_addr = tunnel .domains @@ -134,9 +164,12 @@ impl Control { let local_addr = format_local_addr(&np.local_ip, np.local_port); let mut tm = self.tunnels.lock().await; - let _ = tm - .insert(name.clone(), RegisteredTunnel::Http(tunnel), local_addr) - .await; + if let Err(tunnel) = tm.insert(name.clone(), RegisteredTunnel::Http(tunnel), local_addr) { + drop(tm); + tunnel.close().await; + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(anyhow!("tunnel `{name}` already present in this session")); + } self.note_tunnel_registered(&name, "http"); Ok(remote_addr) } @@ -148,23 +181,27 @@ impl Control { .ok_or_else(|| anyhow!("https tunnel requires server httpsGwPort > 0"))?; let limiter = Self::tunnel_transport(np)?; - let name = np.tunnel_name.clone(); - { - let mut tm = self.tunnels.lock().await; - if let Some(old_ty) = tm.remove(&name).await { - self.metrics.close_tunnel(&name, old_ty); - } - } + let owner = self.owner(); - let tunnel = HttpsTunnel::register( + self.prepare_name_slot(&name).await; + self.tunnel_registry.try_insert(&name, owner.clone())?; + + let tunnel = match HttpsTunnel::register( np, Arc::clone(self), Arc::clone(&gw), &self.cfg.root_domain, limiter, ) - .await?; + .await + { + Ok(t) => t, + Err(e) => { + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(e); + } + }; let remote_addr = tunnel .domains @@ -175,9 +212,12 @@ impl Control { let local_addr = format_local_addr(&np.local_ip, np.local_port); let mut tm = self.tunnels.lock().await; - let _ = tm - .insert(name.clone(), RegisteredTunnel::Https(tunnel), local_addr) - .await; + if let Err(tunnel) = tm.insert(name.clone(), RegisteredTunnel::Https(tunnel), local_addr) { + drop(tm); + tunnel.close().await; + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(anyhow!("tunnel `{name}` already present in this session")); + } self.note_tunnel_registered(&name, "https"); Ok(remote_addr) } @@ -188,44 +228,60 @@ impl Control { } let limiter = Self::tunnel_transport(np)?; - let bind_addr = self.cfg.proxy_addr.clone(); let remote_port = np.remote_port as u16; let name = np.tunnel_name.clone(); - let control = Arc::clone(self); + let owner = self.owner(); let packet_size = self.cfg.udp_packet_size.max(512); - { - let mut tm = self.tunnels.lock().await; - if let Some(old_ty) = tm.remove(&name).await { - self.metrics.close_tunnel(&name, old_ty); - } + self.prepare_name_slot(&name).await; + self.tunnel_registry.try_insert(&name, owner.clone())?; + if let Err(e) = self.udp_ports.claim(remote_port, &name) { + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(e); } - let tunnel = UdpTunnel::start( + let tunnel = match UdpTunnel::start( name.clone(), bind_addr, remote_port, - control, + Arc::clone(self), limiter, packet_size, ) - .await?; - let remote_addr = format!(":{}", remote_port); + .await + { + Ok(t) => t, + Err(e) => { + self.udp_ports.release(remote_port, &name); + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(e); + } + }; + let remote_addr = format!(":{remote_port}"); let local_addr = format_local_addr(&np.local_ip, np.local_port); let mut tm = self.tunnels.lock().await; - let _ = tm - .insert(name.clone(), RegisteredTunnel::Udp(tunnel), local_addr) - .await; + if let Err(tunnel) = tm.insert(name.clone(), RegisteredTunnel::Udp(tunnel), local_addr) { + drop(tm); + tunnel.close().await; + self.udp_ports.release(remote_port, &name); + self.tunnel_registry.remove_if_owner(&name, &owner); + return Err(anyhow!("tunnel `{name}` already present in this session")); + } self.note_tunnel_registered(&name, "udp"); - tracing::info!(tunnel = %np.tunnel_name, port = remote_port, "udp tunnel registered"); + tracing::info!( + tunnel = %np.tunnel_name, + port = remote_port, + session_id = %self.session_id, + generation = self.generation, + "udp tunnel registered" + ); Ok(remote_addr) } pub(super) async fn handle_close_tunnel(&self, cp: CloseTunnel) -> Result<()> { - let mut tm = self.tunnels.lock().await; - if let Some(ty) = tm.remove(&cp.tunnel_name).await { + if let Some(ty) = self.detach_tunnel(&cp.tunnel_name).await { self.metrics.close_tunnel(&cp.tunnel_name, ty); } Ok(()) diff --git a/server/src/dashboard/mod.rs b/server/src/dashboard/mod.rs index 4d25dab..05eee58 100644 --- a/server/src/dashboard/mod.rs +++ b/server/src/dashboard/mod.rs @@ -1,5 +1,6 @@ pub(crate) mod model; mod routes; +mod snapshot; use crate::service::Service; use anyhow::Result; diff --git a/server/src/dashboard/model.rs b/server/src/dashboard/model.rs index f7e058e..0d76faf 100644 --- a/server/src/dashboard/model.rs +++ b/server/src/dashboard/model.rs @@ -76,6 +76,8 @@ pub struct SystemStatus { pub struct ClientInfo { #[serde(rename = "sessionId")] pub session_id: String, + #[serde(rename = "agentId")] + pub agent_id: String, pub user: String, pub hostname: String, pub os: String, @@ -92,6 +94,16 @@ pub struct ClientInfo { pub status: String, } +impl ClientInfo { + pub fn display_id(&self) -> &str { + if self.agent_id.is_empty() { + &self.session_id + } else { + &self.agent_id + } + } +} + #[derive(Serialize)] pub struct TunnelInfo { pub name: String, diff --git a/server/src/dashboard/routes.rs b/server/src/dashboard/routes.rs index a8d34eb..ccf2431 100644 --- a/server/src/dashboard/routes.rs +++ b/server/src/dashboard/routes.rs @@ -225,7 +225,7 @@ async fn get_client( match snap .clients .into_iter() - .find(|c| c.session_id == session_id) + .find(|c| c.session_id == session_id || c.agent_id == session_id) { Some(c) => Ok(Json(ApiResponse::ok(c))), None => Err(StatusCode::NOT_FOUND), diff --git a/server/src/service/dashboard_view.rs b/server/src/dashboard/snapshot.rs similarity index 82% rename from server/src/service/dashboard_view.rs rename to server/src/dashboard/snapshot.rs index 3a3864f..18c379b 100644 --- a/server/src/service/dashboard_view.rs +++ b/server/src/dashboard/snapshot.rs @@ -1,9 +1,10 @@ -use super::Service; +use crate::dashboard::model::{ClientInfo, TunnelInfo}; +use crate::service::Service; use std::collections::BTreeMap; pub struct DashboardSnapshot { - pub clients: Vec, - pub tunnels: Vec, + pub clients: Vec, + pub tunnels: Vec, pub tunnel_type_count: BTreeMap, pub active_connections: usize, pub total_client_counts: usize, @@ -13,16 +14,15 @@ pub struct DashboardSnapshot { impl Service { pub async fn dashboard_snapshot(&self) -> DashboardSnapshot { - use crate::dashboard::model::{ClientInfo, TunnelInfo}; - let controls = self.controls.lock().await; - let offline = self.offline_clients.lock().await; - let mut clients = Vec::with_capacity(controls.len() + offline.len()); + let agents = self.agents.list(); + let mut clients = Vec::with_capacity(controls.len() + agents.len()); let mut tunnels = Vec::new(); let mut tunnel_type_count: BTreeMap = BTreeMap::new(); let mut online_ids = std::collections::HashSet::new(); - for (_, ctrl) in controls.iter() { + for (_, entry) in controls.iter() { + let ctrl = &entry.control; let tunnel_count = ctrl.tunnel_count().await; online_ids.insert(ctrl.session_id.clone()); let mut active_connections = 0usize; @@ -49,6 +49,7 @@ impl Service { } clients.push(ClientInfo { session_id: ctrl.session_id.clone(), + agent_id: ctrl.agent_id.clone(), user: ctrl.user.clone(), hostname: ctrl.hostname.clone(), os: ctrl.os.clone(), @@ -63,12 +64,13 @@ impl Service { tunnels.extend(client_tunnels); } - for (id, rec) in offline.iter() { - if online_ids.contains(id) { + for rec in agents { + if rec.online || online_ids.contains(&rec.session_id) { continue; } clients.push(ClientInfo { session_id: rec.session_id.clone(), + agent_id: rec.agent_id.clone(), user: rec.user.clone(), hostname: rec.hostname.clone(), os: rec.os.clone(), @@ -77,7 +79,10 @@ impl Service { version: rec.version.clone(), tunnel_count: rec.tunnel_count, active_connections: 0, - connected_secs: rec.disconnected_at.elapsed().as_secs(), + connected_secs: rec + .disconnected_at + .map(|t| t.elapsed().as_secs()) + .unwrap_or(0), status: "offline".into(), }); } @@ -85,7 +90,9 @@ impl Service { clients.sort_by(|a, b| { let ao = a.status == "online"; let bo = b.status == "online"; - bo.cmp(&ao).then_with(|| a.session_id.cmp(&b.session_id)) + bo.cmp(&ao) + .then_with(|| a.display_id().cmp(b.display_id())) + .then_with(|| a.session_id.cmp(&b.session_id)) }); tunnels.sort_by(|a, b| a.name.cmp(&b.name).then(a.session_id.cmp(&b.session_id))); diff --git a/server/src/service/agent_registry.rs b/server/src/service/agent_registry.rs new file mode 100644 index 0000000..0f8bc3d --- /dev/null +++ b/server/src/service/agent_registry.rs @@ -0,0 +1,196 @@ +use std::collections::HashMap; +use std::sync::Mutex; +use std::time::Instant; + +pub const MAX_AGENT_ID_LEN: usize = 64; +pub const MAX_SESSION_ID_LEN: usize = 64; +pub const MAX_USER_LEN: usize = 64; + +#[derive(Debug, Clone)] +pub struct AgentEntry { + pub user: String, + pub agent_id: String, + pub session_id: String, + pub generation: u64, + pub hostname: String, + pub os: String, + pub arch: String, + pub client_ip: String, + pub version: String, + pub tunnel_count: usize, + pub online: bool, + pub disconnected_at: Option, +} + +#[derive(Debug, Clone)] +pub struct AgentOnlineSpec { + pub user: String, + pub agent_id: String, + pub session_id: String, + pub generation: u64, + pub hostname: String, + pub os: String, + pub arch: String, + pub client_ip: String, + pub version: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AgentRegisterError { + Conflict, +} + +#[derive(Debug, Default)] +pub struct AgentRegistry { + by_key: Mutex>, + by_session: Mutex>, +} + +impl AgentRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn try_online(&self, spec: AgentOnlineSpec) -> Result { + let now = Instant::now(); + let explicit = !spec.agent_id.is_empty(); + let effective_id = if explicit { + spec.agent_id.as_str() + } else { + spec.session_id.as_str() + }; + let key = compose_key(&spec.user, effective_id); + + let mut by_key = self.by_key.lock().unwrap_or_else(|e| e.into_inner()); + let mut by_session = self.by_session.lock().unwrap_or_else(|e| e.into_inner()); + + if explicit { + if let Some(existing) = by_key.get(&key) { + if existing.online + && !existing.session_id.is_empty() + && existing.session_id != spec.session_id + { + return Err(AgentRegisterError::Conflict); + } + } + } + + if let Some(prev_key) = by_session.get(&spec.session_id).cloned() { + if prev_key != key { + if let Some(prev) = by_key.get_mut(&prev_key) { + if prev.session_id == spec.session_id { + if prev.agent_id.is_empty() { + by_key.remove(&prev_key); + } else { + set_offline(prev, now, prev.tunnel_count); + } + } + } + by_session.remove(&spec.session_id); + } + } + + match by_key.get_mut(&key) { + Some(entry) => { + if !entry.session_id.is_empty() && entry.session_id != spec.session_id { + by_session.remove(&entry.session_id); + } + entry.user = spec.user; + entry.agent_id = spec.agent_id; + entry.session_id = spec.session_id.clone(); + entry.generation = spec.generation; + entry.hostname = spec.hostname; + entry.os = spec.os; + entry.arch = spec.arch; + entry.client_ip = spec.client_ip; + entry.version = spec.version; + entry.online = true; + entry.disconnected_at = None; + } + None => { + by_key.insert( + key.clone(), + AgentEntry { + user: spec.user, + agent_id: spec.agent_id, + session_id: spec.session_id.clone(), + generation: spec.generation, + hostname: spec.hostname, + os: spec.os, + arch: spec.arch, + client_ip: spec.client_ip, + version: spec.version, + tunnel_count: 0, + online: true, + disconnected_at: None, + }, + ); + } + } + + by_session.insert(spec.session_id, key.clone()); + Ok(key) + } + + pub fn release(&self, session_id: &str, generation: u64, tunnel_count: usize) { + let now = Instant::now(); + let mut by_key = self.by_key.lock().unwrap_or_else(|e| e.into_inner()); + let mut by_session = self.by_session.lock().unwrap_or_else(|e| e.into_inner()); + + let Some(key) = by_session.get(session_id).cloned() else { + return; + }; + let Some(entry) = by_key.get_mut(&key) else { + by_session.remove(session_id); + return; + }; + if entry.session_id != session_id || entry.generation != generation { + return; + } + + by_session.remove(session_id); + if entry.agent_id.is_empty() { + by_key.remove(&key); + } else { + set_offline(entry, now, tunnel_count); + } + } + + pub fn list(&self) -> Vec { + self.by_key + .lock() + .unwrap_or_else(|e| e.into_inner()) + .values() + .cloned() + .collect() + } +} + +fn set_offline(entry: &mut AgentEntry, now: Instant, tunnel_count: usize) { + entry.generation = 0; + entry.online = false; + entry.tunnel_count = tunnel_count; + entry.disconnected_at = Some(now); +} + +fn compose_key(user: &str, id: &str) -> String { + match (user.is_empty(), id.is_empty()) { + (true, _) => id.to_string(), + (_, true) => user.to_string(), + (false, false) => format!("{user}.{id}"), + } +} + +pub fn sanitize_wire_id(raw: &str, max_len: usize) -> Result { + let s = raw.trim(); + if s.is_empty() { + return Ok(String::new()); + } + if s.len() > max_len { + return Err("identifier too long"); + } + if s.chars().any(|c| c.is_control()) { + return Err("identifier contains control characters"); + } + Ok(s.to_string()) +} diff --git a/server/src/service/mod.rs b/server/src/service/mod.rs index 8c69a30..7bd0f8c 100644 --- a/server/src/service/mod.rs +++ b/server/src/service/mod.rs @@ -1,46 +1,39 @@ -mod dashboard_view; +mod agent_registry; mod ingress; mod session_registry; +mod session_table; use crate::access::AccessPolicy; -use crate::control::Control; use crate::metrics::MemMetrics; -use crate::tunnel::{run_http_gw_listener, run_https_gw_listener, HttpGw, HttpsGw}; +use crate::tunnel::{ + run_http_gw_listener, run_https_gw_listener, HttpGw, HttpsGw, PortTable, TunnelRegistry, +}; +use agent_registry::AgentRegistry; use anyhow::{anyhow, Result}; use orbien_core::config::ServerConfig; use orbien_core::transport; +use session_table::SessionMap; use std::collections::HashMap; use std::net::SocketAddr; +use std::sync::atomic::AtomicU64; use std::sync::Arc; -use std::time::Instant; use tokio::net::TcpListener; use tokio::sync::{Mutex, Notify}; use tokio::task::JoinSet; -#[allow(unused_imports)] // public API re-export -pub use dashboard_view::DashboardSnapshot; - -struct OfflineClientRecord { - session_id: String, - user: String, - hostname: String, - os: String, - arch: String, - client_ip: String, - version: String, - tunnel_count: usize, - disconnected_at: Instant, -} - pub struct Service { cfg: ServerConfig, access: Arc, - controls: Arc>>>, - offline_clients: Arc>>, + pub(crate) controls: Arc>, + pub(crate) agents: Arc, http_gw: Option>, https_gw: Option>, tls_config: Arc, - metrics: Arc, + pub(crate) metrics: Arc, + tunnel_registry: Arc, + tcp_ports: Arc, + udp_ports: Arc, + next_generation: AtomicU64, } impl Service { @@ -60,17 +53,21 @@ impl Service { let tls_config = transport::new_server_tls_config(&tls.cert_file, &tls.key_file, &tls.trusted_ca_file)?; if tls.force { - tracing::info!("transport.tls.force=true — non-TLS control connections rejected"); + tracing::info!("transport.tls.force=true, rejecting non-TLS control connections"); } Ok(Self { cfg, access, controls: Arc::new(Mutex::new(HashMap::new())), - offline_clients: Arc::new(Mutex::new(HashMap::new())), + agents: Arc::new(AgentRegistry::new()), http_gw, https_gw, tls_config, metrics: MemMetrics::new(), + tunnel_registry: Arc::new(TunnelRegistry::new()), + tcp_ports: Arc::new(PortTable::new()), + udp_ports: Arc::new(PortTable::new()), + next_generation: AtomicU64::new(1), }) } @@ -177,34 +174,44 @@ impl Service { } pub async fn kick_client(&self, session_id: &str) -> Result<()> { - let control = { - let mut map = self.controls.lock().await; - map.remove(session_id) + let (gate, control) = { + let map = self.controls.lock().await; + match map.get(session_id) { + Some(entry) => (Arc::clone(&entry.gate), Arc::clone(&entry.control)), + None => return Err(anyhow!("client not online: {session_id}")), + } }; - match control { - Some(c) => { - let tunnel_count = c.tunnel_count().await; - { - let mut offline = self.offline_clients.lock().await; - offline.insert( - session_id.to_string(), - OfflineClientRecord { - session_id: session_id.to_string(), - user: c.user.clone(), - hostname: c.hostname.clone(), - os: c.os.clone(), - arch: c.arch.clone(), - client_ip: c.client_ip.clone(), - version: c.version.clone(), - tunnel_count, - disconnected_at: Instant::now(), - }, - ); - } - c.kick("kicked from dashboard").await; - Ok(()) + + let guard = gate.lock_owned().await; + { + let map = self.controls.lock().await; + let still = map + .get(session_id) + .map(|e| Arc::ptr_eq(&e.control, &control)) + .unwrap_or(false); + if !still { + return Err(anyhow!("client not online: {session_id}")); } - None => Err(anyhow!("client not online: {session_id}")), } + + let tunnel_count = control.tunnel_count().await; + let generation = control.generation; + control.kick("kicked from dashboard").await; + control.wait_finished().await; + + { + let mut map = self.controls.lock().await; + if map + .get(session_id) + .map(|cur| Arc::ptr_eq(&cur.control, &control)) + .unwrap_or(false) + { + map.remove(session_id); + } + } + drop(guard); + + self.agents.release(session_id, generation, tunnel_count); + Ok(()) } } diff --git a/server/src/service/session_registry.rs b/server/src/service/session_registry.rs index ee5a493..55a2333 100644 --- a/server/src/service/session_registry.rs +++ b/server/src/service/session_registry.rs @@ -1,4 +1,9 @@ -use super::{OfflineClientRecord, Service}; +use super::agent_registry::{ + sanitize_wire_id, AgentOnlineSpec, AgentRegisterError, MAX_AGENT_ID_LEN, MAX_SESSION_ID_LEN, + MAX_USER_LEN, +}; +use super::session_table::{self, remove_if_current, swap_in_locked}; +use super::Service; use crate::control::Control; use crate::metrics::ServerMetrics; use anyhow::{anyhow, Result}; @@ -7,8 +12,8 @@ use orbien_core::msg::{self, Login, LoginResp, Message, NewDataConn}; use orbien_core::transport::DynStream; use orbien_core::VERSION; use std::net::SocketAddr; +use std::sync::atomic::Ordering; use std::sync::Arc; -use std::time::Instant; use uuid::Uuid; impl Service { @@ -32,66 +37,117 @@ impl Service { return Err(anyhow!("authorization failed")); } - let session_id = if login.session_id.is_empty() { + let user = match sanitize_wire_id(&login.user, MAX_USER_LEN) { + Ok(u) => u, + Err(msg) => { + return reject_login(stream, msg).await; + } + }; + let agent_id = match sanitize_wire_id(&login.agent_id, MAX_AGENT_ID_LEN) { + Ok(id) => id, + Err(msg) => { + return reject_login(stream, msg).await; + } + }; + let session_id = if login.session_id.trim().is_empty() { short_session_id() } else { - login.session_id.clone() + match sanitize_wire_id(&login.session_id, MAX_SESSION_ID_LEN) { + Ok(id) if !id.is_empty() => id, + Ok(_) => short_session_id(), + Err(msg) => { + return reject_login(stream, msg).await; + } + } }; - let mut stream = stream; - msg::write_msg( - &mut stream, - &Message::LoginResp(LoginResp { - version: VERSION.into(), - session_id: session_id.clone(), - error: String::new(), - }), - ) - .await?; - - tracing::info!(%session_id, %peer, pool = login.pool_count, "client logged in"); - let max_pool = self.cfg.transport.max_conn_pool.max(0) as usize; let pool_count = (login.pool_count.max(0) as usize).min(max_pool); - let client_ip = peer.ip().to_string(); + let generation = self.next_generation.fetch_add(1, Ordering::SeqCst); let control = Control::new( session_id.clone(), + generation, stream, self.cfg.clone(), pool_count, self.http_gw.clone(), self.https_gw.clone(), Arc::clone(&self.access), - login.user.clone(), + user.clone(), + agent_id.clone(), login.hostname.clone(), login.os.clone(), login.arch.clone(), login.version.clone(), client_ip, Arc::clone(&self.metrics), + Arc::clone(&self.tunnel_registry), + Arc::clone(&self.tcp_ports), + Arc::clone(&self.udp_ports), ); let control = Arc::new(control); - { - let mut offline = self.offline_clients.lock().await; - offline.remove(&session_id); - } + let (session_guard, previous) = + swap_in_locked(&self.controls, &session_id, Arc::clone(&control)).await; - let old = { - let mut map = self.controls.lock().await; - map.insert(session_id.clone(), Arc::clone(&control)) - }; - - if let Some(old) = old { + if let Some(old) = previous { + tracing::info!( + %session_id, + old_generation = old.generation, + new_generation = generation, + "replacing prior control session" + ); old.shutdown().await; + old.wait_finished().await; } + match self.agents.try_online(AgentOnlineSpec { + user: user.clone(), + agent_id: agent_id.clone(), + session_id: session_id.clone(), + generation, + hostname: control.hostname.clone(), + os: control.os.clone(), + arch: control.arch.clone(), + client_ip: control.client_ip.clone(), + version: control.version.clone(), + }) { + Ok(_) => {} + Err(AgentRegisterError::Conflict) => { + drop(session_guard); + let _ = remove_if_current(&self.controls, &session_id, &control).await; + let err = format!("agent_id [{agent_id}] for user [{user}] is already online"); + let _ = control.send_login_err(VERSION, &err).await; + control.shutdown().await; + return Err(anyhow!(err)); + } + } + + if let Err(e) = control.send_login_ok(VERSION).await { + self.agents.release(&session_id, generation, 0); + drop(session_guard); + let _ = remove_if_current(&self.controls, &session_id, &control).await; + control.shutdown().await; + return Err(e); + } + + drop(session_guard); + + tracing::info!( + %session_id, + %agent_id, + generation, + %peer, + pool = login.pool_count, + "client logged in" + ); + self.metrics.new_client(&session_id); let controls = Arc::clone(&self.controls); - let offline_clients = Arc::clone(&self.offline_clients); + let agents = Arc::clone(&self.agents); let metrics = Arc::clone(&self.metrics); let rid = session_id.clone(); let result = Arc::clone(&control).run().await; @@ -99,32 +155,8 @@ impl Service { metrics.close_client(); let tunnel_count = control.tunnel_count().await; - let mut map = controls.lock().await; - if map - .get(&rid) - .map(|c| Arc::ptr_eq(c, &control)) - .unwrap_or(false) - { - map.remove(&rid); - } - if !map.contains_key(&rid) { - drop(map); - let mut offline = offline_clients.lock().await; - offline.insert( - rid.clone(), - OfflineClientRecord { - session_id: rid, - user: control.user.clone(), - hostname: control.hostname.clone(), - os: control.os.clone(), - arch: control.arch.clone(), - client_ip: control.client_ip.clone(), - version: control.version.clone(), - tunnel_count, - disconnected_at: Instant::now(), - }, - ); - } + let _ = remove_if_current(&controls, &rid, &control).await; + agents.release(&rid, generation, tunnel_count); result } @@ -143,23 +175,33 @@ impl Service { nw.session_id )); } - let control = { - let map = self.controls.lock().await; - map.get(&nw.session_id).cloned() - }; - match control { + + match session_table::lookup_accepting(&self.controls, &nw.session_id).await { Some(c) => { c.push_data_conn(stream).await; Ok(()) } None => Err(anyhow!( - "unknown session_id for data conn: {}", + "no accepting control for data conn session_id={}", nw.session_id )), } } } +async fn reject_login(mut stream: DynStream, error: &str) -> Result<()> { + let _ = msg::write_msg( + &mut stream, + &Message::LoginResp(LoginResp { + version: VERSION.into(), + session_id: String::new(), + error: error.into(), + }), + ) + .await; + Err(anyhow!(error.to_string())) +} + fn short_session_id() -> String { let hex = Uuid::new_v4().simple().to_string(); hex[..16].to_owned() diff --git a/server/src/service/session_table.rs b/server/src/service/session_table.rs new file mode 100644 index 0000000..4a72b0d --- /dev/null +++ b/server/src/service/session_table.rs @@ -0,0 +1,93 @@ +use crate::control::Control; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::{Mutex, OwnedMutexGuard}; + +pub(crate) struct SessionEntry { + pub control: Arc, + pub gate: Arc>, +} + +pub(crate) type SessionMap = HashMap; + +pub(super) async fn swap_in_locked( + controls: &Mutex, + session_id: &str, + control: Arc, +) -> (OwnedMutexGuard<()>, Option>) { + loop { + let peeked_gate = { + let map = controls.lock().await; + map.get(session_id).map(|e| Arc::clone(&e.gate)) + }; + + let gate = peeked_gate.unwrap_or_else(|| Arc::new(Mutex::new(()))); + let guard = Arc::clone(&gate).lock_owned().await; + + let mut map = controls.lock().await; + if let Some(entry) = map.get(session_id) { + if !Arc::ptr_eq(&entry.gate, &gate) { + drop(map); + drop(guard); + continue; + } + } + + let previous = map.insert( + session_id.to_string(), + SessionEntry { + control: Arc::clone(&control), + gate: Arc::clone(&gate), + }, + ); + drop(map); + return (guard, previous.map(|e| e.control)); + } +} + +pub(super) async fn remove_if_current( + controls: &Mutex, + session_id: &str, + control: &Arc, +) -> bool { + let gate = { + let map = controls.lock().await; + match map.get(session_id) { + Some(entry) if Arc::ptr_eq(&entry.control, control) => Arc::clone(&entry.gate), + _ => return false, + } + }; + + let _guard = gate.lock_owned().await; + let mut map = controls.lock().await; + if map + .get(session_id) + .map(|e| Arc::ptr_eq(&e.control, control)) + .unwrap_or(false) + { + map.remove(session_id); + true + } else { + false + } +} + +pub(super) async fn lookup_accepting( + controls: &Mutex, + session_id: &str, +) -> Option> { + let (gate, candidate) = { + let map = controls.lock().await; + let entry = map.get(session_id)?; + (Arc::clone(&entry.gate), Arc::clone(&entry.control)) + }; + + let _guard = gate.lock_owned().await; + let map = controls.lock().await; + let entry = map.get(session_id)?; + if Arc::ptr_eq(&entry.control, &candidate) && candidate.is_accepting_data() { + Some(candidate) + } else { + None + } +} diff --git a/server/src/tunnel/http.rs b/server/src/tunnel/http.rs index af0aca2..c307b26 100644 --- a/server/src/tunnel/http.rs +++ b/server/src/tunnel/http.rs @@ -40,19 +40,24 @@ impl HttpTunnel { for domain in &domains { for location in &locations { - gw.register( - domain, - HttpRoute { - tunnel_name: name.clone(), - control: Arc::downgrade(&control), - location: location.clone(), - host_header_rewrite: rewrite.clone(), - basic_auth_user: basic_auth_user.clone(), - basic_auth_password: basic_auth_password.clone(), - limiter: limiter.clone(), - }, - ) - .await?; + if let Err(e) = gw + .register( + domain, + HttpRoute { + tunnel_name: name.clone(), + control: Arc::downgrade(&control), + location: location.clone(), + host_header_rewrite: rewrite.clone(), + basic_auth_user: basic_auth_user.clone(), + basic_auth_password: basic_auth_password.clone(), + limiter: limiter.clone(), + }, + ) + .await + { + gw.unregister_tunnel(&name).await; + return Err(e); + } } } diff --git a/server/src/tunnel/https.rs b/server/src/tunnel/https.rs index 7f9c706..d8f9dfa 100644 --- a/server/src/tunnel/https.rs +++ b/server/src/tunnel/https.rs @@ -81,15 +81,20 @@ impl HttpsTunnel { gw.unregister_tunnel(&name).await; for domain in &domains { - gw.register( - domain, - HttpsRoute { - tunnel_name: name.clone(), - control: Arc::downgrade(&control), - limiter: limiter.clone(), - }, - ) - .await?; + if let Err(e) = gw + .register( + domain, + HttpsRoute { + tunnel_name: name.clone(), + control: Arc::downgrade(&control), + limiter: limiter.clone(), + }, + ) + .await + { + gw.unregister_tunnel(&name).await; + return Err(e); + } } tracing::info!( diff --git a/server/src/tunnel/manager.rs b/server/src/tunnel/manager.rs index dcf8a82..a7919e0 100644 --- a/server/src/tunnel/manager.rs +++ b/server/src/tunnel/manager.rs @@ -18,6 +18,14 @@ impl RegisteredTunnel { } } + pub fn remote_port(&self) -> Option { + match self { + Self::Tcp(t) => Some(t.remote_port), + Self::Udp(u) => Some(u.remote_port), + Self::Http(_) | Self::Https(_) => None, + } + } + pub async fn close(&self) { match self { Self::Tcp(p) => p.close().await, @@ -28,6 +36,11 @@ impl RegisteredTunnel { } } +pub struct DetachedTunnel { + pub tunnel_type: &'static str, + pub remote_port: Option, +} + struct TunnelEntry { tunnel: RegisteredTunnel, local_addr: String, @@ -44,41 +57,61 @@ impl TunnelManager { } } - pub async fn insert( + pub fn insert( &mut self, name: String, tunnel: RegisteredTunnel, local_addr: String, - ) -> Option<&'static str> { - let entry = TunnelEntry { tunnel, local_addr }; - if let Some(old) = self.tunnels.insert(name, entry) { - let ty = old.tunnel.tunnel_type(); - old.tunnel.close().await; - Some(ty) - } else { - None + ) -> Result<(), RegisteredTunnel> { + use std::collections::hash_map::Entry; + match self.tunnels.entry(name) { + Entry::Vacant(slot) => { + slot.insert(TunnelEntry { tunnel, local_addr }); + Ok(()) + } + Entry::Occupied(_) => Err(tunnel), } } - pub async fn remove(&mut self, name: &str) -> Option<&'static str> { + pub async fn remove(&mut self, name: &str) -> Option { if let Some(entry) = self.tunnels.remove(name) { - let ty = entry.tunnel.tunnel_type(); + let detached = DetachedTunnel { + tunnel_type: entry.tunnel.tunnel_type(), + remote_port: entry.tunnel.remote_port(), + }; entry.tunnel.close().await; - Some(ty) + Some(detached) } else { None } } - pub async fn close_all(&mut self) -> Vec<(String, &'static str)> { + pub async fn close_all(&mut self) -> Vec<(String, DetachedTunnel)> { let mut closed = Vec::with_capacity(self.tunnels.len()); for (name, entry) in self.tunnels.drain() { - closed.push((name, entry.tunnel.tunnel_type())); + let detached = DetachedTunnel { + tunnel_type: entry.tunnel.tunnel_type(), + remote_port: entry.tunnel.remote_port(), + }; entry.tunnel.close().await; + closed.push((name, detached)); } closed } + pub fn abandon_all(&mut self) -> Vec<(String, DetachedTunnel)> { + self.tunnels + .drain() + .map(|(name, entry)| { + let detached = DetachedTunnel { + tunnel_type: entry.tunnel.tunnel_type(), + remote_port: entry.tunnel.remote_port(), + }; + (name, detached) + }) + .collect() + } + pub fn summaries(&self) -> Vec { self.tunnels .iter() diff --git a/server/src/tunnel/mod.rs b/server/src/tunnel/mod.rs index da7336f..2ada1c9 100644 --- a/server/src/tunnel/mod.rs +++ b/server/src/tunnel/mod.rs @@ -2,12 +2,18 @@ mod gw; mod http; mod https; mod manager; +mod ports; +mod registry; mod tcp; mod udp; pub use gw::HttpGw; pub use http::{run_http_gw_listener, HttpTunnel}; pub use https::{run_https_gw_listener, HttpsGw, HttpsTunnel}; -pub use manager::{format_local_addr, RegisteredTunnel, TunnelManager, TunnelSummary}; +pub use manager::{ + format_local_addr, DetachedTunnel, RegisteredTunnel, TunnelManager, TunnelSummary, +}; +pub use ports::PortTable; +pub use registry::{TunnelOwner, TunnelRegistry}; pub use tcp::TcpTunnel; pub use udp::UdpTunnel; diff --git a/server/src/tunnel/ports.rs b/server/src/tunnel/ports.rs new file mode 100644 index 0000000..4662452 --- /dev/null +++ b/server/src/tunnel/ports.rs @@ -0,0 +1,35 @@ +use anyhow::{anyhow, Result}; +use std::collections::HashMap; +use std::sync::Mutex; + +#[derive(Debug, Default)] +pub struct PortTable { + by_port: Mutex>, +} + +impl PortTable { + pub fn new() -> Self { + Self::default() + } + + pub fn claim(&self, port: u16, tunnel_name: &str) -> Result<()> { + if port == 0 { + return Err(anyhow!("invalid remote port 0")); + } + let mut map = self.by_port.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(existing) = map.get(&port) { + return Err(anyhow!( + "remote port {port} is already in use by tunnel `{existing}`" + )); + } + map.insert(port, tunnel_name.to_string()); + Ok(()) + } + + pub fn release(&self, port: u16, tunnel_name: &str) { + let mut map = self.by_port.lock().unwrap_or_else(|e| e.into_inner()); + if map.get(&port).map(|n| n == tunnel_name).unwrap_or(false) { + map.remove(&port); + } + } +} diff --git a/server/src/tunnel/registry.rs b/server/src/tunnel/registry.rs new file mode 100644 index 0000000..385f34d --- /dev/null +++ b/server/src/tunnel/registry.rs @@ -0,0 +1,44 @@ +use anyhow::{anyhow, Result}; +use std::collections::HashMap; +use std::sync::Mutex; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TunnelOwner { + pub session_id: String, + pub generation: u64, +} + +#[derive(Debug, Default)] +pub struct TunnelRegistry { + by_name: Mutex>, +} + +impl TunnelRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn try_insert(&self, name: &str, owner: TunnelOwner) -> Result<()> { + let name = name.trim(); + if name.is_empty() { + return Err(anyhow!("empty tunnel name")); + } + let mut map = self.by_name.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(existing) = map.get(name) { + return Err(anyhow!( + "tunnel `{name}` is already registered (session={}, generation={})", + existing.session_id, + existing.generation + )); + } + map.insert(name.to_string(), owner); + Ok(()) + } + + pub fn remove_if_owner(&self, name: &str, owner: &TunnelOwner) { + let mut map = self.by_name.lock().unwrap_or_else(|e| e.into_inner()); + if map.get(name) == Some(owner) { + map.remove(name); + } + } +}