src/main.rs

   1 use std::ffi::{OsStr, OsString};
   2 use std::future::Future;
   3 use std::io as StdIo;
   4 use std::io::{stderr, stdout, Write};
   5 use std::os::unix::ffi::OsStrExt;
   6
   7 use anyhow::{bail, format_err, Error};
   8 use nix::sys::socket::SockAddr;
   9
  10 #[macro_use]
  11 mod macros;
  12
  13 pub mod apparmor;
  14 pub mod capability;
  15 pub mod client;
  16 pub mod error;
  17 pub mod fork;
  18 pub mod io;
  19 pub mod lxcseccomp;
  20 pub mod nsfd;
  21 pub mod poll_fn;
  22 pub mod process;
  23 pub mod seccomp;
  24 pub mod sys_mknod;
  25 pub mod sys_quotactl;
  26 pub mod syscall;
  27 pub mod tools;
  28
  29 use crate::io::seq_packet::SeqPacketListener;
  30
  31 pub fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
  32     tokio::spawn(fut);
  33 }
  34
  35 fn usage(status: i32, program: &OsStr, out: &mut dyn Write) -> ! {
  36     let _ = out.write_all("usage: ".as_bytes());
  37     let _ = out.write_all(program.as_bytes());
  38     let _ = out.write_all(
  39         concat!(
  40             "[options] SOCKET_PATH\n",
  41             "options:\n",
  42             "    -h, --help      show this help message\n",
  43             "    --system        \
  44                      run as systemd daemon (use sd_notify() when ready to accept connections)\n",
  45         )
  46         .as_bytes(),
  47     );
  48     std::process::exit(status);
  49 }
  50
  51 fn main() {
  52     let mut args = std::env::args_os();
  53     let program = args.next().unwrap(); // program name always exists
  54
  55     let mut use_sd_notify = false;
  56     let mut path = None;
  57
  58     let mut nonopt_arg = |arg: OsString| {
  59         if path.is_some() {
  60             let _ = stderr().write_all(b"unexpected extra parameter: ");
  61             let _ = stderr().write_all(arg.as_bytes());
  62             let _ = stderr().write_all(b"\n");
  63             usage(1, &program, &mut stderr());
  64         }
  65
  66         path = Some(arg);
  67     };
  68
  69     for arg in &mut args {
  70         if arg == "-h" || arg == "--help" {
  71             usage(0, &program, &mut stdout());
  72         }
  73
  74         if arg == "--" {
  75             break;
  76         } else if arg == "--system" {
  77             use_sd_notify = true;
  78         } else {
  79             if arg.as_bytes().starts_with(b"-") {
  80                 let _ = stderr().write_all(b"unexpected option: ");
  81                 let _ = stderr().write_all(arg.as_bytes());
  82                 let _ = stderr().write_all(b"\n");
  83                 usage(1, &program, &mut stderr());
  84             }
  85
  86             nonopt_arg(arg);
  87         }
  88     }
  89
  90     for arg in &mut args {
  91         nonopt_arg(arg);
  92     }
  93     drop(nonopt_arg);
  94
  95     let path = match path {
  96         Some(path) => path,
  97         None => {
  98             eprintln!("missing path");
  99             usage(1, &program, &mut stderr());
 100         }
 101     };
 102
 103     let rt = tokio::runtime::Runtime::new().expect("failed to spawn tokio runtime");
 104
 105     if let Err(err) = rt.block_on(do_main(use_sd_notify, path)) {
 106         eprintln!("error: {}", err);
 107         std::process::exit(1);
 108     }
 109 }
 110
 111 async fn do_main(use_sd_notify: bool, socket_path: OsString) -> Result<(), Error> {
 112     match std::fs::remove_file(&socket_path) {
 113         Ok(_) => (),
 114         Err(ref e) if e.kind() == StdIo::ErrorKind::NotFound => (), // Ok
 115         Err(e) => bail!("failed to remove previous socket: {}", e),
 116     }
 117
 118     let address =
 119         SockAddr::new_unix(socket_path.as_os_str()).expect("cannot create struct sockaddr_un?");
 120
 121     let mut listener = SeqPacketListener::bind(&address)
 122         .map_err(|e| format_err!("failed to create listening socket: {}", e))?;
 123
 124     if use_sd_notify {
 125         notify_systemd()?;
 126     }
 127
 128     loop {
 129         let client = listener.accept().await?;
 130         let client = client::Client::new(client);
 131         spawn(client.main());
 132     }
 133 }
 134
 135 #[link(name = "systemd")]
 136 extern "C" {
 137     fn sd_notify(unset_environment: libc::c_int, state: *const libc::c_char) -> libc::c_int;
 138 }
 139
 140 fn notify_systemd() -> StdIo::Result<()> {
 141     let err = unsafe { sd_notify(0, c_str!("READY=1\n").as_ptr()) };
 142     if err >= 0 {
 143         Ok(())
 144     } else {
 145         Err(StdIo::Error::from_raw_os_error(-err))
 146     }
 147 }