feat: release OpsXShell 2.1.0 native Tauri edition

2026-07-30 23:32:51 +08:00
1 parent e9a0d21758
commit a3dba0c6ed
@@ -0,0 +1,432 @@
+use std::{
+    collections::HashMap,
+    io::{self, Write},
+    net::{TcpListener, TcpStream},
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc, Mutex,
+    },
+    thread,
+    time::Duration,
+};
+
+use ssh2::ErrorCode;
+use uuid::Uuid;
+
+use crate::domain::ssh::{ForwardStatus, SshConnectRequest};
+
+use super::ssh::establish_session;
+
+mod model;
+mod relay;
+mod socks5;
+
+use model::{descriptor, ConnectionGuard, TunnelCounters, TunnelRuntime};
+use relay::relay_channel;
+use socks5::read_target;
+
+/// 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<HashMap<String, TunnelRuntime>>,
+}
+
+impl TunnelManager {
+    pub fn start_local(
+        &self,
+        forward_id: Option<String>,
+        connection_id: &str,
+        config: SshConnectRequest,
+        local_port: u16,
+        remote_host: String,
+        remote_port: u16,
+    ) -> Result<ForwardStatus, String> {
+        validate_target(&remote_host, remote_port)?;
+        let forward_id = self.available_id(forward_id)?;
+        let listener = TcpListener::bind(("127.0.0.1", 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",
+            "127.0.0.1",
+            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 listener_thread = thread::spawn(move || {
+            while !thread_cancel.load(Ordering::Acquire) {
+                match listener.accept() {
+                    Ok((socket, _)) => {
+                        let config = config.clone();
+                        let host = remote_host.clone();
+                        let cancel = thread_cancel.clone();
+                        let counters = thread_counters.clone();
+                        thread::spawn(move || {
+                            let Some(_guard) = ConnectionGuard::start(counters.clone(), &cancel)
+                            else {
+                                return;
+                            };
+                            let Ok(session) = establish_session(&config) else {
+                                return;
+                            };
+                            let Ok(channel) =
+                                session.channel_direct_tcpip(&host, remote_port, None)
+                            else {
+                                return;
+                            };
+                            relay_channel(session, channel, socket, cancel, counters);
+                        });
+                    }
+                    Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
+                        thread::sleep(Duration::from_millis(25));
+                    }
+                    Err(_) => break,
+                }
+            }
+            thread_running.store(false, Ordering::Release);
+        });
+        self.register(
+            forward_id,
+            TunnelRuntime {
+                descriptor,
+                cancel,
+                running,
+                counters,
+                listener_thread: Some(listener_thread),
+            },
+        )
+    }
+
+    pub fn start_dynamic(
+        &self,
+        forward_id: Option<String>,
+        connection_id: &str,
+        config: SshConnectRequest,
+        local_port: u16,
+    ) -> Result<ForwardStatus, String> {
+        let forward_id = self.available_id(forward_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 listener_thread = thread::spawn(move || {
+            while !thread_cancel.load(Ordering::Acquire) {
+                match listener.accept() {
+                    Ok((mut socket, _)) => {
+                        let config = config.clone();
+                        let cancel = thread_cancel.clone();
+                        let counters = thread_counters.clone();
+                        thread::spawn(move || {
+                            let Some(_guard) = ConnectionGuard::start(counters.clone(), &cancel)
+                            else {
+                                return;
+                            };
+                            let (host, port) = match read_target(&mut socket) {
+                                Ok(target) => target,
+                                Err(error) => {
+                                    tracing::warn!("SOCKS5 请求解析失败: {error}");
+                                    return;
+                                }
+                            };
+                            if cancel.load(Ordering::Acquire) {
+                                return;
+                            }
+                            let session = match establish_session(&config) {
+                                Ok(session) => session,
+                                Err(error) => {
+                                    tracing::warn!("SOCKS5 专用 SSH 连接失败: {error}");
+                                    let _ = socket.write_all(&[5, 1, 0, 1, 0, 0, 0, 0, 0, 0]);
+                                    return;
+                                }
+                            };
+                            let channel = match session.channel_direct_tcpip(&host, port, None) {
+                                Ok(channel) => channel,
+                                Err(error) => {
+                                    tracing::warn!("SOCKS5 目标 {host}:{port} 连接失败: {error}");
+                                    let _ = socket.write_all(&[5, 1, 0, 1, 0, 0, 0, 0, 0, 0]);
+                                    return;
+                                }
+                            };
+                            let _ = socket.write_all(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 0]);
+                            relay_channel(session, channel, socket, cancel, counters);
+                        });
+                    }
+                    Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
+                        thread::sleep(Duration::from_millis(25));
+                    }
+                    Err(_) => break,
+                }
+            }
+            thread_running.store(false, Ordering::Release);
+        });
+        self.register(
+            forward_id,
+            TunnelRuntime {
+                descriptor,
+                cancel,
+                running,
+                counters,
+                listener_thread: Some(listener_thread),
+            },
+        )
+    }
+
+    pub fn start_remote(
+        &self,
+        forward_id: Option<String>,
+        connection_id: &str,
+        config: SshConnectRequest,
+        remote_port: u16,
+        local_host: String,
+        local_port: u16,
+    ) -> Result<ForwardStatus, String> {
+        validate_target(&local_host, local_port)?;
+        let forward_id = self.available_id(forward_id)?;
+        let session = establish_session(&config)?;
+        let (mut listener, bound_port) = session
+            .channel_forward_listen(remote_port, None, None)
+            .map_err(|error| format!("建立远程端口转发失败: {error}"))?;
+        session.set_blocking(false);
+        let descriptor = descriptor(
+            &forward_id,
+            connection_id,
+            "remote",
+            "remote",
+            bound_port,
+            Some(local_host.clone()),
+            Some(local_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 listener_thread = thread::spawn(move || {
+            while !thread_cancel.load(Ordering::Acquire) {
+                match listener.accept() {
+                    Ok(channel) => {
+                        let Ok(socket) = TcpStream::connect((&*local_host, local_port)) else {
+                            continue;
+                        };
+                        let session = session.clone();
+                        let cancel = thread_cancel.clone();
+                        let counters = thread_counters.clone();
+                        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(_) => break,
+                }
+            }
+            drop(listener);
+            let _ = session.disconnect(None, "port forward stopped", None);
+            thread_running.store(false, Ordering::Release);
+        });
+        self.register(
+            forward_id,
+            TunnelRuntime {
+                descriptor,
+                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::<Vec<_>>();
+            ids.into_iter()
+                .filter_map(|id| all.remove(&id))
+                .collect::<Vec<_>>()
+        };
+        for runtime in runtimes {
+            runtime.stop();
+        }
+        Ok(())
+    }
+
+    pub fn list(&self, connection_id: Option<&str>) -> Result<Vec<ForwardStatus>, String> {
+        let mut statuses = self
+            .runtimes
+            .lock()
+            .map_err(lock_error)?
+            .values()
+            .filter(|runtime| {
+                connection_id.map_or(true, |id| runtime.descriptor.connection_id == id)
+            })
+            .map(TunnelRuntime::snapshot)
+            .collect::<Vec<_>>();
+        statuses.sort_by(|left, right| left.forward_id.cmp(&right.forward_id));
+        Ok(statuses)
+    }
+
+    fn available_id(&self, requested: Option<String>) -> Result<String, String> {
+        let id = requested
+            .filter(|value| !value.trim().is_empty())
+            .unwrap_or_else(|| format!("forward-{}", Uuid::new_v4()));
+        if self.runtimes.lock().map_err(lock_error)?.contains_key(&id) {
+            return Err(format!("端口转发已在运行: {id}"));
+        }
+        Ok(id)
+    }
+
+    fn register(&self, id: String, runtime: TunnelRuntime) -> Result<ForwardStatus, String> {
+        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();
+            }
+        }
+    }
+}
+
+fn validate_target(host: &str, port: u16) -> Result<(), String> {
+    if host.trim().is_empty() || port == 0 {
+        return Err("转发目标主机和端口不能为空".to_string());
+    }
+    Ok(())
+}
+
+fn lock_error<T>(_: std::sync::PoisonError<T>) -> String {
+    "端口转发状态锁异常".to_string()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    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_rejected() {
+        let manager = TunnelManager::default();
+        manager
+            .start_dynamic(
+                Some("same-id".to_string()),
+                "connection-test",
+                unreachable_config(),
+                0,
+            )
+            .unwrap();
+        let error = manager
+            .start_dynamic(
+                Some("same-id".to_string()),
+                "connection-test",
+                unreachable_config(),
+                0,
+            )
+            .unwrap_err();
+        assert!(error.contains("已在运行"));
+        manager.stop_connection("connection-test").unwrap();
+    }
+}