rss/src/daemon.rs

use std::collections::HashMap;
use std::fs;
use std::io::{BufRead, BufReader, Write};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::{UnixListener, UnixStream};
use std::os::unix::process::CommandExt;
use std::process::{Child, Command, Stdio};
use std::time::{Duration, SystemTime};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use crate::config::{Config, RestartPolicy, ServiceConfig};
use crate::state;

static STOP: AtomicBool = AtomicBool::new(false);

extern "C" fn handle_sigterm(_: i32) {
    STOP.store(true, Ordering::Relaxed);
}

struct Proc {
    child: Child,
    pid: u32,
    cfg: ServiceConfig,
    restart_count: u32,
    stdin: Option<Arc<Mutex<std::process::ChildStdin>>>,
}

pub fn run(config_path: String) {
    // Signal handler for graceful shutdown
    let handler = nix::sys::signal::SigHandler::Handler(handle_sigterm);
    let sig_action = nix::sys::signal::SigAction::new(
        handler,
        nix::sys::signal::SaFlags::SA_RESTART,
        nix::sys::signal::SigSet::empty(),
    );
    unsafe {
        let _ = nix::sys::signal::sigaction(nix::sys::signal::Signal::SIGTERM, &sig_action);
        let _ = nix::sys::signal::sigaction(nix::sys::signal::Signal::SIGINT, &sig_action);
    }

    // Write daemon PID
    let pid_path = state::daemon_pid_path();
    if let Err(e) = state::write_pid(&pid_path) {
        eprintln!("error: write PID at {:?}: {}", pid_path, e);
        return;
    }

    // Ensure socket path is clean
    let sock_path = state::socket_path();
    if sock_path.exists() {
        let _ = fs::remove_file(&sock_path);
    }
    if let Some(parent) = sock_path.parent() {
        if let Err(e) = fs::create_dir_all(parent) {
            eprintln!("error: create {}: {}", parent.display(), e);
            let _ = fs::remove_file(&pid_path);
            return;
        }
        if let Err(e) = fs::set_permissions(parent, fs::Permissions::from_mode(0o770)) {
            eprintln!("error: chmod {}: {}", parent.display(), e);
            let _ = fs::remove_file(&pid_path);
            return;
        }
    }

    // Load config
    let mut config = match Config::load(&config_path) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("error: config: {}", e);
            let _ = fs::remove_file(&pid_path);
            return;
        }
    };

    let listener = match UnixListener::bind(&sock_path) {
        Ok(l) => l,
        Err(e) => {
            eprintln!("error: bind socket at {:?}: {}", sock_path, e);
            let _ = fs::remove_file(&pid_path);
            return;
        }
    };
    let _ = fs::set_permissions(&sock_path, fs::Permissions::from_mode(0o770));
    listener.set_nonblocking(true).ok();

    let mut procs: HashMap<String, Proc> = HashMap::new();
    let mut last_mtime: Option<SystemTime> = config_mtime(&config_path);

    // Start enabled services
    start_enabled(&config, &mut procs);

    // Main loop
    loop {
        if STOP.load(Ordering::Relaxed) {
            stop_all(&mut procs);
            let _ = std::fs::remove_file(&sock_path);
            let _ = std::fs::remove_file(&pid_path);
            return;
        }

        // Poll config for changes
        if let Some(mtime) = config_mtime(&config_path) {
            if Some(mtime) != last_mtime {
                last_mtime = Some(mtime);
                let new_config = Config::load(&config_path);
                if let Ok(new_config) = new_config {
                    reconcile(&new_config, &mut procs);
                    config = new_config;
                }
            }
        }

        // Check children
        check_children(&mut procs);

        // Accept socket connections
        while let Ok((stream, _)) = listener.accept() {
            handle_client(&stream, &config, &mut procs, &config_path);
        }

        std::thread::sleep(Duration::from_millis(200));
    }
}

fn config_mtime(path: &str) -> Option<SystemTime> {
    std::fs::metadata(path).ok().and_then(|m| m.modified().ok())
}

fn start_enabled(config: &Config, procs: &mut HashMap<String, Proc>) {
    if let Ok(enabled) = crate::state::list_enabled() {
        for name in &enabled {
            if let Some(cfg) = config.service.get(name) {
                if !procs.contains_key(name) {
                    if let Some(p) = spawn(name, cfg) {
                        procs.insert(name.clone(), p);
                    }
                }
            }
        }
    }
}

fn reconcile(new_config: &Config, procs: &mut HashMap<String, Proc>) {
    let old_names: Vec<String> = procs.keys().cloned().collect();

    for name in &old_names {
        if !new_config.service.contains_key(name) {
            stop(name, procs);
        }
    }

    for (name, cfg) in &new_config.service {
        if !procs.contains_key(name) && crate::state::is_enabled(name) {
            if let Some(p) = spawn(name, cfg) {
                procs.insert(name.clone(), p);
            }
        }
    }
}

fn spawn(name: &str, cfg: &ServiceConfig) -> Option<Proc> {
    let log_path = crate::state::log_path(name);
    if let Some(parent) = log_path.parent() {
        let _ = std::fs::create_dir_all(parent);
    }
    let log_path = log_path.to_string_lossy().to_string();

    let cmd = cfg.resolve_command(name, &log_path);
    let mut child = {
        let mut sh = Command::new("sh");
        sh.arg("-c")
            .arg(&cmd)
            .current_dir(&cfg.workdir)
            .envs(&cfg.env)
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .stdin(Stdio::piped());
        // SAFETY: pre_exec runs in the child after fork; minimal ops only
        unsafe {
            sh.pre_exec(|| {
                // Parent died before we could set PDEATHSIG — nothing left to do
                if libc::getppid() == 1 {
                    libc::_exit(1);
                }
                libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM);
                Ok(())
            });
        }
        match sh.spawn() {
            Ok(c) => c,
            Err(e) => {
                append_log(&log_path, &format!("E: spawn failed: {}", e));
                return None;
            }
        }
    };

    let pid = child.id();
    let child_stdin = child.stdin.take().map(|s| Arc::new(Mutex::new(s)));
    let log_out = log_path.clone();
    let log_err = log_path.clone();

    if let Some(stdout) = child.stdout.take() {
        std::thread::spawn(move || {
            for line in BufReader::new(stdout).lines() {
                if let Ok(l) = line {
                    append_log(&log_out, &format!("O: {}", l));
                }
            }
        });
    }
    if let Some(stderr) = child.stderr.take() {
        std::thread::spawn(move || {
            for line in BufReader::new(stderr).lines() {
                if let Ok(l) = line {
                    append_log(&log_err, &format!("E: {}", l));
                }
            }
        });
    }

    append_log(&log_path, &format!("O: started (pid {})", pid));

    Some(Proc {
        child,
        pid,
        cfg: cfg.clone(),
        restart_count: 0,
        stdin: child_stdin,
    })
}

fn stop(name: &str, procs: &mut HashMap<String, Proc>) {
    let p = match procs.get(name) {
        Some(p) => p,
        None => return,
    };

    let pid = p.pid;
    let log_path = crate::state::log_path(name).to_string_lossy().to_string();

    // stop_command
    if let Some(cmd) = p.cfg.resolve_stop_command(name, pid, &log_path) {
        let _ = Command::new("sh").arg("-c").arg(&cmd).output();
    }

    // signal
    let sig = crate::supervisor::parse_signal(&p.cfg.stop_signal);
    let _ = nix::sys::signal::kill(nix::unistd::Pid::from_raw(pid as i32), sig);

    // wait
    let deadline = std::time::Instant::now() + Duration::from_secs(p.cfg.timeout_stop);
    while std::time::Instant::now() < deadline {
        if let Ok(Some(_)) = procs.get_mut(name).unwrap().child.try_wait() {
            break;
        }
        std::thread::sleep(Duration::from_millis(100));
    }

    // force kill
    if let Ok(None) = procs.get_mut(name).unwrap().child.try_wait() {
        let _ = nix::sys::signal::kill(
            nix::unistd::Pid::from_raw(pid as i32),
            nix::sys::signal::Signal::SIGKILL,
        );
        let _ = procs.get_mut(name).unwrap().child.wait();
    } else {
        let _ = procs.get_mut(name).unwrap().child.wait();
    }

    append_log(&log_path, "O: stopped");
    procs.remove(name);
}

fn stop_all(procs: &mut HashMap<String, Proc>) {
    let names: Vec<String> = procs.keys().cloned().collect();
    for name in &names {
        if procs.contains_key(name) {
            stop(name, procs);
        }
    }
}

fn check_children(procs: &mut HashMap<String, Proc>) {
    let mut dead = Vec::new();
    let mut restarts = Vec::new();

    for (name, p) in procs.iter_mut() {
        if let Ok(Some(status)) = p.child.try_wait() {
            let log_path = crate::state::log_path(name).to_string_lossy().to_string();

            let code = status
                .code()
                .map(|c| format!("code {}", c))
                .unwrap_or_else(|| "signal".to_string());

            append_log(&log_path, &format!("O: exited with {}", code));

            if p.cfg.oneshot {
                dead.push(name.clone());
                continue;
            }

            let should = match p.cfg.restart {
                RestartPolicy::No => false,
                RestartPolicy::Always | RestartPolicy::UnlessStopped => true,
                RestartPolicy::OnFailure => !status.success(),
            };

            let max = if p.cfg.restart_max == 0 {
                u32::MAX
            } else {
                p.cfg.restart_max
            };

            if should && p.restart_count < max {
                p.restart_count += 1;
                let max_str = if p.cfg.restart_max == 0 {
                    "unlimited".to_string()
                } else {
                    p.cfg.restart_max.to_string()
                };
                append_log(
                    &log_path,
                    &format!(
                        "O: restarting in {}s (attempt {}/{})",
                        p.cfg.restart_delay,
                        p.restart_count,
                        max_str
                    ),
                );
                restarts.push((name.clone(), p.cfg.clone()));
                dead.push(name.clone());
            } else {
                if p.cfg.restart != RestartPolicy::No {
                    append_log(&log_path, "O: max restarts reached, exiting");
                }
                dead.push(name.clone());
            }
        }
    }

    for name in &dead {
        procs.remove(name);
    }
    for (name, cfg) in &restarts {
        std::thread::sleep(Duration::from_secs(cfg.restart_delay));
        if let Some(p) = spawn(name, cfg) {
            procs.insert(name.clone(), p);
        }
    }
}

fn handle_client(
    mut stream: &UnixStream,
    config: &Config,
    procs: &mut HashMap<String, Proc>,
    _config_path: &str,
) {
    let mut line = String::new();
    let mut reader = BufReader::new(stream);
    if reader.read_line(&mut line).is_err() {
        return;
    }
    let line = line.trim().to_string();

    if line.starts_with("stdin ") {
        drop(reader);
        handle_stdin(&line, stream, procs);
        return;
    }
    if line.starts_with("signal ") {
        let rest = &line[7..];
        let response = if let Some(space) = rest.find(' ') {
            let name = &rest[..space];
            let sig_name = &rest[space + 1..];
            if let Some(p) = procs.get(name) {
                let sig = crate::supervisor::parse_signal(sig_name);
                match nix::sys::signal::kill(nix::unistd::Pid::from_raw(p.pid as i32), sig) {
                    Ok(_) => "ok".to_string(),
                    Err(e) => format!("error: signal failed: {}", e),
                }
            } else {
                format!("error: '{}' not running", name)
            }
        } else {
            "error: usage: signal <name> <signal>".to_string()
        };
        let _ = writeln!(stream, "{}", response);
        let _ = stream.flush();
        return;
    }

    let response = match line.as_str() {
        l if l.starts_with("start ") => {
            let names: Vec<&str> = l[6..].split_whitespace().collect();
            let mut out = Vec::new();
            for name in &names {
                if !config.service.contains_key(*name) {
                    out.push(format!("error: service '{}' not found", name));
                } else if procs.contains_key(*name) {
                    out.push(format!("error: '{}' already running", name));
                } else {
                    match spawn(name, &config.service[*name]) {
                        Some(p) => {
                            let pid = p.pid;
                            procs.insert(name.to_string(), p);
                            out.push(format!("ok {} pid={}", name, pid));
                        }
                        None => out.push(format!("error: failed to spawn '{}'", name)),
                    }
                }
            }
            out.join("\n")
        }
        l if l.starts_with("stop ") => {
            let names: Vec<&str> = l[5..].split_whitespace().collect();
            let mut out = Vec::new();
            for name in &names {
                if !procs.contains_key(*name) {
                    out.push(format!("error: '{}' not running", name));
                } else {
                    stop(name, procs);
                    out.push(format!("ok {}", name));
                }
            }
            out.join("\n")
        }
        l if l.starts_with("restart ") => {
            let names: Vec<&str> = l[8..].split_whitespace().collect();
            let mut out = Vec::new();
            for name in &names {
                if procs.contains_key(*name) {
                    stop(name, procs);
                }
                if let Some(cfg) = config.service.get(*name) {
                    match spawn(name, cfg) {
                        Some(p) => {
                            let pid = p.pid;
                            procs.insert(name.to_string(), p);
                            out.push(format!("ok {} pid={}", name, pid));
                        }
                        None => out.push(format!("error: failed to spawn '{}'", name)),
                    }
                } else {
                    out.push(format!("error: service '{}' not found", name));
                }
            }
            out.join("\n")
        }
        "list" | "list " => {
            let mut out = Vec::new();
            for (name, _cfg) in &config.service {
                let enabled = if crate::state::is_enabled(name) { "yes" } else { "no" };
                let info = if let Some(p) = procs.get(name) {
                    format!("running {} pid={} enabled={}", name, p.pid, enabled)
                } else {
                    format!("stopped {} pid=0 enabled={}", name, enabled)
                };
                out.push(info);
            }
            out.join("\n")
        }
        l if l.starts_with("status ") => {
            let name = &l[7..];
            let enabled = if crate::state::is_enabled(name) { "yes" } else { "no" };
            if let Some(p) = procs.get(name) {
                format!("running {} pid={} enabled={}", name, p.pid, enabled)
            } else if config.service.contains_key(name) {
                format!("stopped {} pid=0 enabled={}", name, enabled)
            } else {
                format!("error: service '{}' not found", name)
            }
        }
        "status" => {
            let mut out = Vec::new();
            for (name, _) in &config.service {
                let enabled = if crate::state::is_enabled(name) { "yes" } else { "no" };
                let info = if let Some(p) = procs.get(name) {
                    format!("running {} pid={} enabled={}", name, p.pid, enabled)
                } else {
                    format!("stopped {} pid=0 enabled={}", name, enabled)
                };
                out.push(info);
            }
            out.join("\n")
        }
        l if l.starts_with("enable ") => {
            let rest = &l[7..];
            let (names_str, now) = if let Some(n) = rest.strip_suffix(" now") {
                (n, true)
            } else {
                (rest, false)
            };
            let names: Vec<&str> = names_str.split_whitespace().collect();
            let mut out = Vec::new();
            for name in &names {
                if !config.service.contains_key(*name) {
                    out.push(format!("error: service '{}' not found", name));
                } else if let Err(e) = crate::state::enable(name) {
                    out.push(format!("error: enable '{}': {}", name, e));
                } else {
                    if now {
                        if !procs.contains_key(*name) {
                            if let Some(p) = spawn(name, &config.service[*name]) {
                                procs.insert(name.to_string(), p);
                            }
                        }
                    }
                    out.push(format!("ok {}", name));
                }
            }
            out.join("\n")
        }
        l if l.starts_with("disable ") => {
            let names: Vec<&str> = l[8..].split_whitespace().collect();
            let mut out = Vec::new();
            for name in &names {
                crate::state::disable(name);
                out.push(format!("ok {}", name));
            }
            out.join("\n")
        }
        _ => format!("error: unknown command '{}'", line),
    };

    let _ = writeln!(stream, "{}", response);
    let _ = stream.flush();
}

fn handle_stdin(line: &str, mut stream: &UnixStream, procs: &mut HashMap<String, Proc>) {
    let rest = &line[6..];
    if let Some(space) = rest.find(' ') {
        let name = &rest[..space];
        let msg = &rest[space + 1..];
        match procs.get(name) {
            Some(p) => match &p.stdin {
                Some(stdin) => {
                    let mut stdin = stdin.lock().unwrap();
                    let _ = writeln!(stdin, "{}", msg);
                    let _ = stdin.flush();
                    let _ = writeln!(stream, "ok");
                }
                None => {
                    let _ = writeln!(stream, "error: '{}' has no stdin", name);
                }
            },
            None => {
                let _ = writeln!(stream, "error: '{}' not running", name);
            }
        }
        let _ = stream.flush();
    } else {
        let name = rest;
        let has_stdin = procs.get(name).and_then(|p| p.stdin.as_ref()).is_some();
        if !procs.contains_key(name) {
            let _ = writeln!(stream, "error: '{}' not running", name);
            let _ = stream.flush();
            return;
        }
        if !has_stdin {
            let _ = writeln!(stream, "error: '{}' has no stdin", name);
            let _ = stream.flush();
            return;
        }
        let stdin_arc = procs.get(name).and_then(|p| p.stdin.clone()).unwrap();
        match stream.try_clone() {
            Ok(clone) => {
                let _ = writeln!(stream, "ok");
                let _ = stream.flush();
                std::thread::spawn(move || {
                    let mut reader = BufReader::new(clone);
                    let mut line = String::new();
                    loop {
                        line.clear();
                        match reader.read_line(&mut line) {
                            Ok(0) => break,
                            Ok(_) => {
                                let t = line.trim();
                                if !t.is_empty() {
                                    let mut s = match stdin_arc.lock() {
                                        Ok(s) => s,
                                        Err(_) => break,
                                    };
                                    if writeln!(s, "{}", t).is_err() {
                                        break;
                                    }
                                    let _ = s.flush();
                                }
                            }
                            Err(_) => break,
                        }
                    }
                });
            }
            Err(e) => {
                let _ = writeln!(stream, "error: {}", e);
                let _ = stream.flush();
            }
        }
    }
}

fn append_log(path: &str, msg: &str) {
    let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S");
    if let Ok(mut f) = std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(path)
    {
        let _ = writeln!(f, "[{}] {}", ts, msg);
        let _ = f.flush();
    }
}