src/main.rs

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