starred/orbien
1
0
Fork 0
mirror of https://github.com/orbien-org/orbien.git synced 2026-09-22 00:01:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Harden session and tunnel registration, add agent_id identity management, drop session persistence, and support multi-client login (#21)

Browse Source
This commit is contained in:
lxien
2026-09-03 18:21:31 +08:00
parent 78b51168ba
commit af60f723ba
28 changed files with 940 additions and 286 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/src/control/mod.rs
+1 -1
View File
@@ -1,3 +1,3 @@
mod session;
pub use session::{ActiveSession, Control, SessionEnd};
pub use session::{ActiveSession, Control, LoginRejected, SessionEnd};
client/src/control/session.rs
+18 -8
View File
@@ -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<RwLock<ClientConfig>>,
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(_) => {
client/src/handle.rs
+2 -2
View File
@@ -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);
}
client/src/lib.rs
-1
View File
@@ -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};
client/src/service.rs
+13 -15
View File
@@ -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<PathBuf>) -> 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::<LoginRejected>() {
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(());
client/src/session_id.rs
-33
View File
@@ -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<String> {
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)
}
core/src/config/client.rs
+4
View File
@@ -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
core/src/msg/types.rs
+2
View File
@@ -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,
desktop/src/config_bridge.rs
+2
View File
@@ -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,
server-ui/src/types/api.ts
+1
View File
@@ -41,6 +41,7 @@ export interface SystemStatus {
export interface ClientInfo {
sessionId: string
agentId?: string
user: string
hostname: string
os: string
server-ui/src/views/ClientDetail.vue
+3 -2
View File
@@ -195,7 +195,7 @@ onUnmounted(() => {
{{ t('nav.clients') }}
</button>
<span class="crumb-sep" aria-hidden="true">/</span>
<span class="crumb-current mono">{{ client?.sessionId || sessionId }}</span>
<span class="crumb-current mono">{{ client?.agentId || client?.sessionId || sessionId }}</span>
</nav>
<div v-if="loading && !client" class="empty-card">{{ t('traffic.loading') }}</div>
@@ -209,11 +209,12 @@ onUnmounted(() => {
</div>
<div class="head-body">
<div class="title-row">
<h2 class="name mono">{{ client.sessionId }}</h2>
<h2 class="name mono">{{ client.agentId || client.sessionId }}</h2>
<span v-if="client.version" class="tag version">v{{ client.version }}</span>
<span v-if="client.user" class="tag">{{ client.user }}</span>
</div>
<div class="meta">
<span v-if="client.agentId" class="mono">session {{ client.sessionId }}</span>
<span v-if="client.clientIP" class="mono">{{ client.clientIP }}</span>
<OsBadge :os="client.os" :arch="client.arch" size="md" text-only/>
</div>
server-ui/src/views/Clients.vue
+2 -1
View File
@@ -137,7 +137,8 @@ async function onKick(sessionId: string, evt: Event) {
<div class="client-body">
<div class="client-title">
<h3 class="client-id">{{ c.sessionId }}</h3>
<h3 class="client-id">{{ c.agentId || c.sessionId }}</h3>
<span v-if="c.agentId && c.agentId !== c.sessionId" class="tag mono">{{ c.sessionId }}</span>
<span v-if="c.hostname" class="tag">{{ c.hostname }}</span>
<span v-if="c.user" class="tag">{{ c.user }}</span>
<span v-if="c.version" class="tag version">v{{ c.version }}</span>
server/src/control/session/mod.rs
+132 -5
View File
@@ -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<DynStream>;
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<TunnelManager>,
tunnel_registry: Arc<TunnelRegistry>,
tcp_ports: Arc<PortTable>,
udp_ports: Arc<PortTable>,
bg_tasks: Mutex<JoinSet<()>>,
closed: AtomicBool,
finished: watch::Sender<bool>,
activated: AtomicBool,
pool_count: usize,
http_gw: Option<Arc<HttpGw>>,
https_gw: Option<Arc<HttpsGw>>,
@@ -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<Arc<HttpsGw>>,
access: Arc<AccessPolicy>,
user: String,
agent_id: String,
hostname: String,
os: String,
arch: String,
version: String,
client_ip: String,
metrics: Arc<MemMetrics>,
tunnel_registry: Arc<TunnelRegistry>,
tcp_ports: Arc<PortTable>,
udp_ports: Arc<PortTable>,
) -> 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<crate::tunnel::TunnelSummary> {
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<Self>) -> 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<String>) {
@@ -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();
}
}
server/src/control/session/register.rs
+112 -56
View File
@@ -51,6 +51,9 @@ impl Control {
}
async fn register_tunnel(self: &Arc<Self>, np: &NewTunnel) -> Result<String> {
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<Self>, np: &NewTunnel) -> Result<String> {
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(())
server/src/dashboard/mod.rs
+1
View File
@@ -1,5 +1,6 @@
pub(crate) mod model;
mod routes;
mod snapshot;
use crate::service::Service;
use anyhow::Result;
server/src/dashboard/model.rs
+12
View File
@@ -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,
server/src/dashboard/routes.rs
+1 -1
View File
@@ -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),
server/src/service/dashboard_view.rs → server/src/dashboard/snapshot.rs
+19 -12
View File
@@ -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<crate::dashboard::model::ClientInfo>,
pub tunnels: Vec<crate::dashboard::model::TunnelInfo>,
pub clients: Vec<ClientInfo>,
pub tunnels: Vec<TunnelInfo>,
pub tunnel_type_count: BTreeMap<String, usize>,
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<String, usize> = 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)));
server/src/service/agent_registry.rs
+196
View File
@@ -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<Instant>,
}
#[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<HashMap<String, AgentEntry>>,
by_session: Mutex<HashMap<String, String>>,
}
impl AgentRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn try_online(&self, spec: AgentOnlineSpec) -> Result<String, AgentRegisterError> {
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<AgentEntry> {
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<String, &'static str> {
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())
}
server/src/service/mod.rs
+57 -50
View File
@@ -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<AccessPolicy>,
controls: Arc<Mutex<HashMap<String, Arc<Control>>>>,
offline_clients: Arc<Mutex<HashMap<String, OfflineClientRecord>>>,
pub(crate) controls: Arc<Mutex<SessionMap>>,
pub(crate) agents: Arc<AgentRegistry>,
http_gw: Option<Arc<HttpGw>>,
https_gw: Option<Arc<HttpsGw>>,
tls_config: Arc<rustls::ServerConfig>,
metrics: Arc<MemMetrics>,
pub(crate) metrics: Arc<MemMetrics>,
tunnel_registry: Arc<TunnelRegistry>,
tcp_ports: Arc<PortTable>,
udp_ports: Arc<PortTable>,
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(())
}
}
server/src/service/session_registry.rs
+104 -62
View File
@@ -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()
server/src/service/session_table.rs
+93
View File
@@ -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<Control>,
pub gate: Arc<Mutex<()>>,
}
pub(crate) type SessionMap = HashMap<String, SessionEntry>;
pub(super) async fn swap_in_locked(
controls: &Mutex<SessionMap>,
session_id: &str,
control: Arc<Control>,
) -> (OwnedMutexGuard<()>, Option<Arc<Control>>) {
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<SessionMap>,
session_id: &str,
control: &Arc<Control>,
) -> 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<SessionMap>,
session_id: &str,
) -> Option<Arc<Control>> {
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
}
}
server/src/tunnel/http.rs
+18 -13
View File
@@ -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);
}
}
}
server/src/tunnel/https.rs
+14 -9
View File
@@ -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!(
server/src/tunnel/manager.rs
+47 -14
View File
@@ -18,6 +18,14 @@ impl RegisteredTunnel {
}
}
pub fn remote_port(&self) -> Option<u16> {
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<u16>,
}
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<DetachedTunnel> {
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<TunnelSummary> {
self.tunnels
.iter()
server/src/tunnel/mod.rs
+7 -1
View File
@@ -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;
server/src/tunnel/ports.rs
+35
View File
@@ -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<HashMap<u16, String>>,
}
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);
}
}
}
server/src/tunnel/registry.rs
+44
View File
@@ -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<HashMap<String, TunnelOwner>>,
}
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);
}
}
}