network-infra-utility 0.3.0 → 0.5.0

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

@@ -0,0 +1,389 @@
+//! IPC 网关 — TCP 监听、JSON-RPC 路由、推送、反向 RPC。
+//! 对应 Erlang 端的 ssh_ipc_gateway.erl。
+//!
+//! 核心设计（与 Erlang 端完全对齐）：
+//!   - 监听 127.0.0.1 随机端口，写入端点文件供 Ruby 读取
+//!   - hello 握手认证（auth_token）
+//!   - 路由表：method → 处理函数
+//!   - 推送：push_event / push_batch（notification，无 id）
+//!   - 反向 RPC：synchronous_push（Rust→Ruby 带 id 的请求，阻塞等响应）
+
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::atomic::{AtomicI64, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use serde_json::{json, Value};
+use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
+use tokio::net::{TcpListener, TcpStream};
+use tokio::sync::{oneshot, Mutex, RwLock};
+use tracing::{debug, error, info, warn};
+
+use crate::proto;
+
+/// RPC 处理函数类型：接收 params，返回 Result<Value, String>
+type RpcHandler = Arc<
+    dyn Fn(Value) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Value, String>> + Send>>
+        + Send
+        + Sync,
+>;
+
+/// 反向 RPC 待处理条目
+struct ReversePending {
+    tx: oneshot::Sender<Result<Value, String>>,
+}
+
+/// 客户端连接信息
+struct ClientInfo {
+    writer_tx: tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
+    authed: bool,
+    reader_handle: Option<tokio::task::JoinHandle<()>>,
+}
+
+/// 网关共享状态
+pub struct Gateway {
+    auth_token: String,
+    routes: RwLock<HashMap<String, RpcHandler>>,
+    clients: RwLock<Vec<Arc<Mutex<ClientInfo>>>>,
+    reverse_pending: Mutex<HashMap<i64, ReversePending>>,
+    next_reverse_id: AtomicI64,
+}
+
+impl Gateway {
+    /// 创建网关实例
+    pub fn new(auth_token: String) -> Arc<Self> {
+        Arc::new(Self {
+            auth_token,
+            routes: RwLock::new(HashMap::new()),
+            clients: RwLock::new(Vec::new()),
+            reverse_pending: Mutex::new(HashMap::new()),
+            next_reverse_id: AtomicI64::new(1000000),
+        })
+    }
+
+    /// 注册 RPC 路由
+    pub async fn register_route<F, Fut>(&self, method: &str, handler: F)
+    where
+        F: Fn(Value) -> Fut + Send + Sync + 'static,
+        Fut: std::future::Future<Output = Result<Value, String>> + Send + 'static,
+    {
+        let handler_arc: RpcHandler = Arc::new(move |params| {
+            let f = handler(params);
+            Box::pin(f)
+        });
+        self.routes.write().await.insert(method.to_string(), handler_arc);
+    }
+
+    /// 启动 TCP 监听并写入端点文件
+    pub async fn start(self: Arc<Self>, endpoint_file: PathBuf) -> Result<(), String> {
+        // 监听 127.0.0.1:0（OS 分配端口）
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .map_err(|e| format!("bind failed: {}", e))?;
+
+        let port = listener
+            .local_addr()
+            .map_err(|e| format!("local_addr failed: {}", e))?
+            .port();
+
+        // 写入端点文件：第一行 tcp://127.0.0.1:<port>，第二行 auth_token
+        let endpoint_content = format!("tcp://127.0.0.1:{}\n{}", port, self.auth_token);
+        std::fs::write(&endpoint_file, &endpoint_content)
+            .map_err(|e| format!("write endpoint file failed: {}", e))?;
+
+        // 设置文件权限（Unix: 600）
+        #[cfg(unix)]
+        {
+            use std::os::unix::fs::PermissionsExt;
+            let mut perms = std::fs::metadata(&endpoint_file)
+                .map_err(|e| format!("metadata failed: {}", e))?
+                .permissions();
+            perms.set_mode(0o600);
+            let _ = std::fs::set_permissions(&endpoint_file, perms);
+        }
+
+        info!("IPC gateway listening on 127.0.0.1:{}", port);
+        info!("endpoint file: {}", endpoint_file.display());
+
+        // 接受连接循环
+        loop {
+            match listener.accept().await {
+                Ok((stream, addr)) => {
+                    debug!("new client connection from {}", addr);
+                    let gw = Arc::clone(&self);
+                    tokio::spawn(async move {
+                        gw.handle_client(stream).await;
+                    });
+                }
+                Err(e) => {
+                    error!("accept failed: {}", e);
+                }
+            }
+        }
+    }
+
+    /// 处理单个客户端连接
+    async fn handle_client(self: Arc<Self>, stream: TcpStream) {
+        let (reader, mut writer) = stream.into_split();
+
+        // 创建写通道
+        let (writer_tx, mut writer_rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
+
+        let client = Arc::new(Mutex::new(ClientInfo {
+            writer_tx: writer_tx.clone(),
+            authed: false,
+            reader_handle: None,
+        }));
+
+        // 注册客户端
+        self.clients.write().await.push(Arc::clone(&client));
+
+        // 启动写任务
+        let write_task = tokio::spawn(async move {
+            while let Some(data) = writer_rx.recv().await {
+                if writer.write_all(&data).await.is_err() {
+                    break;
+                }
+            }
+        });
+
+        // 读取循环
+        let buf_reader = BufReader::new(reader);
+        let mut lines = buf_reader.lines();
+        let mut leftover = String::new();
+
+        loop {
+            match lines.next_line().await {
+                Ok(Some(line)) => {
+                    let frame = if leftover.is_empty() {
+                        line
+                    } else {
+                        format!("{}{}", leftover, line)
+                    };
+                    leftover.clear();
+
+                    if frame.is_empty() {
+                        continue;
+                    }
+
+                    // 检查消息大小
+                    if frame.len() > proto::MAX_MSG_SIZE {
+                        warn!("message too large ({} bytes), discarding", frame.len());
+                        continue;
+                    }
+
+                    let msg = match proto::decode(frame.as_bytes()) {
+                        Ok(m) => m,
+                        Err(e) => {
+                            warn!("decode error: {}", e);
+                            continue;
+                        }
+                    };
+
+                    self.handle_message(msg, &client).await;
+                }
+                Ok(None) => {
+                    // EOF
+                    debug!("client disconnected");
+                    break;
+                }
+                Err(e) => {
+                    debug!("read error: {}", e);
+                    break;
+                }
+            }
+        }
+
+        // 清理：从客户端列表移除
+        self.clients.write().await.retain(|c| {
+            !Arc::ptr_eq(c, &client)
+        });
+
+        write_task.abort();
+    }
+
+    /// 处理收到的消息
+    async fn handle_message(&self, msg: proto::Message, client: &Arc<Mutex<ClientInfo>>) {
+        match msg {
+            proto::Message::Request { id, method, params } => {
+                if method == "hello" {
+                    // hello 握手认证
+                    let result = self.handle_hello(&params).await;
+                    match result {
+                        Ok(resp) => {
+                            // 认证成功
+                            client.lock().await.authed = true;
+                            self.send_response(id, &resp, client).await;
+                        }
+                        Err(e) => {
+                            self.send_error(id, -32001, &e, client).await;
+                        }
+                    }
+                    return;
+                }
+
+                // 检查是否已认证
+                if !client.lock().await.authed {
+                    self.send_error(id, -32000, "not authenticated", client).await;
+                    return;
+                }
+
+                // 路由到处理函数
+                let handler = self.routes.read().await.get(&method).cloned();
+                match handler {
+                    Some(h) => {
+                        let client_clone = Arc::clone(client);
+                        let id_clone = id;
+                        let self_clone = Arc::new(()); // 不需要 self，直接用闭包
+                        tokio::spawn(async move {
+                            let _ = self_clone; // 消除警告
+                            match h(params).await {
+                                Ok(result) => {
+                                    client_clone
+                                        .lock()
+                                        .await
+                                        .writer_tx
+                                        .send(proto::frame(&proto::encode_response(id_clone, &result)))
+                                        .ok();
+                                }
+                                Err(e) => {
+                                    client_clone
+                                        .lock()
+                                        .await
+                                        .writer_tx
+                                        .send(proto::frame(&proto::encode_error(id_clone, -32603, &e)))
+                                        .ok();
+                                }
+                            }
+                        });
+                    }
+                    None => {
+                        self.send_error(id, -32601, &format!("method not found: {}", method), client)
+                            .await;
+                    }
+                }
+            }
+            proto::Message::Response { id, result, error } => {
+                // 反向 RPC 的响应
+                let mut pending = self.reverse_pending.lock().await;
+                if let Some(entry) = pending.remove(&id) {
+                    if let Some(err) = error {
+                        let _ = entry.tx.send(Err(err.to_string()));
+                    } else {
+                        let _ = entry.tx.send(Ok(result.unwrap_or(Value::Null)));
+                    }
+                }
+            }
+            proto::Message::Notification { .. } => {
+                // Ruby→Rust 的通知，目前忽略
+                debug!("received notification from Ruby (ignored)");
+            }
+        }
+    }
+
+    /// 处理 hello 握手
+    async fn handle_hello(&self, params: &Value) -> Result<Value, String> {
+        let token = proto::get_str(params, "auth_token").unwrap_or("");
+        if token != self.auth_token {
+            return Err("auth_failed".to_string());
+        }
+
+        Ok(json!({
+            "capabilities": ["coalesce"],
+            "ver": proto::PROTOCOL_VER
+        }))
+    }
+
+    /// 发送响应给客户端
+    async fn send_response(&self, id: i64, result: &Value, client: &Arc<Mutex<ClientInfo>>) {
+        let data = proto::frame(&proto::encode_response(id, result));
+        let _ = client.lock().await.writer_tx.send(data);
+    }
+
+    /// 发送错误给客户端
+    async fn send_error(&self, id: i64, code: i64, message: &str, client: &Arc<Mutex<ClientInfo>>) {
+        let data = proto::frame(&proto::encode_error(id, code, message));
+        let _ = client.lock().await.writer_tx.send(data);
+    }
+
+    /// 推送通知给所有已认证客户端
+    pub async fn push_event(&self, method: &str, params: Value) {
+        let data = proto::frame(&proto::encode_push(method, &params));
+        let clients = self.clients.read().await;
+        for client in clients.iter() {
+            let c = client.lock().await;
+            if c.authed {
+                let _ = c.writer_tx.send(data.clone());
+            }
+        }
+    }
+
+    /// 推送批量通知（channel.data.batch）
+    pub async fn push_batch(&self, params: Value) {
+        self.push_event("channel.data.batch", params).await;
+    }
+
+    /// 同步反向 RPC（Rust→Ruby）：发送带 id 的请求，阻塞等响应。
+    /// 用于 hostkey.resolve。
+    pub async fn synchronous_push(
+        &self,
+        method: &str,
+        params: Value,
+        timeout_ms: u64,
+    ) -> Result<Value, String> {
+        let id = self.next_reverse_id.fetch_add(1, Ordering::SeqCst);
+
+        let (tx, rx) = oneshot::channel();
+        self.reverse_pending
+            .lock()
+            .await
+            .insert(id, ReversePending { tx });
+
+        // 构造反向 RPC 请求消息（有 id 和 method）
+        let msg = json!({
+            "jsonrpc": "2.0",
+            "id": id,
+            "method": method,
+            "params": params
+        });
+        let data = proto::frame(&msg.to_string());
+
+        // 发送给所有已认证客户端
+        let sent = {
+            let clients = self.clients.read().await;
+            let mut sent = false;
+            for client in clients.iter() {
+                let c = client.lock().await;
+                if c.authed {
+                    let _ = c.writer_tx.send(data.clone());
+                    sent = true;
+                }
+            }
+            sent
+        };
+
+        if !sent {
+            self.reverse_pending.lock().await.remove(&id);
+            return Err("no_clients".to_string());
+        }
+
+        // 等待响应或超时
+        match tokio::time::timeout(Duration::from_millis(timeout_ms), rx).await {
+            Ok(Ok(result)) => result,
+            Ok(Err(_)) => {
+                self.reverse_pending.lock().await.remove(&id);
+                Err("channel closed".to_string())
+            }
+            Err(_) => {
+                self.reverse_pending.lock().await.remove(&id);
+                Err("timeout".to_string())
+            }
+        }
+    }
+
+    /// 获取已连接客户端数
+    pub async fn client_count(&self) -> usize {
+        self.clients.read().await.len()
+    }
+}

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

@@ -0,0 +1,293 @@
+//! russh Handler 实现 — SSH 协议事件回调，桥接到 IPC 推送。
+//! 对应 dssh/src/handler.rs + Erlang 端 channel_stm 的数据推送逻辑。
+//!
+//! 架构说明：
+//!   russh 0.61.2 的 Handler trait 使用 `&mut self`（引用），方法返回
+//!   `Result<(), Self::Error>`。handler 持有 Arc<ConnShared>，通过共享状态
+//!   将通道数据推送到 channel.rs 的接收循环，将连接事件推送到 conn.rs 的事件循环。
+
+use std::collections::HashMap;
+use std::future::Future;
+use std::sync::Arc;
+
+use russh::client;
+use russh::keys::HashAlg;
+use russh::keys::PublicKey;
+use russh::ChannelId;
+use serde_json::json;
+use tokio::sync::{mpsc, RwLock};
+use tracing::{debug, warn};
+
+use crate::gateway::Gateway;
+
+/// 连接级事件（从 handler 转发到连接管理器的事件循环）
+#[derive(Debug)]
+pub enum ConnEvent {
+    /// 连接被远端关闭
+    Disconnected,
+}
+
+/// 单个通道的数据接收器。
+/// channel.rs 在 open 时通过 register_channel_tx 注册，
+/// handler 的 data/extended_data 回调通过此 tx 推送数据到 channel.rs 的接收循环。
+pub type ChannelDataTx = mpsc::UnboundedSender<ChannelData>;
+
+/// 来自 SSH 通道的数据消息
+#[derive(Debug)]
+pub enum ChannelData {
+    /// 标准输出数据
+    Stdout { data: Vec<u8> },
+    /// 扩展数据（stderr 等）
+    Extended { data: Vec<u8>, ext: u32 },
+    /// 通道 EOF/关闭
+    Eof { reason: String },
+}
+
+/// 连接级共享状态——由 handler 持有 Arc，在 &mut self 回调中共享
+pub struct ConnShared {
+    /// conn_id（用于推送和日志）
+    pub conn_id: String,
+    /// 主机地址（用于 hostkey.resolve）
+    pub host: String,
+    /// 端口（用于 hostkey.resolve）
+    pub port: u16,
+    /// IPC 网关（用于 check_server_key 反向 RPC）
+    pub gateway: Arc<Gateway>,
+    /// 通道数据发送器映射：russh ChannelId → channel.rs 的接收端
+    channel_txs: RwLock<HashMap<ChannelId, ChannelDataTx>>,
+    /// 连接级事件发送器（连接关闭等）
+    conn_event_tx: mpsc::UnboundedSender<ConnEvent>,
+    /// 服务器主机密钥指纹（check_server_key 回调中写入，authenticate 中读取）
+    server_fingerprint: RwLock<Option<String>>,
+}
+
+impl ConnShared {
+    pub fn new(
+        conn_id: String,
+        host: String,
+        port: u16,
+        gateway: Arc<Gateway>,
+        conn_event_tx: mpsc::UnboundedSender<ConnEvent>,
+    ) -> Arc<Self> {
+        Arc::new(Self {
+            conn_id,
+            host,
+            port,
+            gateway,
+            channel_txs: RwLock::new(HashMap::new()),
+            conn_event_tx,
+            server_fingerprint: RwLock::new(None),
+        })
+    }
+
+    /// 注册通道数据发送器（channel.rs 在 channel.open 时调用）
+    pub async fn register_channel_tx(&self, russh_ch_id: ChannelId, tx: ChannelDataTx) {
+        self.channel_txs.write().await.insert(russh_ch_id, tx);
+    }
+
+    /// 注销通道数据发送器（channel.rs 在 channel.close 时调用）
+    pub async fn unregister_channel_tx(&self, russh_ch_id: ChannelId) {
+        self.channel_txs.write().await.remove(&russh_ch_id);
+    }
+
+    /// 向指定通道推送数据
+    async fn send_channel_data(&self, channel: ChannelId, msg: ChannelData) {
+        let txs = self.channel_txs.read().await;
+        if let Some(tx) = txs.get(&channel) {
+            if tx.send(msg).is_err() {
+                debug!("conn {} channel {}: data rx dropped", self.conn_id, channel);
+            }
+        }
+    }
+
+    /// 获取服务器指纹（authenticate 完成后可读取）
+    pub async fn get_server_fingerprint(&self) -> Option<String> {
+        self.server_fingerprint.read().await.clone()
+    }
+}
+
+/// russh Handler 实现
+pub struct SshHandler {
+    /// 共享状态 Arc——在 &mut self 回调中共享
+    pub shared: Arc<ConnShared>,
+}
+
+impl SshHandler {
+    pub fn new(shared: Arc<ConnShared>) -> Self {
+        Self { shared }
+    }
+}
+
+/// server key 验证超时（毫秒）
+const HOSTKEY_RESOLVE_TIMEOUT_MS: u64 = 30_000;
+
+/// russh 0.61.2 Handler trait 使用原生 AFIT（`fn -> impl Future` 返回类型）。
+/// 不启用 async-trait 特性，方法用 `fn -> impl Future { async { ... } }` 风格实现。
+impl client::Handler for SshHandler {
+    type Error = russh::Error;
+
+    fn check_server_key(
+        &mut self,
+        server_public_key: &PublicKey,
+    ) -> impl Future<Output = Result<bool, Self::Error>> + Send {
+        async {
+            let fingerprint = server_public_key
+                .fingerprint(HashAlg::Sha256)
+                .to_string();
+            debug!(
+                "conn {}: server key fingerprint: {}",
+                self.shared.conn_id, fingerprint
+            );
+
+            // 存储指纹到共享状态（authenticate 完成后读取返回给 Ruby）
+            *self.shared.server_fingerprint.write().await = Some(fingerprint.clone());
+
+            // 通过反向 RPC 询问 Ruby 是否接受此主机密钥
+            let resp = self
+                .shared
+                .gateway
+                .synchronous_push(
+                    "hostkey.resolve",
+                    json!({
+                        "host": self.shared.host,
+                        "port": self.shared.port,
+                        "fingerprint": fingerprint,
+                    }),
+                    HOSTKEY_RESOLVE_TIMEOUT_MS,
+                )
+                .await;
+
+            let accepted = match resp {
+                Ok(result) => {
+                    let action = result
+                        .get("action")
+                        .and_then(|v| v.as_str())
+                        .unwrap_or("reject");
+                    match action {
+                        "accept" => true,
+                        "once" => true,
+                        "reject" => false,
+                        _ => {
+                            warn!(
+                                "conn {}: hostkey.resolve returned unknown action: {}, reject",
+                                self.shared.conn_id, action
+                            );
+                            false
+                        }
+                    }
+                }
+                Err(e) => {
+                    warn!(
+                        "conn {}: hostkey.resolve failed: {}, reject (safe default)",
+                        self.shared.conn_id, e
+                    );
+                    false
+                }
+            };
+
+            if !accepted {
+                warn!(
+                    "conn {}: server key rejected by Ruby",
+                    self.shared.conn_id
+                );
+            }
+            Ok(accepted)
+        }
+    }
+
+    fn data(
+        &mut self,
+        channel: ChannelId,
+        data: &[u8],
+        _session: &mut client::Session,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        async move {
+            tracing::trace!(
+                "conn {} channel {}: {} bytes data",
+                self.shared.conn_id,
+                channel,
+                data.len()
+            );
+            self.shared
+                .send_channel_data(
+                    channel,
+                    ChannelData::Stdout { data: data.to_vec() },
+                )
+                .await;
+            Ok(())
+        }
+    }
+
+    fn extended_data(
+        &mut self,
+        channel: ChannelId,
+        ext: u32,
+        data: &[u8],
+        _session: &mut client::Session,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        async move {
+            tracing::trace!(
+                "conn {} channel {}: {} bytes extended data (ext={})",
+                self.shared.conn_id,
+                channel,
+                data.len(),
+                ext
+            );
+            self.shared
+                .send_channel_data(
+                    channel,
+                    ChannelData::Extended { data: data.to_vec(), ext },
+                )
+                .await;
+            Ok(())
+        }
+    }
+
+    fn channel_close(
+        &mut self,
+        channel: ChannelId,
+        _session: &mut client::Session,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        async move {
+            debug!("conn {} channel {}: close", self.shared.conn_id, channel);
+            self.shared
+                .send_channel_data(
+                    channel,
+                    ChannelData::Eof { reason: "channel_closed".to_string() },
+                )
+                .await;
+            Ok(())
+        }
+    }
+
+    fn channel_eof(
+        &mut self,
+        channel: ChannelId,
+        _session: &mut client::Session,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        async move {
+            debug!("conn {} channel {}: eof", self.shared.conn_id, channel);
+            self.shared
+                .send_channel_data(
+                    channel,
+                    ChannelData::Eof { reason: "eof".to_string() },
+                )
+                .await;
+            Ok(())
+        }
+    }
+
+    fn disconnected(
+        &mut self,
+        reason: russh::client::DisconnectReason<Self::Error>,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
+        async move {
+            warn!("conn {} disconnected by remote: {:?}", self.shared.conn_id, reason);
+            let _ = self.shared.conn_event_tx.send(ConnEvent::Disconnected);
+            match reason {
+                russh::client::DisconnectReason::ReceivedDisconnect(_) => Ok(()),
+                russh::client::DisconnectReason::Error(e) => Err(e),
+            }
+        }
+    }
+}