network-infra-utility 0.3.0 → 0.6.0

data/service/ssh/ext/ssh_core_rs/src/keepalive.rs

@@ -0,0 +1,194 @@
+//! 保活管理器 — 周期性检查所有活动连接，发送 keepalive 心跳。
+//! 对应 Erlang 端的 ssh_keepalive_mgr.erl。
+//!
+//! 核心设计（与 Erlang 端完全对齐）：
+//!   - 定时扫描所有 Ready 状态的连接，发送 keepalive 全局请求
+//!   - 连续 KEEPALIVE_MAX_FAIL（3）次失败 → 推送 conn.closed
+//!   - 重连决策由 Ruby 决定（通过 conn.reconnect RPC），核心不自主重连
+//!   - 活动感知：有数据流时跳过心跳（简化版：由 channel.rs 通知活动）
+//!
+//! RPC 方法：
+//!   - keepalive.set_interval：设置心跳间隔
+//!   - keepalive.get_interval：获取当前间隔
+//!   - keepalive.get_status：获取所有连接的保活状态
+
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+
+use serde_json::{json, Value};
+use tokio::sync::{Mutex, RwLock};
+use tracing::{debug, warn};
+
+use crate::codec;
+use crate::conn::ConnManager;
+use crate::gateway::Gateway;
+use crate::proto;
+
+/// 默认保活间隔（毫秒）
+const KEEPALIVE_DEFAULT_INTERVAL_MS: u64 = 30_000;
+
+/// 最大连续失败次数
+const KEEPALIVE_MAX_FAIL: u32 = 3;
+
+/// 保活管理器
+pub struct KeepaliveMgr {
+    /// 心跳间隔（毫秒）
+    interval_ms: RwLock<u64>,
+    /// 失败计数：conn_id → count
+    fail_counts: Mutex<HashMap<String, u32>>,
+    /// 最近活动时间戳：conn_id → unix_ms
+    last_activity: Mutex<HashMap<String, u64>>,
+    /// 上次检查时间戳：conn_id → unix_ms
+    last_check: Mutex<HashMap<String, u64>>,
+    /// 连接管理器引用
+    conn_mgr: Arc<ConnManager>,
+    /// 网关引用
+    gateway: Arc<Gateway>,
+}
+
+impl KeepaliveMgr {
+    pub fn new(conn_mgr: Arc<ConnManager>, gateway: Arc<Gateway>) -> Arc<Self> {
+        Arc::new(Self {
+            interval_ms: RwLock::new(KEEPALIVE_DEFAULT_INTERVAL_MS),
+            fail_counts: Mutex::new(HashMap::new()),
+            last_activity: Mutex::new(HashMap::new()),
+            last_check: Mutex::new(HashMap::new()),
+            conn_mgr,
+            gateway,
+        })
+    }
+
+    /// 启动保活循环
+    pub fn start(self: &Arc<Self>) {
+        let this = Arc::clone(self);
+        tokio::spawn(async move {
+            this.run_loop().await;
+        });
+    }
+
+    /// 保活主循环
+    async fn run_loop(&self) {
+        loop {
+            let ms = *self.interval_ms.read().await;
+            tokio::time::sleep(Duration::from_millis(ms)).await;
+            self.tick().await;
+        }
+    }
+
+    /// 单次保活检查
+    async fn tick(&self) {
+        let now = codec::unix_ms();
+        let conn_ids = self.conn_mgr.ready_conn_ids().await;
+
+        // 记录检查时间
+        {
+            let mut checks = self.last_check.lock().await;
+            for id in &conn_ids {
+                checks.insert(id.clone(), now);
+            }
+        }
+
+        for conn_id in &conn_ids {
+            self.check_keepalive(conn_id).await;
+        }
+
+        // 清理已断开的连接
+        let active_ids: std::collections::HashSet<String> = conn_ids.into_iter().collect();
+        {
+            let mut fails = self.fail_counts.lock().await;
+            fails.retain(|k, _| active_ids.contains(k));
+        }
+        {
+            let mut acts = self.last_activity.lock().await;
+            acts.retain(|k, _| active_ids.contains(k));
+        }
+    }
+
+    /// 检查单个连接的保活
+    async fn check_keepalive(&self, conn_id: &str) {
+        // 发送 keepalive 全局请求
+        let handle = self.conn_mgr.get_handle(conn_id).await;
+        if handle.is_none() {
+            // 连接已断开
+            return;
+        }
+        let handle = handle.unwrap();
+
+        let result = handle
+            .send_keepalive(true)
+            .await;
+
+        match result {
+            Ok(_) => {
+                // 成功——重置失败计数
+                self.fail_counts.lock().await.remove(conn_id);
+                debug!("keepalive ok for conn {}", conn_id);
+            }
+            Err(e) => {
+                let mut fails = self.fail_counts.lock().await;
+                let count = fails.entry(conn_id.to_string()).or_insert(0);
+                *count += 1;
+
+                if *count >= KEEPALIVE_MAX_FAIL {
+                    warn!(
+                        "conn {} keepalive failed {}/{}: {}, marking closed",
+                        conn_id, *count, KEEPALIVE_MAX_FAIL, e
+                    );
+                    drop(fails);
+                    self.conn_mgr
+                        .mark_conn_closed(conn_id, "keepalive_failed")
+                        .await;
+                } else {
+                    warn!(
+                        "conn {} keepalive failed {}/{}: {}",
+                        conn_id, *count, KEEPALIVE_MAX_FAIL, e
+                    );
+                }
+            }
+        }
+    }
+
+    /// 通知活动（由 channel.rs 在数据收发时调用）
+    pub async fn notify_activity(&self, conn_id: &str) {
+        let now = codec::unix_ms();
+        self.last_activity.lock().await.insert(conn_id.to_string(), now);
+        self.fail_counts.lock().await.remove(conn_id);
+    }
+
+    /// RPC: keepalive.set_interval
+    pub async fn rpc_set_interval(&self, params: Value) -> Result<Value, String> {
+        let interval_ms = proto::get_num(&params, "interval_ms")
+            .ok_or("missing interval_ms")? as u64;
+        *self.interval_ms.write().await = interval_ms;
+        Ok(json!({"ok": true}))
+    }
+
+    /// RPC: keepalive.get_interval
+    pub async fn rpc_get_interval(&self, _params: Value) -> Result<Value, String> {
+        let ms = *self.interval_ms.read().await;
+        Ok(json!({"interval_ms": ms}))
+    }
+
+    /// RPC: keepalive.get_status
+    pub async fn rpc_get_status(&self, _params: Value) -> Result<Value, String> {
+        let fails = self.fail_counts.lock().await;
+        let acts = self.last_activity.lock().await;
+        let checks = self.last_check.lock().await;
+
+        let conn_ids = self.conn_mgr.all_conn_ids().await;
+        let status: Vec<Value> = conn_ids
+            .iter()
+            .map(|id| {
+                json!({
+                    "conn_id": id,
+                    "fail_count": fails.get(id).copied().unwrap_or(0),
+                    "last_activity": acts.get(id).copied().unwrap_or(0),
+                    "last_check": checks.get(id).copied().unwrap_or(0),
+                })
+            })
+            .collect();
+
+        Ok(json!({"connections": status}))
+    }
+}

data/service/ssh/ext/ssh_core_rs/src/main.rs

@@ -0,0 +1,351 @@
+//! SSH Core RS — 主入口
+//! 对应 Erlang 端的 ssh_core_app.erl + ssh_infra_sup.erl + ssh_ipc_gateway.erl:start_link/1。
+//!
+//! 启动流程：
+//!   1. 生成 auth_token
+//!   2. 删除旧的端点文件
+//!   3. 创建 Gateway（TCP 监听 127.0.0.1:0，分配端口）
+//!   4. 写入端点文件（tcp://127.0.0.1:port\n auth_token）
+//!   5. 创建 ConnManager、Coalescer、ChannelManager、KeepaliveMgr
+//!   6. 注册所有 RPC 路由
+//!   7. 启动 coalesce tick 循环和 keepalive 循环
+//!   8. 进入 gateway accept 循环（阻塞）
+
+mod codec;
+mod coalesce;
+mod channel;
+mod conn;
+mod gateway;
+mod handler;
+mod keepalive;
+mod portfwd;
+mod proto;
+mod sftp;
+
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use serde_json::{json, Value};
+use tracing::{error, info};
+
+use channel::ChannelManager;
+use conn::ConnManager;
+use gateway::Gateway;
+use keepalive::KeepaliveMgr;
+use coalesce::Coalescer;
+use portfwd::PortFwdManager;
+use sftp::SftpManager;
+use std::time::Instant;
+
+/// 端点文件路径：与 Erlang 端完全一致
+/// Windows: %TEMP%\ssh_core_<USERNAME>.endpoint
+/// Unix:    /tmp/ssh_core_<uid>.endpoint
+fn endpoint_file_path() -> PathBuf {
+    let tmp = std::env::temp_dir();
+    let uid = if cfg!(windows) {
+        std::env::var("USERNAME").unwrap_or_else(|_| "default".to_string())
+    } else {
+        std::env::var("UID").unwrap_or_else(|_| {
+            // 尝试通过 id -u 获取 uid
+            std::process::Command::new("id")
+                .arg("-u")
+                .output()
+                .ok()
+                .and_then(|o| String::from_utf8(o.stdout).ok())
+                .map(|s| s.trim().to_string())
+                .unwrap_or_else(|| "0".to_string())
+        })
+    };
+    tmp.join(format!("ssh_core_{}.endpoint", uid))
+}
+
+#[tokio::main]
+async fn main() {
+    // 初始化日志
+    tracing_subscriber::fmt()
+        .with_env_filter(
+            tracing_subscriber::EnvFilter::try_from_default_env()
+                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
+        )
+        .init();
+
+    info!("ssh_core_rs starting up...");
+
+    // 0. 记录启动时间（用于 engine.stats 的 uptime_ms）
+    let start_time = Instant::now();
+
+    // 1. 生成 auth_token
+    let auth_token = codec::gen_auth_token();
+
+    // 2. 删除旧端点文件
+    let endpoint_file = endpoint_file_path();
+    if endpoint_file.exists() {
+        let _ = std::fs::remove_file(&endpoint_file);
+    }
+    info!("endpoint file: {}", endpoint_file.display());
+
+    // 3. 创建 Gateway
+    let gateway = Gateway::new(auth_token);
+
+    // 4. 创建各管理器
+    let conn_mgr = ConnManager::new(Arc::clone(&gateway));
+    let coalesce = Coalescer::new(Arc::clone(&gateway));
+    let channel_mgr = ChannelManager::new(Arc::clone(&conn_mgr), Arc::clone(&coalesce));
+    let keepalive_mgr = KeepaliveMgr::new(Arc::clone(&conn_mgr), Arc::clone(&gateway));
+    let sftp_mgr = SftpManager::new(Arc::clone(&conn_mgr));
+    let portfwd_mgr = PortFwdManager::new(Arc::clone(&conn_mgr));
+
+    // 5. 注册所有 RPC 路由
+    register_routes(
+        &gateway,
+        &conn_mgr,
+        &channel_mgr,
+        &keepalive_mgr,
+        &sftp_mgr,
+        &portfwd_mgr,
+        start_time,
+    )
+    .await;
+
+    // 6. 启动 coalesce tick 循环
+    coalesce.start();
+
+    // 7. 启动 keepalive 循环
+    keepalive_mgr.start();
+
+    // 8. 启动 Gateway（写入端点文件 + accept 循环）
+    info!("starting IPC gateway...");
+    if let Err(e) = gateway.start(endpoint_file).await {
+        error!("gateway failed: {}", e);
+        std::process::exit(1);
+    }
+}
+
+/// 注册所有 JSON-RPC 路由
+async fn register_routes(
+    gateway: &Arc<Gateway>,
+    conn_mgr: &Arc<ConnManager>,
+    channel_mgr: &Arc<ChannelManager>,
+    keepalive_mgr: &Arc<KeepaliveMgr>,
+    sftp_mgr: &Arc<SftpManager>,
+    portfwd_mgr: &Arc<PortFwdManager>,
+    start_time: Instant,
+) {
+    // --- 连接管理 ---
+    let cm = Arc::clone(conn_mgr);
+    gateway
+        .register_route("conn.connect", move |params| {
+            let cm = Arc::clone(&cm);
+            async move { cm.rpc_connect(params).await }
+        })
+        .await;
+
+    let cm = Arc::clone(conn_mgr);
+    gateway
+        .register_route("conn.disconnect", move |params| {
+            let cm = Arc::clone(&cm);
+            async move { cm.rpc_disconnect(params).await }
+        })
+        .await;
+
+    let cm = Arc::clone(conn_mgr);
+    gateway
+        .register_route("conn.list", move |params| {
+            let cm = Arc::clone(&cm);
+            async move { cm.rpc_list(params).await }
+        })
+        .await;
+
+    let cm = Arc::clone(conn_mgr);
+    gateway
+        .register_route("conn.reconnect", move |params| {
+            let cm = Arc::clone(&cm);
+            async move { cm.rpc_reconnect(params).await }
+        })
+        .await;
+
+    // --- 通道管理 ---
+    let chm = Arc::clone(channel_mgr);
+    gateway
+        .register_route("channel.open", move |params| {
+            let chm = Arc::clone(&chm);
+            async move { chm.rpc_open(params).await }
+        })
+        .await;
+
+    let chm = Arc::clone(channel_mgr);
+    gateway
+        .register_route("channel.send", move |params| {
+            let chm = Arc::clone(&chm);
+            async move { chm.rpc_send(params).await }
+        })
+        .await;
+
+    let chm = Arc::clone(channel_mgr);
+    gateway
+        .register_route("channel.close", move |params| {
+            let chm = Arc::clone(&chm);
+            async move { chm.rpc_close(params).await }
+        })
+        .await;
+
+    let chm = Arc::clone(channel_mgr);
+    gateway
+        .register_route("channel.window_change", move |params| {
+            let chm = Arc::clone(&chm);
+            async move { chm.rpc_window_change(params).await }
+        })
+        .await;
+
+    // --- 保活管理 ---
+    let km = Arc::clone(keepalive_mgr);
+    gateway
+        .register_route("keepalive.set_interval", move |params| {
+            let km = Arc::clone(&km);
+            async move { km.rpc_set_interval(params).await }
+        })
+        .await;
+
+    let km = Arc::clone(keepalive_mgr);
+    gateway
+        .register_route("keepalive.get_interval", move |params| {
+            let km = Arc::clone(&km);
+            async move { km.rpc_get_interval(params).await }
+        })
+        .await;
+
+    let km = Arc::clone(keepalive_mgr);
+    gateway
+        .register_route("keepalive.get_status", move |params| {
+            let km = Arc::clone(&km);
+            async move { km.rpc_get_status(params).await }
+        })
+        .await;
+
+    // --- 引擎管理 ---
+    gateway
+        .register_route("engine.ping", move |_params| {
+            async move { Ok(json!({"ok": true, "timestamp": codec::unix_ms()})) }
+        })
+        .await;
+
+    let cm2 = Arc::clone(conn_mgr);
+    gateway
+        .register_route("engine.stats", move |_params| {
+            let cm2 = Arc::clone(&cm2);
+            let st = start_time;
+            async move {
+                let conn_count = cm2.all_conn_ids().await.len();
+                let uptime_ms = st.elapsed().as_millis() as u64;
+                Ok(json!({
+                    "connections": conn_count,
+                    "uptime_ms": uptime_ms
+                }))
+            }
+        })
+        .await;
+
+    gateway
+        .register_route("engine.shutdown", move |_params| {
+            async move {
+                info!("engine.shutdown received, exiting...");
+                tokio::spawn(async {
+                    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+                    std::process::exit(0);
+                });
+                Ok(json!({"ok": true}))
+            }
+        })
+        .await;
+
+    // --- bye ---
+    gateway
+        .register_route("bye", move |_params| {
+            async move { Ok(json!({"ok": true})) }
+        })
+        .await;
+
+    // --- SFTP ---
+    let sm = Arc::clone(sftp_mgr);
+    gateway
+        .register_route("sftp.open", move |params| {
+            let sm = Arc::clone(&sm);
+            async move { sm.rpc_open(params).await }
+        })
+        .await;
+
+    let sm = Arc::clone(sftp_mgr);
+    gateway
+        .register_route("sftp.list_dir", move |params| {
+            let sm = Arc::clone(&sm);
+            async move { sm.rpc_list_dir(params).await }
+        })
+        .await;
+
+    let sm = Arc::clone(sftp_mgr);
+    gateway
+        .register_route("sftp.download", move |params| {
+            let sm = Arc::clone(&sm);
+            async move { sm.rpc_download(params).await }
+        })
+        .await;
+
+    let sm = Arc::clone(sftp_mgr);
+    gateway
+        .register_route("sftp.upload", move |params| {
+            let sm = Arc::clone(&sm);
+            async move { sm.rpc_upload(params).await }
+        })
+        .await;
+
+    let sm = Arc::clone(sftp_mgr);
+    gateway
+        .register_route("sftp.mkdir", move |params| {
+            let sm = Arc::clone(&sm);
+            async move { sm.rpc_mkdir(params).await }
+        })
+        .await;
+
+    let sm = Arc::clone(sftp_mgr);
+    gateway
+        .register_route("sftp.remove", move |params| {
+            let sm = Arc::clone(&sm);
+            async move { sm.rpc_remove(params).await }
+        })
+        .await;
+
+    let sm = Arc::clone(sftp_mgr);
+    gateway
+        .register_route("sftp.stat", move |params| {
+            let sm = Arc::clone(&sm);
+            async move { sm.rpc_stat(params).await }
+        })
+        .await;
+
+    // --- 端口转发 ---
+    let pm = Arc::clone(portfwd_mgr);
+    gateway
+        .register_route("portfwd.add", move |params| {
+            let pm = Arc::clone(&pm);
+            async move { pm.rpc_add(params).await }
+        })
+        .await;
+
+    let pm = Arc::clone(portfwd_mgr);
+    gateway
+        .register_route("portfwd.remove", move |params| {
+            let pm = Arc::clone(&pm);
+            async move { pm.rpc_remove(params).await }
+        })
+        .await;
+
+    let pm = Arc::clone(portfwd_mgr);
+    gateway
+        .register_route("portfwd.list", move |params| {
+            let pm = Arc::clone(&pm);
+            async move { pm.rpc_list(params).await }
+        })
+        .await;
+
+    info!("all RPC routes registered");
+}