update to tokio 1.0

[pve-lxc-syscalld.git] / src / main.rs

diff --git a/src/main.rs b/src/main.rs
index f1051282fc41070ed9ab2462e79a3f0b7771f13e..ca4366d81b3f41dfc9e84ec0121938d22b762f7b 100644 (file)
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,65 +1,147 @@
-#![feature(async_await)]
+use std::ffi::{OsStr, OsString};
+use std::future::Future;
+use std::io as StdIo;
+use std::io::{stderr, stdout, Write};
+use std::os::unix::ffi::OsStrExt;
 
-use std::ffi::OsString;
-use std::io;
-
-use failure::{bail, format_err, Error};
+use anyhow::{bail, format_err, Error};
 use nix::sys::socket::SockAddr;
 
+#[macro_use]
+mod macros;
+
 pub mod apparmor;
+pub mod capability;
 pub mod client;
+pub mod error;
 pub mod fork;
+pub mod io;
 pub mod lxcseccomp;
 pub mod nsfd;
-pub mod pidfd;
+pub mod poll_fn;
+pub mod process;
 pub mod seccomp;
-pub mod socket;
 pub mod sys_mknod;
 pub mod sys_quotactl;
 pub mod syscall;
 pub mod tools;
 
-use socket::SeqPacketListener;
+use crate::io::seq_packet::SeqPacketListener;
+
+pub fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
+    tokio::spawn(fut);
+}
+
+fn usage(status: i32, program: &OsStr, out: &mut dyn Write) -> ! {
+    let _ = out.write_all("usage: ".as_bytes());
+    let _ = out.write_all(program.as_bytes());
+    let _ = out.write_all(
+        concat!(
+            "[options] SOCKET_PATH\n",
+            "options:\n",
+            "    -h, --help      show this help message\n",
+            "    --system        \
+                     run as systemd daemon (use sd_notify() when ready to accept connections)\n",
+        )
+        .as_bytes(),
+    );
+    std::process::exit(status);
+}
 
 fn main() {
-    if let Err(err) = run() {
+    let mut args = std::env::args_os();
+    let program = args.next().unwrap(); // program name always exists
+
+    let mut use_sd_notify = false;
+    let mut path = None;
+
+    let mut nonopt_arg = |arg: OsString| {
+        if path.is_some() {
+            let _ = stderr().write_all(b"unexpected extra parameter: ");
+            let _ = stderr().write_all(arg.as_bytes());
+            let _ = stderr().write_all(b"\n");
+            usage(1, &program, &mut stderr());
+        }
+
+        path = Some(arg);
+    };
+
+    for arg in &mut args {
+        if arg == "-h" || arg == "--help" {
+            usage(0, &program, &mut stdout());
+        }
+
+        if arg == "--" {
+            break;
+        } else if arg == "--system" {
+            use_sd_notify = true;
+        } else {
+            if arg.as_bytes().starts_with(b"-") {
+                let _ = stderr().write_all(b"unexpected option: ");
+                let _ = stderr().write_all(arg.as_bytes());
+                let _ = stderr().write_all(b"\n");
+                usage(1, &program, &mut stderr());
+            }
+
+            nonopt_arg(arg);
+        }
+    }
+
+    for arg in &mut args {
+        nonopt_arg(arg);
+    }
+    drop(nonopt_arg);
+
+    let path = match path {
+        Some(path) => path,
+        None => {
+            eprintln!("missing path");
+            usage(1, &program, &mut stderr());
+        }
+    };
+
+    let rt = tokio::runtime::Runtime::new().expect("failed to spawn tokio runtime");
+
+    if let Err(err) = rt.block_on(do_main(use_sd_notify, path)) {
         eprintln!("error: {}", err);
         std::process::exit(1);
     }
 }
 
-fn run() -> Result<(), Error> {
-    let socket_path = std::env::args_os()
-        .skip(1)
-        .next()
-        .ok_or_else(|| format_err!("missing parameter: socket path to listen on"))?;
-
+async fn do_main(use_sd_notify: bool, socket_path: OsString) -> Result<(), Error> {
     match std::fs::remove_file(&socket_path) {
         Ok(_) => (),
-        Err(ref e) if e.kind() == io::ErrorKind::NotFound => (), // Ok
+        Err(ref e) if e.kind() == StdIo::ErrorKind::NotFound => (), // Ok
         Err(e) => bail!("failed to remove previous socket: {}", e),
     }
 
-    tokio::run(async_run(socket_path));
-
-    Ok(())
-}
-
-async fn async_run(socket_path: OsString) {
-    if let Err(err) = async_run_do(socket_path).await {
-        eprintln!("error accepting clients, bailing out: {}", err);
-    }
-}
-
-async fn async_run_do(socket_path: OsString) -> Result<(), Error> {
     let address =
         SockAddr::new_unix(socket_path.as_os_str()).expect("cannot create struct sockaddr_un?");
 
     let mut listener = SeqPacketListener::bind(&address)
         .map_err(|e| format_err!("failed to create listening socket: {}", e))?;
+
+    if use_sd_notify {
+        notify_systemd()?;
+    }
+
     loop {
         let client = listener.accept().await?;
         let client = client::Client::new(client);
-        tokio::spawn(client.main());
+        spawn(client.main());
+    }
+}
+
+#[link(name = "systemd")]
+extern "C" {
+    fn sd_notify(unset_environment: libc::c_int, state: *const libc::c_char) -> libc::c_int;
+}
+
+fn notify_systemd() -> StdIo::Result<()> {
+    let err = unsafe { sd_notify(0, c_str!("READY=1\n").as_ptr()) };
+    if err >= 0 {
+        Ok(())
+    } else {
+        Err(StdIo::Error::from_raw_os_error(-err))
     }
 }