zinit 0.3.6

//! zinit-pid1 - The init process (PID 1).
//!
//! Minimal init that spawns and monitors zinit-server, reaps zombies,
//! forwards signals, and handles system shutdown/reboot.
//!
//! **Critical:** On VM/bare-metal, PID 1 must NEVER exit (causes kernel panic).
//! Only containers can safely exit.
//!
//! **Note:** This crate only compiles on Linux. On other platforms, it produces
//! a stub binary that prints an error message.

#[cfg(not(target_os = "linux"))]
fn main() {
    eprintln!("zinit-pid1 is only supported on Linux");
    std::process::exit(1);
}

#[cfg(target_os = "linux")]
mod linux {
    use std::ffi::CString;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::thread;
    use std::time::Duration;

    use clap::Parser;
    use nix::sys::reboot::{RebootMode, reboot, set_cad_enabled};
    use nix::sys::signal::{self, SaFlags, SigAction, SigHandler, SigSet, Signal};
    use nix::sys::wait::{WaitPidFlag, WaitStatus, waitpid};
    use nix::unistd::{ForkResult, Pid, execv, fork, getpid};

    use zinit::sdk::ZinitClient;

    /// zinit-pid1 - Init process that spawns and monitors zinit-server.
    #[derive(Parser, Debug)]
    #[command(name = "zinit-pid1")]
    #[command(version, about, long_about = None)]
    struct Args {
        /// Run in container mode (exit cleanly instead of reboot syscall).
        #[arg(short, long)]
        container: bool,
    }

    // Signal flags - set by signal handler, read by main loop
    static SIGTERM_RECEIVED: AtomicBool = AtomicBool::new(false);
    static SIGINT_RECEIVED: AtomicBool = AtomicBool::new(false);
    static SIGUSR1_RECEIVED: AtomicBool = AtomicBool::new(false);
    static SIGUSR2_RECEIVED: AtomicBool = AtomicBool::new(false);
    static SIGCHLD_RECEIVED: AtomicBool = AtomicBool::new(false);

    /// Shutdown mode determined by signal received.
    #[derive(Debug, Clone, Copy, PartialEq)]
    enum ShutdownMode {
        None,
        Reboot,   // SIGINT
        Poweroff, // SIGTERM
    }

    pub fn run() {
        let args = Args::parse();

        // Initialize logging with optional level from env
        let log_level = std::env::var("ZINIT_LOG_LEVEL").unwrap_or_else(|_| "info".to_string());
        let filter = tracing_subscriber::EnvFilter::try_new(&log_level)
            .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));

        tracing_subscriber::fmt()
            .with_env_filter(filter)
            .with_target(false)
            .with_level(true)
            .init();

        let is_pid1 = getpid().as_raw() == 1;
        if !is_pid1 {
            tracing::warn!("Not running as PID 1 (pid={})", getpid());
        }

        // Disable immediate reboot on Ctrl+Alt+Del - make kernel send SIGINT to PID 1 instead
        if is_pid1 {
            if let Err(e) = set_cad_enabled(false) {
                tracing::warn!(error = %e, "Failed to disable Ctrl+Alt+Del immediate reboot");
            } else {
                tracing::debug!("Disabled Ctrl+Alt+Del immediate reboot (kernel will send SIGINT)");
            }
        }

        let container_mode = args.container || is_container();
        if container_mode {
            tracing::info!("Running in container mode");
        }

        setup_signals();

        // Determine if we should run in PID1 mode (bare-metal/VM, not container)
        let pid1_server_mode = is_pid1 && !container_mode;
        if pid1_server_mode {
            tracing::info!(
                "Will start zinit-server with --pid1-mode (system services from /etc/zinit/system)"
            );
        }

        let mut server_pid = spawn_server(pid1_server_mode);

        #[allow(unused_assignments)]
        let mut shutdown_mode = ShutdownMode::None;

        // Main loop
        loop {
            // Handle SIGCHLD - reap zombies
            if SIGCHLD_RECEIVED.swap(false, Ordering::SeqCst) {
                server_pid = reap_zombies(server_pid);
            }

            // Handle SIGTERM - poweroff
            if SIGTERM_RECEIVED.swap(false, Ordering::SeqCst) {
                tracing::info!("Received SIGTERM, initiating poweroff");
                shutdown_mode = ShutdownMode::Poweroff;
                break;
            }

            // Handle SIGINT - reboot
            if SIGINT_RECEIVED.swap(false, Ordering::SeqCst) {
                tracing::info!("Received SIGINT, initiating reboot");
                shutdown_mode = ShutdownMode::Reboot;
                break;
            }

            // Handle SIGUSR1 - soft restart zinit-server
            if SIGUSR1_RECEIVED.swap(false, Ordering::SeqCst) {
                tracing::info!("Received SIGUSR1, soft-restarting zinit-server");
                soft_restart_server(&mut server_pid, pid1_server_mode);
            }

            // Handle SIGUSR2 - self-update (placeholder)
            if SIGUSR2_RECEIVED.swap(false, Ordering::SeqCst) {
                tracing::info!("Received SIGUSR2, checking for self-update");
                handle_self_update();
            }

            // Respawn server if it died
            if server_pid.is_none() {
                tracing::error!("zinit-server died, respawning in 1s...");
                thread::sleep(Duration::from_secs(1));
                server_pid = spawn_server(pid1_server_mode);
            }

            // Sleep to avoid busy-looping
            thread::sleep(Duration::from_millis(100));
        }

        do_shutdown(server_pid, shutdown_mode, container_mode);
    }

    /// Signal handler - just sets atomic flags.
    /// Must be async-signal-safe (no allocations, no complex operations).
    extern "C" fn signal_handler(sig: i32) {
        match sig {
            libc::SIGTERM => SIGTERM_RECEIVED.store(true, Ordering::SeqCst),
            libc::SIGINT => SIGINT_RECEIVED.store(true, Ordering::SeqCst),
            libc::SIGUSR1 => SIGUSR1_RECEIVED.store(true, Ordering::SeqCst),
            libc::SIGUSR2 => SIGUSR2_RECEIVED.store(true, Ordering::SeqCst),
            libc::SIGCHLD => SIGCHLD_RECEIVED.store(true, Ordering::SeqCst),
            _ => {}
        }
    }

    /// Set up signal handlers.
    fn setup_signals() {
        let handler = SigHandler::Handler(signal_handler);
        let flags = SaFlags::SA_RESTART;
        let action = SigAction::new(handler, flags, SigSet::empty());

        unsafe {
            signal::sigaction(Signal::SIGTERM, &action).expect("Failed to set SIGTERM handler");
            signal::sigaction(Signal::SIGINT, &action).expect("Failed to set SIGINT handler");
            signal::sigaction(Signal::SIGUSR1, &action).expect("Failed to set SIGUSR1 handler");
            signal::sigaction(Signal::SIGUSR2, &action).expect("Failed to set SIGUSR2 handler");
            signal::sigaction(Signal::SIGCHLD, &action).expect("Failed to set SIGCHLD handler");
        }

        tracing::debug!("Signal handlers installed");
    }

    /// Reap all zombie processes and check if server died.
    fn reap_zombies(server_pid: Option<Pid>) -> Option<Pid> {
        let mut server_alive = server_pid;

        loop {
            match waitpid(Pid::from_raw(-1), Some(WaitPidFlag::WNOHANG)) {
                Ok(WaitStatus::Exited(pid, code)) => {
                    tracing::info!(pid = pid.as_raw(), exit_code = code, "Process exited");
                    if Some(pid) == server_pid {
                        server_alive = None;
                    }
                }
                Ok(WaitStatus::Signaled(pid, sig, _)) => {
                    tracing::info!(pid = pid.as_raw(), signal = ?sig, "Process killed by signal");
                    if Some(pid) == server_pid {
                        server_alive = None;
                    }
                }
                Ok(WaitStatus::StillAlive) => break,
                Err(nix::errno::Errno::ECHILD) => break, // No children
                _ => break,
            }
        }

        server_alive
    }

    /// Spawn zinit-server as a child process.
    /// If `pid1_mode` is true, adds --pid1-mode flag to load system services.
    fn spawn_server(pid1_mode: bool) -> Option<Pid> {
        // Try multiple paths for the server binary (system paths first)
        let paths = [
            "/sbin/zinit-server",
            "/usr/sbin/zinit-server",
            "/usr/bin/zinit-server",
            "/usr/local/bin/zinit-server",
        ];

        let server_path = paths
            .iter()
            .find(|p| std::path::Path::new(p).exists())
            .copied()
            .unwrap_or("/sbin/zinit-server");

        // Check for config overrides from environment
        let config_dir = std::env::var("ZINIT_CONFIG_DIR").unwrap_or_else(|_| {
            zinit::sdk::socket::system_config_dir()
                .to_string_lossy()
                .to_string()
        });
        let socket_path = std::env::var("ZINIT_SOCKET").unwrap_or_else(|_| {
            zinit::sdk::socket::system_path()
                .to_string_lossy()
                .to_string()
        });

        // Check for debug mode from kernel cmdline (zinitdebug=1)
        let debug_mode = is_debug_mode();
        let log_level = if debug_mode {
            tracing::info!("Debug mode enabled via kernel cmdline (zinitdebug=1)");
            "debug".to_string()
        } else {
            std::env::var("ZINIT_LOG_LEVEL").unwrap_or_else(|_| "info".to_string())
        };

        match unsafe { fork() } {
            Ok(ForkResult::Child) => {
                // Child process - exec into zinit-server
                let prog = CString::new(server_path).expect("Invalid server path");
                let arg_c = CString::new("-c").unwrap();
                let arg_config = CString::new(config_dir).unwrap();
                let arg_s = CString::new("-s").unwrap();
                let arg_socket = CString::new(socket_path).unwrap();
                let arg_log = CString::new("--log-level").unwrap();
                let arg_log_level = CString::new(log_level).unwrap();

                // Build args list - conditionally include --pid1-mode
                let args: Vec<CString> = if pid1_mode {
                    let arg_pid1 = CString::new("--pid1-mode").unwrap();
                    vec![
                        prog.clone(),
                        arg_c,
                        arg_config,
                        arg_s,
                        arg_socket,
                        arg_log,
                        arg_log_level,
                        arg_pid1,
                    ]
                } else {
                    vec![
                        prog.clone(),
                        arg_c,
                        arg_config,
                        arg_s,
                        arg_socket,
                        arg_log,
                        arg_log_level,
                    ]
                };

                // execv never returns on success
                #[allow(unreachable_code)]
                match execv(&prog, &args) {
                    Ok(infallible) => match infallible {},
                    Err(e) => {
                        eprintln!("Failed to exec zinit-server: {}", e);
                        std::process::exit(1);
                    }
                }
            }
            Ok(ForkResult::Parent { child }) => {
                tracing::info!(
                    pid = child.as_raw(),
                    path = server_path,
                    pid1_mode = pid1_mode,
                    "Spawned zinit-server"
                );
                Some(child)
            }
            Err(e) => {
                tracing::error!(error = %e, "Failed to fork");
                None
            }
        }
    }

    /// Soft-restart zinit-server (for updates).
    /// Calls prepare_restart RPC to save state, then respawns.
    fn soft_restart_server(server_pid: &mut Option<Pid>, pid1_mode: bool) {
        if let Some(pid) = *server_pid {
            tracing::info!(
                pid = pid.as_raw(),
                "Initiating soft restart of zinit-server"
            );

            // Call prepare_restart RPC to save state
            if let Ok(mut client) = ZinitClient::connect_default() {
                match client.prepare_restart() {
                    Ok(result) => {
                        tracing::info!(
                            state_path = %result.state_path,
                            ready = result.ready,
                            "State saved for restart"
                        );
                    }
                    Err(e) => {
                        tracing::warn!(error = %e, "Failed to save state, continuing anyway");
                    }
                }
            } else {
                tracing::warn!("Could not connect to zinit-server, proceeding with restart");
            }

            // Wait for server to exit (it should exit after prepare_restart)
            for _ in 0..50 {
                thread::sleep(Duration::from_millis(100));
                match waitpid(pid, Some(WaitPidFlag::WNOHANG)) {
                    Ok(WaitStatus::Exited(_, code)) => {
                        tracing::info!(exit_code = code, "zinit-server exited gracefully");
                        *server_pid = None;
                        break;
                    }
                    Ok(WaitStatus::Signaled(_, sig, _)) => {
                        tracing::info!(signal = ?sig, "zinit-server killed by signal");
                        *server_pid = None;
                        break;
                    }
                    _ => {}
                }
            }

            // Force kill if still running
            if server_pid.is_some() {
                tracing::warn!(
                    pid = pid.as_raw(),
                    "zinit-server didn't exit, sending SIGKILL"
                );
                let _ = signal::kill(pid, Signal::SIGKILL);
                let _ = waitpid(pid, None);
                *server_pid = None;
            }
        }

        // Spawn new server (it will restore from saved state)
        *server_pid = spawn_server(pid1_mode);
    }

    /// Handle self-update request (placeholder).
    fn handle_self_update() {
        // TODO: Implement self-update
        // 1. Check for new zinit-pid1 binary at /usr/bin/zinit-pid1.new
        // 2. Verify checksum/signature
        // 3. Move to /usr/bin/zinit-pid1
        // 4. Re-exec self: execv("/usr/bin/zinit-pid1", ["--adopt-server", pid])
        //    This replaces process image but keeps PID 1

        tracing::warn!("Self-update not yet implemented");
    }

    /// Detect if we're running in a container via environment variable.
    fn is_container() -> bool {
        std::env::var("ZINIT_CONTAINER").is_ok()
    }

    /// Check kernel cmdline for zinitdebug=1 parameter.
    fn is_debug_mode() -> bool {
        std::fs::read_to_string("/proc/cmdline")
            .map(|cmdline| {
                cmdline
                    .split_whitespace()
                    .any(|param| param == "zinitdebug=1")
            })
            .unwrap_or(false)
    }

    /// Perform system shutdown.
    fn do_shutdown(server_pid: Option<Pid>, mode: ShutdownMode, container_mode: bool) {
        tracing::info!(mode = ?mode, "Initiating shutdown");

        // 1. Tell zinit-server to stop all services
        if let Ok(mut client) = ZinitClient::connect_default() {
            tracing::info!("Requesting zinit-server shutdown");
            let _ = client.shutdown();
        }

        // 2. Wait for server to exit (30s timeout)
        if let Some(pid) = server_pid {
            tracing::info!(pid = pid.as_raw(), "Waiting for zinit-server to exit");

            for i in 0..300 {
                thread::sleep(Duration::from_millis(100));
                match waitpid(pid, Some(WaitPidFlag::WNOHANG)) {
                    Ok(WaitStatus::Exited(_, code)) => {
                        tracing::info!(exit_code = code, "zinit-server exited");
                        break;
                    }
                    Ok(WaitStatus::Signaled(_, sig, _)) => {
                        tracing::info!(signal = ?sig, "zinit-server killed by signal");
                        break;
                    }
                    _ => {
                        if i == 299 {
                            tracing::warn!("zinit-server didn't exit in 30s, sending SIGKILL");
                            let _ = signal::kill(pid, Signal::SIGKILL);
                            let _ = waitpid(pid, None);
                        }
                    }
                }
            }
        }

        // 3. Reap all remaining zombies
        loop {
            match waitpid(Pid::from_raw(-1), Some(WaitPidFlag::WNOHANG)) {
                Err(nix::errno::Errno::ECHILD) => break,
                Ok(WaitStatus::StillAlive) => break,
                _ => continue,
            }
        }

        // 4. Sync filesystems
        tracing::info!("Syncing filesystems");
        unsafe {
            libc::sync();
        }

        // 5. Container vs bare-metal/VM
        if container_mode {
            tracing::info!("Container mode, exiting normally");
            match mode {
                ShutdownMode::Poweroff | ShutdownMode::Reboot => std::process::exit(0),
                ShutdownMode::None => std::process::exit(1),
            }
        }

        // 6. Bare-metal/VM: use reboot syscall - NEVER exit!
        match mode {
            ShutdownMode::Reboot => {
                tracing::info!("Rebooting system...");
                let _ = reboot(RebootMode::RB_AUTOBOOT);
            }
            ShutdownMode::Poweroff => {
                tracing::info!("Powering off system...");
                let _ = reboot(RebootMode::RB_POWER_OFF);
            }
            ShutdownMode::None => {
                tracing::error!("Unexpected shutdown state, rebooting as fallback");
                let _ = reboot(RebootMode::RB_AUTOBOOT);
            }
        }

        // If reboot syscall failed, loop forever (never exit as PID 1!)
        tracing::error!("Reboot syscall failed! Halting (infinite loop to prevent kernel panic)");
        loop {
            thread::sleep(Duration::from_secs(3600));
        }
    }
} // mod linux

#[cfg(target_os = "linux")]
fn main() {
    linux::run();
}