sinit 0.1.2

//! # Simple Init
//!
//! Simple Init (sinit) is a small init system that runs as PID 1 and proxies
//! signals to its child processes.
//!
//! ## Usage
//!
//!   ./sinit python -m http.server
//!
//! ## Signal Handling
//!
//! Simple Init proxies any signals received to its child process(es).
//! In setsid mode, signals are forwarded to the entire process tree managed by
//! sinit. In non-setsid mode signals are only proxied to the first child.
//!
//! If a job control signal is received, sinit itself goes to sleep. This allows
//! it to manage foreground processes started by users without messing up their
//! shell (eg. they can still use ^C to kill the process).
//!
//! From Yelp's dumb-init comments:
//!
//! > When running in setsid mode, however, it is not sufficient to forward
//! > SIGTSTP/SIGTTIN/SIGTTOU in most cases. If the process has not added a custom
//! > signal handler for these signals, then the kernel will not apply default
//! > signal handling behavior (which would be suspending the process) since it is
//! > a member of an orphaned process group.
//! >
//! > Sadly this doesn't appear to be well documented except in the kernel itself:
//! > https://github.com/torvalds/linux/blob/v4.2/kernel/signal.c#L2296-L2299
//! >
//! > Forwarding SIGSTOP instead is effective, though not ideal; unlike SIGTSTP,
//! > SIGSTOP cannot be caught, and so it doesn't allow processes a change to
//! > clean up before suspending. In non-setsid mode, we proxy the original signal
//! > instead of SIGSTOP for this reason.

#![deny(missing_docs)]
#![doc(html_root_url = "https://docs.rs/sinit/0.1.2")]

use std::{env, mem};
use std::process::Command;
use std::os::unix::process::CommandExt;

extern crate getopts;
extern crate libc;

macro_rules! debugln {
    ($fmt:expr) => (println!(concat!("DEBUG:", $fmt, "\n")));
    ($fmt:expr, $($arg:tt)*) => (print!(concat!("DEBUG:", $fmt, "\n"), $($arg)*));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "Print this help menu.");
    opts.optflag("c",
                 "single-child",
                 "Run in single-child mode.
                  In this mode, signals are only proxied to the direct child and not to any of its
                  decendants.");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let sid = matches.opt_present("c");
    let cmd = if matches.free.is_empty() {
        print_usage(&program, opts);
        return;
    } else {
        matches.free[0].clone()
    };

    // TODO: Cleanup: https://github.com/rust-lang/rust/issues/15023
    let mut command = Command::new(cmd);
    let mut command = command.before_exec(move || {
        if sid {
            unsafe {
                let _ = libc::setsid();
            }
        }
        Ok(())
    });

    let mut command = if matches.free.len() > 1 {
        command.args(&matches.free[1..])
    } else {
        command
    };

    // Here there be dragons.
    unsafe {
        let mut sigset: libc::sigset_t = mem::uninitialized();
        let mut oldsigset: libc::sigset_t = mem::uninitialized();
        libc::sigfillset(&mut sigset);

        libc::pthread_sigmask(libc::SIG_BLOCK,
                              &mut sigset as *const libc::sigset_t,
                              &mut oldsigset as *mut libc::sigset_t);

        let child = command.spawn()
            .unwrap_or_else(|e| panic!("failed to execute child: {}", e));

        // TODO: Why does child.id() return a u32 but Linux and Command think PIDs are i32?
        let pid = child.id() as i32;
        println!("Child process spawned with pid {}", pid);

        let mut signum: i32 = 0;
        loop {
            libc::sigwait(&sigset as *const libc::sigset_t, &mut signum);
            handle_signal(matches.opt_present("c"), pid, signum);
        }
    }
}

fn print_usage(program: &str, opts: getopts::Options) {
    let brief = format!("Usage: {} [options] CMD", program);
    print!("{}", opts.usage(&brief));
}

unsafe fn forward_signal(setsid: bool, pid: i32, sig: i32) {
    libc::kill(if setsid { -pid } else { pid }, sig);
    debugln!("Forwarded signal {} to pid {}", sig, pid);
}

unsafe fn handle_signal(setsid: bool, pid: i32, sig: i32) {
    debugln!("Received signal {}", sig);
    match sig {
        libc::SIGCHLD => {
            let mut status = 0;
            let raw_status = &mut status as *mut i32;

            let mut killed_pid;
            loop {
                killed_pid = libc::waitpid(-1, raw_status, libc::WNOHANG);
                if killed_pid < 1 {
                    break;
                }

                let exit_status;
                // TODO: Only the first branch will ever be hit here because we didn't call fork(3) on
                // the child directly (so we're not the chlid that just exited).
                if libc::WIFEXITED(status) {
                    exit_status = libc::WEXITSTATUS(status);
                    debugln!("A child with PID {} exited with exit status {}.",
                             killed_pid,
                             exit_status);
                } else {
                    // TODO: assert!(libc::WIFSIGNALED(status));
                    exit_status = 128 + libc::WTERMSIG(status);
                    debugln!("A child with PID {} was terminated by signal {}.",
                             killed_pid,
                             exit_status - 128);
                }

                if killed_pid == pid {
                    // send SIGTERM to any remaining children
                    forward_signal(setsid, pid, libc::SIGTERM);
                    debugln!("Child exited with status {}. Goodbye.", exit_status);
                    std::process::exit(exit_status);
                }
            }
        }
        libc::SIGTSTP | libc::SIGTTIN | libc::SIGTTOU => {
            if setsid {
                debugln!("Replacing signal with SIGSTOP in setsid mode.");
                forward_signal(setsid, pid, libc::SIGSTOP);
            } else {
                forward_signal(setsid, pid, sig);
            }

            debugln!("Suspending self due to TTY signal.");
            forward_signal(setsid, libc::getpid(), libc::SIGSTOP);
        }
        _ => {
            forward_signal(setsid, pid, sig);
        }
    }
}