kanit

Toy init system
Log | Files | Refs | README | LICENSE

cli.rs (4423B)

      1 use std::fs::{read_to_string, write};
      2 use std::os::unix::prelude::ExitStatusExt;
      3 use std::path::Path;
      4 use std::process;
      5 use std::process::{ExitCode, ExitStatus};
      6 
      7 use nix::errno::Errno;
      8 use nix::libc::pid_t;
      9 use nix::sched::{CloneFlags, unshare};
     10 use nix::sys::signal;
     11 use nix::sys::signal::{Signal, kill};
     12 use nix::sys::signalfd::{SfdFlags, SigSet, SignalFd};
     13 use nix::sys::stat::Mode;
     14 use nix::sys::wait::{WaitPidFlag, WaitStatus, waitpid};
     15 use nix::unistd::{Pid, mkdir};
     16 
     17 use kanit_common::error::{Context, Result};
     18 
     19 use crate::{Supervisor, spawn, spawn_restart};
     20 
     21 fn cgroup_signal(name: &str, signal: Signal) -> Result<()> {
     22     let procs_path = Path::new("/sys/fs/cgroup/system.slice")
     23         .join(name)
     24         .join("cgroup.procs");
     25 
     26     read_to_string(procs_path)
     27         .context("failed to read procs")?
     28         .split('\n')
     29         .filter_map(|pid| pid.parse::<u32>().ok())
     30         .try_for_each(|pid| {
     31             if pid == process::id() {
     32                 return Ok(());
     33             }
     34 
     35             let _ = kill(Pid::from_raw(pid as pid_t), signal);
     36 
     37             Ok(())
     38         })
     39 }
     40 
     41 pub fn handle_cli() -> ExitCode {
     42     match Supervisor::from_env() {
     43         Ok(mut cfg) => {
     44             let mut mask = SigSet::empty();
     45             mask.add(signal::SIGCHLD);
     46             mask.add(signal::SIGTERM);
     47             mask.thread_block().unwrap();
     48 
     49             let sfd = SignalFd::with_flags(&mask, SfdFlags::empty()).unwrap();
     50 
     51             if let Some(ref cgroup) = cfg.cgroup {
     52                 unshare(CloneFlags::CLONE_NEWCGROUP).expect("unshare cgroup");
     53 
     54                 let supervisor_cgroup = Path::new("/sys/fs/cgroup/system.slice").join(cgroup);
     55 
     56                 mkdir(
     57                     &supervisor_cgroup,
     58                     Mode::S_IRWXU | Mode::S_IRGRP | Mode::S_IRUSR,
     59                 )
     60                 .expect("create directory");
     61 
     62                 write(
     63                     supervisor_cgroup.join("cgroup.procs"),
     64                     process::id().to_string(),
     65                 )
     66                 .expect("write cgroup proc");
     67             }
     68 
     69             #[allow(clippy::zombie_processes)] // we use `waitpid` and not `.wait`
     70             let mut child = spawn(&cfg).expect("spawn child");
     71 
     72             loop {
     73                 match sfd.read_signal() {
     74                     Ok(Some(sig)) => match Signal::try_from(sig.ssi_signo as i32) {
     75                         Ok(Signal::SIGCHLD) => {
     76                             loop {
     77                                 let pid = waitpid(None, Some(WaitPidFlag::WNOHANG));
     78 
     79                                 match pid {
     80                                     Ok(WaitStatus::StillAlive) => break,
     81                                     Err(Errno::ECHILD) => break,
     82                                     Err(e) => {
     83                                         eprintln!("failed to waitpid: {e}");
     84                                         return ExitCode::FAILURE;
     85                                     }
     86                                     _ => {}
     87                                 }
     88                             }
     89 
     90                             if let Some(c) =
     91                                 spawn_restart(&mut cfg, ExitStatus::from_raw(sig.ssi_status), true)
     92                                     .expect("restart child")
     93                             {
     94                                 child = c;
     95                             } else {
     96                                 return ExitCode::SUCCESS;
     97                             }
     98                         }
     99                         Ok(Signal::SIGTERM) => {
    100                             child.kill().expect("kill child");
    101                             if let Some(attempts) = cfg.restart_attempts {
    102                                 cfg.restart_attempts = Some(attempts + 1);
    103                             }
    104                             if let Some(ref cgroup) = cfg.cgroup {
    105                                 cgroup_signal(cgroup.as_ref(), Signal::SIGKILL)
    106                                     .expect("kill child");
    107                             }
    108                         }
    109                         _ => {}
    110                     },
    111                     Ok(None) => unreachable!(),
    112                     Err(e) => {
    113                         eprintln!("{e}");
    114                         return ExitCode::FAILURE;
    115                     }
    116                 }
    117             }
    118         }
    119         Err(e) => {
    120             eprintln!("{e}");
    121             ExitCode::FAILURE
    122         }
    123     }
    124 }