Genarrative/server-rs/crates/server-manager-panel/src/remote.rs

use std::io::Write;
use std::process::{Command, Stdio};
use std::sync::mpsc;
use std::thread;
use std::time::{Duration, Instant};

use crate::health::{HEALTH_SCRIPT, ServerHealthReport, parse_health_report};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ServiceAction {
    Start,
    Stop,
    Restart,
}

impl ServiceAction {
    pub fn as_systemctl_arg(self) -> &'static str {
        match self {
            ServiceAction::Start => "start",
            ServiceAction::Stop => "stop",
            ServiceAction::Restart => "restart",
        }
    }

    pub fn label(self) -> &'static str {
        match self {
            ServiceAction::Start => "启动",
            ServiceAction::Stop => "关闭",
            ServiceAction::Restart => "重启",
        }
    }
}

#[derive(Debug, Clone)]
pub struct RemoteCommandResult {
    pub success: bool,
    pub summary: String,
    pub stdout: String,
    pub stderr: String,
}

#[derive(Debug)]
pub enum RemoteEvent {
    Health {
        alias: String,
        result: Result<ServerHealthReport, String>,
    },
    ServiceAction {
        alias: String,
        service: String,
        action: ServiceAction,
        result: RemoteCommandResult,
    },
}

pub type RemoteSender = mpsc::Sender<RemoteEvent>;
pub type RemoteReceiver = mpsc::Receiver<RemoteEvent>;

pub fn channel() -> (RemoteSender, RemoteReceiver) {
    mpsc::channel()
}

pub fn spawn_health_check(alias: String, tx: RemoteSender) {
    thread::spawn(move || {
        let result =
            run_ssh_script(&alias, HEALTH_SCRIPT, Duration::from_secs(20)).and_then(|output| {
                if output.success {
                    Ok(parse_health_report(&output.stdout))
                } else {
                    Err(format_remote_error(&output))
                }
            });
        let _ = tx.send(RemoteEvent::Health { alias, result });
    });
}

pub fn spawn_service_action(
    alias: String,
    service: String,
    action: ServiceAction,
    tx: RemoteSender,
) {
    thread::spawn(move || {
        let result = if is_safe_service_name(&service) {
            run_ssh_script(
                &alias,
                &build_service_action_script(&service, action),
                Duration::from_secs(20),
            )
            .unwrap_or_else(|error| RemoteCommandResult {
                success: false,
                summary: error,
                stdout: String::new(),
                stderr: String::new(),
            })
        } else {
            RemoteCommandResult {
                success: false,
                summary: "服务名包含不允许的字符".to_owned(),
                stdout: String::new(),
                stderr: String::new(),
            }
        };
        let _ = tx.send(RemoteEvent::ServiceAction {
            alias,
            service,
            action,
            result,
        });
    });
}

pub fn is_safe_service_name(service: &str) -> bool {
    !service.is_empty()
        && service.len() <= 128
        && service.bytes().all(|byte| {
            matches!(
                byte,
                b'a'..=b'z'
                    | b'A'..=b'Z'
                    | b'0'..=b'9'
                    | b'.'
                    | b'_'
                    | b'-'
                    | b'@'
                    | b':'
            )
        })
}

fn build_service_action_script(service: &str, action: ServiceAction) -> String {
    format!(
        r#"set -eu
service='{service}'
action='{action}'
if [ "$(id -u)" = "0" ]; then
  systemctl "$action" "$service"
else
  sudo -n systemctl "$action" "$service"
fi
systemctl is-active "$service" || true
systemctl status "$service" --no-pager -l -n 12 || true
"#,
        service = service,
        action = action.as_systemctl_arg()
    )
}

fn run_ssh_script(
    alias: &str,
    script: &str,
    timeout: Duration,
) -> Result<RemoteCommandResult, String> {
    let started = Instant::now();
    let mut child = Command::new("ssh")
        .arg("-o")
        .arg("BatchMode=yes")
        .arg("-o")
        .arg("ConnectTimeout=8")
        .arg(alias)
        .arg("sh")
        .arg("-s")
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .map_err(|error| format!("无法启动 ssh: {error}"))?;

    {
        // 中文注释：写完脚本后必须关闭 stdin，让远端 `sh -s` 收到 EOF 并开始退出。
        let Some(mut stdin) = child.stdin.take() else {
            return Err("无法写入 ssh stdin".to_owned());
        };
        stdin
            .write_all(script.as_bytes())
            .map_err(|error| format!("写入远端脚本失败: {error}"))?;
    }

    loop {
        match child.try_wait() {
            Ok(Some(_status)) => {
                let output = child
                    .wait_with_output()
                    .map_err(|error| format!("读取 ssh 输出失败: {error}"))?;
                let success = output.status.success();
                return Ok(RemoteCommandResult {
                    success,
                    summary: if success {
                        "执行成功".to_owned()
                    } else {
                        format!("ssh 退出码 {:?}", output.status.code())
                    },
                    stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
                    stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
                });
            }
            Ok(None) if started.elapsed() >= timeout => {
                let _ = child.kill();
                let _ = child.wait();
                return Err(format!("ssh 执行超过 {} 秒", timeout.as_secs()));
            }
            Ok(None) => thread::sleep(Duration::from_millis(80)),
            Err(error) => return Err(format!("等待 ssh 进程失败: {error}")),
        }
    }
}

fn format_remote_error(result: &RemoteCommandResult) -> String {
    let stderr = result.stderr.trim();
    let stdout = result.stdout.trim();
    if !stderr.is_empty() {
        format!("{}: {}", result.summary, stderr)
    } else if !stdout.is_empty() {
        format!("{}: {}", result.summary, stdout)
    } else {
        result.summary.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn allows_systemd_unit_names_only() {
        assert!(is_safe_service_name("genarrative-api.service"));
        assert!(is_safe_service_name("worker@1.service"));
        assert!(!is_safe_service_name("api.service;rm -rf /"));
        assert!(!is_safe_service_name(""));
    }
}