use std::{ collections::{HashMap, HashSet}, io::{self, Write}, net::{TcpListener, TcpStream}, sync::{ atomic::{AtomicBool, Ordering}, Arc, Mutex, }, thread, time::{Duration, Instant}, }; use ssh2::ErrorCode; use uuid::Uuid; use crate::domain::ssh::{ForwardStatus, SshConnectRequest}; use super::ssh::establish_session; mod model; mod relay; mod session; mod socks5; use model::{descriptor, ConnectionGuard, TunnelCounters, TunnelRuntime}; use relay::relay_channel; use session::TunnelSession; use socks5::read_target; enum IdClaim { Fresh(String), Existing(ForwardStatus), } /// Owns forwarding listeners independently from persisted forwarding rules. /// A rule may exist without a runtime; a runtime always has a cancellable id. #[derive(Default)] pub struct TunnelManager { runtimes: Mutex>, } impl TunnelManager { pub fn start_local( &self, forward_id: Option, connection_id: &str, config: SshConnectRequest, local_bind_host: String, local_port: u16, remote_host: String, remote_port: u16, ) -> Result { let remote_host = validate_target(&remote_host, remote_port)?; let bind_host = validate_bind_host(&local_bind_host)?; let forward_id = match self.claim_id(forward_id, connection_id, "local")? { IdClaim::Existing(status) => return Ok(status), IdClaim::Fresh(id) => id, }; let listener = TcpListener::bind((bind_host.as_str(), local_port)) .map_err(|error| format!("监听本地端口失败: {error}"))?; listener .set_nonblocking(true) .map_err(|error| format!("设置本地端口非阻塞模式失败: {error}"))?; let bound_port = listener .local_addr() .map_err(|error| error.to_string())? .port(); let descriptor = descriptor( &forward_id, connection_id, "local", &bind_host, bound_port, Some(remote_host.clone()), Some(remote_port), ); let cancel = Arc::new(AtomicBool::new(false)); let running = Arc::new(AtomicBool::new(true)); let counters = Arc::new(TunnelCounters::default()); let thread_cancel = cancel.clone(); let thread_running = running.clone(); let thread_counters = counters.clone(); let fingerprint = connection_fingerprint(&config); let listener_thread = thread::spawn(move || { run_direct_listener( listener, config, DirectTarget::Fixed { host: remote_host, port: remote_port, }, thread_cancel, thread_running, thread_counters, ); }); self.register( forward_id, TunnelRuntime { descriptor, fingerprint, cancel, running, counters, listener_thread: Some(listener_thread), }, ) } pub fn start_dynamic( &self, forward_id: Option, connection_id: &str, config: SshConnectRequest, local_port: u16, ) -> Result { let forward_id = match self.claim_id(forward_id, connection_id, "dynamic")? { IdClaim::Existing(status) => return Ok(status), IdClaim::Fresh(id) => id, }; let listener = TcpListener::bind(("127.0.0.1", local_port)) .map_err(|error| format!("监听 SOCKS5 端口失败: {error}"))?; listener .set_nonblocking(true) .map_err(|error| format!("设置 SOCKS5 端口非阻塞模式失败: {error}"))?; let bound_port = listener .local_addr() .map_err(|error| error.to_string())? .port(); let descriptor = descriptor( &forward_id, connection_id, "dynamic", "127.0.0.1", bound_port, None, None, ); let cancel = Arc::new(AtomicBool::new(false)); let running = Arc::new(AtomicBool::new(true)); let counters = Arc::new(TunnelCounters::default()); let thread_cancel = cancel.clone(); let thread_running = running.clone(); let thread_counters = counters.clone(); let fingerprint = connection_fingerprint(&config); let listener_thread = thread::spawn(move || { run_direct_listener( listener, config, DirectTarget::Dynamic, thread_cancel, thread_running, thread_counters, ); }); self.register( forward_id, TunnelRuntime { descriptor, fingerprint, cancel, running, counters, listener_thread: Some(listener_thread), }, ) } pub fn start_remote( &self, forward_id: Option, connection_id: &str, config: SshConnectRequest, remote_bind_host: String, remote_port: u16, local_host: String, local_port: u16, ) -> Result { let local_host = validate_target(&local_host, local_port)?; let remote_bind_host = validate_remote_bind_host(&remote_bind_host)?; let forward_id = match self.claim_id(forward_id, connection_id, "remote")? { IdClaim::Existing(status) => return Ok(status), IdClaim::Fresh(id) => id, }; let session = establish_session(&config)?; let want_all_interfaces = is_all_interfaces_bind(&remote_bind_host); let (listener, bound_port, bind_host) = listen_remote_forward(&session, remote_port, &remote_bind_host).map_err(|error| { map_remote_forward_error(error, remote_port, want_all_interfaces) })?; session.set_blocking(false); let descriptor = descriptor( &forward_id, connection_id, "remote", if bind_host.is_empty() { "localhost" } else { bind_host.as_str() }, bound_port, Some(local_host.clone()), Some(local_port), ); let fingerprint = connection_fingerprint(&config); let cancel = Arc::new(AtomicBool::new(false)); let running = Arc::new(AtomicBool::new(true)); let counters = Arc::new(TunnelCounters::default()); let thread_cancel = cancel.clone(); let thread_running = running.clone(); let thread_counters = counters.clone(); let thread_config = config.clone(); let thread_bind = remote_bind_host.clone(); let listener_thread = thread::spawn(move || { run_remote_listener( session, listener, thread_config, remote_port, thread_bind, local_host, local_port, thread_cancel, thread_running, thread_counters, ); }); self.register( forward_id, TunnelRuntime { descriptor, fingerprint, cancel, running, counters, listener_thread: Some(listener_thread), }, ) } pub fn stop(&self, forward_id: &str) -> Result<(), String> { let runtime = self .runtimes .lock() .map_err(lock_error)? .remove(forward_id) .ok_or_else(|| format!("端口转发不存在或已停止: {forward_id}"))?; runtime.stop(); Ok(()) } pub fn stop_connection(&self, connection_id: &str) -> Result<(), String> { let runtimes = { let mut all = self.runtimes.lock().map_err(lock_error)?; let ids = all .iter() .filter_map(|(id, runtime)| { (runtime.descriptor.connection_id == connection_id).then_some(id.clone()) }) .collect::>(); ids.into_iter() .filter_map(|id| all.remove(&id)) .collect::>() }; for runtime in runtimes { runtime.stop(); } Ok(()) } pub fn list(&self, connection_id: Option<&str>) -> Result, String> { let mut statuses = self .runtimes .lock() .map_err(lock_error)? .values() .filter(|runtime| connection_id.is_none_or(|id| runtime.descriptor.connection_id == id)) .map(TunnelRuntime::snapshot) .collect::>(); statuses.sort_by(|left, right| left.forward_id.cmp(&right.forward_id)); Ok(statuses) } /// Rebind listeners left behind when the pooled SSH connection was replaced. /// Only takes runtimes whose SSH identity matches and whose previous connection /// is no longer live — a second session to the same host must not steal them. pub fn adopt_orphans( &self, new_connection_id: &str, fingerprint: &str, live_connection_ids: &HashSet, ) -> Result, String> { let mut adopted = Vec::new(); let mut runtimes = self.runtimes.lock().map_err(lock_error)?; for runtime in runtimes.values_mut() { if runtime.fingerprint != fingerprint { continue; } if runtime.descriptor.connection_id == new_connection_id { continue; } if live_connection_ids.contains(&runtime.descriptor.connection_id) { continue; } runtime.descriptor.connection_id = new_connection_id.to_string(); adopted.push(runtime.snapshot()); } Ok(adopted) } /// Claim a forward id for start. /// - Same connection + kind still running → return existing status (idempotent). /// - Same id dead, owned by another connection, or a different kind → stop and reuse. /// Session-scoped rules keep a stable id across reconnects; orphans must not block start. fn claim_id( &self, requested: Option, connection_id: &str, kind: &str, ) -> Result { let id = requested .filter(|value| !value.trim().is_empty()) .unwrap_or_else(|| format!("forward-{}", Uuid::new_v4())); let stale = { let mut runtimes = self.runtimes.lock().map_err(lock_error)?; match runtimes.get(&id) { Some(runtime) if runtime.descriptor.connection_id == connection_id && runtime.descriptor.kind == kind && runtime.running.load(Ordering::Acquire) => { return Ok(IdClaim::Existing(runtime.snapshot())); } Some(_) => runtimes.remove(&id), None => None, } }; if let Some(runtime) = stale { runtime.stop(); } Ok(IdClaim::Fresh(id)) } fn register(&self, id: String, runtime: TunnelRuntime) -> Result { let status = runtime.snapshot(); let mut runtimes = self.runtimes.lock().map_err(lock_error)?; match runtimes.entry(id) { std::collections::hash_map::Entry::Vacant(entry) => { entry.insert(runtime); Ok(status) } std::collections::hash_map::Entry::Occupied(_) => { runtime.stop(); Err("端口转发 ID 冲突".to_string()) } } } } impl Drop for TunnelManager { fn drop(&mut self) { if let Ok(runtimes) = self.runtimes.get_mut() { for (_, runtime) in runtimes.drain() { runtime.stop(); } } } } pub(crate) fn connection_fingerprint(config: &SshConnectRequest) -> String { format!( "{}@{}:{}", config.username.trim(), config.host.trim().to_ascii_lowercase(), config.port ) } #[derive(Clone)] enum DirectTarget { Fixed { host: String, port: u16 }, Dynamic, } /// Local / SOCKS listener loop. The TCP bind stays up until the user stops /// the rule; SSH session flaps only rebuild the dedicated transport. fn run_direct_listener( listener: TcpListener, config: SshConnectRequest, target: DirectTarget, cancel: Arc, running: Arc, counters: Arc, ) { running.store(true, Ordering::Release); let tunnel = std::sync::Arc::new(TunnelSession::new(config)); let mut last_keepalive = Instant::now(); while !cancel.load(Ordering::Acquire) { if last_keepalive.elapsed() >= Duration::from_secs(15) { tunnel.keepalive(); last_keepalive = Instant::now(); } match listener.accept() { Ok((socket, _)) => { let cancel = cancel.clone(); let counters = counters.clone(); let tunnel = tunnel.clone(); let target = target.clone(); thread::spawn(move || { handle_direct_client(socket, target, tunnel, cancel, counters); }); } Err(error) if error.kind() == io::ErrorKind::WouldBlock || error.kind() == io::ErrorKind::Interrupted => { thread::sleep(Duration::from_millis(25)); } Err(error) => { tracing::warn!("端口转发本地监听 accept 失败,继续等待: {error}"); thread::sleep(Duration::from_millis(100)); } } } tunnel.close(); running.store(false, Ordering::Release); } fn handle_direct_client( mut socket: TcpStream, target: DirectTarget, tunnel: std::sync::Arc, cancel: Arc, counters: Arc, ) { let Some(_guard) = ConnectionGuard::start(counters.clone(), &cancel) else { return; }; let is_dynamic = matches!(target, DirectTarget::Dynamic); let (host, port) = match target { DirectTarget::Fixed { host, port } => (host, port), DirectTarget::Dynamic => match read_target(&mut socket) { Ok(destination) => destination, Err(error) => { tracing::warn!("SOCKS5 请求解析失败: {error}"); return; } }, }; if cancel.load(Ordering::Acquire) { return; } match tunnel.open_direct(&host, port, &cancel) { Ok((session, channel)) => { if is_dynamic { let _ = socket.write_all(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 0]); } relay_channel(session, channel, socket, cancel, counters); } Err(error) => { tracing::warn!("端口转发目标 {host}:{port} 连接失败: {error}"); if is_dynamic { let _ = socket.write_all(&[5, 1, 0, 1, 0, 0, 0, 0, 0, 0]); } } } } fn run_remote_listener( mut session: ssh2::Session, mut listener: ssh2::Listener, config: SshConnectRequest, remote_port: u16, remote_bind_host: String, local_host: String, local_port: u16, cancel: Arc, running: Arc, counters: Arc, ) { running.store(true, Ordering::Release); while !cancel.load(Ordering::Acquire) { let died = accept_remote_until_dead( &session, &mut listener, &local_host, local_port, &cancel, &counters, ); if cancel.load(Ordering::Acquire) || !died { break; } tracing::warn!("远程端口转发会话断开,正在重连 {remote_bind_host}:{remote_port}"); drop(listener); let _ = session.disconnect(None, "port forward reconnect", None); let Some((new_session, new_listener)) = reconnect_remote_listener(&config, remote_port, &remote_bind_host, &cancel) else { running.store(false, Ordering::Release); return; }; session = new_session; listener = new_listener; } drop(listener); let _ = session.disconnect(None, "port forward stopped", None); running.store(false, Ordering::Release); } fn accept_remote_until_dead( session: &ssh2::Session, listener: &mut ssh2::Listener, local_host: &str, local_port: u16, cancel: &Arc, counters: &Arc, ) -> bool { let mut last_keepalive = Instant::now(); while !cancel.load(Ordering::Acquire) { if last_keepalive.elapsed() >= Duration::from_secs(15) { let _ = session.keepalive_send(); last_keepalive = Instant::now(); } match listener.accept() { Ok(channel) => { let Ok(socket) = TcpStream::connect((local_host, local_port)) else { continue; }; let session = session.clone(); let cancel = Arc::clone(cancel); let counters = Arc::clone(counters); thread::spawn(move || { let Some(_guard) = ConnectionGuard::start(counters.clone(), &cancel) else { return; }; relay_channel(session, channel, socket, cancel, counters); }); } Err(error) if error.code() == ErrorCode::Session(-37) => { thread::sleep(Duration::from_millis(25)); } Err(_) => return true, } } false } fn reconnect_remote_listener( config: &SshConnectRequest, remote_port: u16, remote_bind_host: &str, cancel: &AtomicBool, ) -> Option<(ssh2::Session, ssh2::Listener)> { let mut delay = Duration::from_millis(400); while !cancel.load(Ordering::Acquire) { match establish_session(config) { Ok(session) => match listen_remote_forward(&session, remote_port, remote_bind_host) { Ok((listener, _, _)) => { session.set_blocking(false); return Some((session, listener)); } Err(error) => { tracing::warn!("远程转发重听失败: {error}"); let _ = session.disconnect(None, "remote re-listen failed", None); } }, Err(error) => { tracing::warn!("远程转发重连失败: {error}"); } } thread::sleep(delay); delay = (delay.saturating_mul(2)).min(Duration::from_secs(15)); } None } fn validate_target(host: &str, port: u16) -> Result { let host = host.trim(); if host.is_empty() || port == 0 { return Err("转发目标主机和端口不能为空".to_string()); } // 0.0.0.0 / :: are listen-only; treat as remote loopback (common user mix-up). if host == "0.0.0.0" || host == "::" || host.eq_ignore_ascii_case("[::]") { return Ok("127.0.0.1".to_string()); } Ok(host.to_string()) } fn validate_bind_host(host: &str) -> Result { let host = host.trim(); if host.is_empty() { return Ok("127.0.0.1".to_string()); } if host.chars().any(char::is_whitespace) { return Err("本地绑定地址不能包含空格".to_string()); } Ok(host.to_string()) } fn validate_remote_bind_host(host: &str) -> Result { let host = host.trim(); if host.is_empty() || host.eq_ignore_ascii_case("localhost") { return Ok("127.0.0.1".to_string()); } if host.chars().any(char::is_whitespace) { return Err("远端监听地址不能包含空格".to_string()); } Ok(host.to_string()) } fn is_all_interfaces_bind(host: &str) -> bool { let host = host.trim(); host == "0.0.0.0" || host == "*" || host == "::" || host.eq_ignore_ascii_case("[::]") } fn lock_error(_: std::sync::PoisonError) -> String { "端口转发状态锁异常".to_string() } /// Ask sshd to listen for remote forwards. /// Loopback: try localhost / 127.0.0.1 / empty (OpenSSH -R default). /// All-interfaces (`0.0.0.0`): requires GatewayPorts yes/clientspecified. fn listen_remote_forward( session: &ssh2::Session, remote_port: u16, prefer_bind: &str, ) -> Result<(ssh2::Listener, u16, String), ssh2::Error> { let prefer = prefer_bind.trim(); let candidates: Vec<&str> = if is_all_interfaces_bind(prefer) { // OpenSSH clients send 0.0.0.0 (or *) when GatewayPorts allows it. vec!["0.0.0.0"] } else if prefer.is_empty() || prefer == "127.0.0.1" || prefer.eq_ignore_ascii_case("localhost") { // Order matters: PermitListen often lists "localhost:port"; empty string is // what `ssh -R port:host:port` sends when no bind address is given. vec!["localhost", "127.0.0.1", ""] } else { vec![prefer] }; let mut denied: Option = None; for host in candidates { match session.channel_forward_listen(remote_port, Some(host), Some(1)) { Ok((listener, bound)) => { let reported = if host.is_empty() { "localhost".to_string() } else { host.to_string() }; return Ok((listener, bound, reported)); } Err(error) if error.code() == ErrorCode::Session(-32) => { denied = Some(error); } Err(error) => return Err(error), } } Err(denied.unwrap_or_else(|| ssh2::Error::from_errno(ErrorCode::Session(-32)))) } /// libssh2 Session(-32) == LIBSSH2_ERROR_REQUEST_DENIED: sshd refused tcpip-forward. fn map_remote_forward_error( error: ssh2::Error, remote_port: u16, want_all_interfaces: bool, ) -> String { match error.code() { ErrorCode::Session(-32) if want_all_interfaces => format!( "建立远程端口转发失败: 远端拒绝 0.0.0.0:{remote_port}(通常需 GatewayPorts yes/clientspecified,且端口空闲)。仅本机访问可改用 127.0.0.1:{remote_port}。" ), ErrorCode::Session(-32) if (1..1024).contains(&remote_port) => format!( "建立远程端口转发失败: 远端拒绝特权端口 {remote_port}(需 root 或改用 ≥1024;检查 AllowTcpForwarding / PermitListen)。" ), ErrorCode::Session(-32) => format!( "建立远程端口转发失败: 远端拒绝端口 {remote_port}(AllowTcpForwarding / PermitListen / 端口占用)。可用 ssh -R {remote_port}:127.0.0.1:<本地端口> 核对。" ), _ => format!("建立远程端口转发失败: {error}"), } } #[cfg(test)] mod tests { use super::*; use std::collections::HashSet; fn unreachable_config() -> SshConnectRequest { SshConnectRequest { host: "127.0.0.1".to_string(), port: 1, username: "test".to_string(), password: None, private_key: None, passphrase: None, proxy: None, jump_host: None, jump_hosts: None, timeout_seconds: Some(3), keep_alive: Some(true), } } #[test] fn dynamic_listener_has_cancellable_runtime_state() { let manager = TunnelManager::default(); let status = manager .start_dynamic( Some("proxy-test".to_string()), "connection-test", unreachable_config(), 0, ) .unwrap(); assert_eq!(status.forward_id, "proxy-test"); assert!(status.bound_port > 0); assert!(manager.list(Some("connection-test")).unwrap()[0].running); manager.stop("proxy-test").unwrap(); assert!(manager.list(None).unwrap().is_empty()); } #[test] fn duplicate_runtime_id_is_idempotent_for_same_connection() { let manager = TunnelManager::default(); let first = manager .start_dynamic( Some("same-id".to_string()), "connection-test", unreachable_config(), 0, ) .unwrap(); let second = manager .start_dynamic( Some("same-id".to_string()), "connection-test", unreachable_config(), 0, ) .unwrap(); assert_eq!(first.forward_id, second.forward_id); assert_eq!(first.bound_port, second.bound_port); assert_eq!(manager.list(Some("connection-test")).unwrap().len(), 1); manager.stop_connection("connection-test").unwrap(); } #[test] fn stale_runtime_id_from_other_connection_is_reclaimed() { let manager = TunnelManager::default(); let first = manager .start_dynamic( Some("shared-rule".to_string()), "connection-old", unreachable_config(), 0, ) .unwrap(); let second = manager .start_dynamic( Some("shared-rule".to_string()), "connection-new", unreachable_config(), 0, ) .unwrap(); assert_eq!(second.forward_id, "shared-rule"); assert_eq!(second.connection_id, "connection-new"); assert!(manager.list(Some("connection-old")).unwrap().is_empty()); assert_eq!(manager.list(Some("connection-new")).unwrap().len(), 1); // Old listener was stopped; new bind may reuse the port or pick another (port 0). assert_ne!(second.bound_port, 0); let _ = first; manager.stop_connection("connection-new").unwrap(); } #[test] fn adopt_orphans_rebinds_to_new_connection() { let manager = TunnelManager::default(); let config = unreachable_config(); manager .start_dynamic(Some("rule-1".to_string()), "old-conn", config.clone(), 0) .unwrap(); let fingerprint = connection_fingerprint(&config); let live = HashSet::from(["new-conn".to_string()]); let adopted = manager .adopt_orphans("new-conn", &fingerprint, &live) .unwrap(); assert_eq!(adopted.len(), 1); assert_eq!(adopted[0].connection_id, "new-conn"); assert!(manager.list(Some("old-conn")).unwrap().is_empty()); assert_eq!( manager.list(Some("new-conn")).unwrap()[0].forward_id, "rule-1" ); manager.stop_connection("new-conn").unwrap(); } #[test] fn adopt_skips_runtimes_still_owned_by_live_connection() { let manager = TunnelManager::default(); let config = unreachable_config(); manager .start_dynamic(Some("rule-live".to_string()), "old-conn", config.clone(), 0) .unwrap(); let fingerprint = connection_fingerprint(&config); let live = HashSet::from(["old-conn".to_string(), "new-conn".to_string()]); let adopted = manager .adopt_orphans("new-conn", &fingerprint, &live) .unwrap(); assert!(adopted.is_empty()); assert_eq!(manager.list(Some("old-conn")).unwrap().len(), 1); assert!(manager.list(Some("new-conn")).unwrap().is_empty()); manager.stop_connection("old-conn").unwrap(); } #[test] fn dead_runtime_is_reclaimed_on_start() { let manager = TunnelManager::default(); manager .start_dynamic( Some("same-id".to_string()), "connection-test", unreachable_config(), 0, ) .unwrap(); { let runtimes = manager.runtimes.lock().unwrap(); runtimes .get("same-id") .unwrap() .running .store(false, Ordering::Release); } let restarted = manager .start_dynamic( Some("same-id".to_string()), "connection-test", unreachable_config(), 0, ) .unwrap(); assert_eq!(restarted.forward_id, "same-id"); assert!(restarted.running); assert_eq!(manager.list(Some("connection-test")).unwrap().len(), 1); manager.stop_connection("connection-test").unwrap(); } }