parameter handling:

the getopts crate is wrong, it uses String for free
arguments. We want a Path argument, which is not utf8.

clap is too big for our use case

so we're doing this manually

Signed-off-by: Wolfgang Bumiller <w.bumiller@proxmox.com>
This commit is contained in:
Wolfgang Bumiller 2020-01-22 10:50:04 +01:00
parent 31d748b335
commit d54e9e5a31
2 changed files with 81 additions and 6 deletions

View File

@ -1,6 +1,7 @@
// c_str!() from the byte-strings crate is implemented via a proc macro which seems a bit excessive
macro_rules! c_str {
($data:expr) => {{
#![allow(unused_unsafe)]
let bytes = concat!($data, "\0");
unsafe { std::ffi::CStr::from_bytes_with_nul_unchecked(bytes.as_bytes()) }
}};

View File

@ -1,5 +1,8 @@
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 failure::{bail, format_err, Error};
use nix::sys::socket::SockAddr;
@ -29,20 +32,72 @@ 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() {
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;
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 {
let bytes = arg.as_bytes();
if bytes.starts_with(b"-") {
let _ = stderr().write_all(b"unexpected option: ");
let _ = stderr().write_all(arg.as_bytes());
usage(1, &program, &mut stderr());
}
if path.is_some() {
let _ = stderr().write_all(b"unexpected extra parameter: ");
let _ = stderr().write_all(arg.as_bytes());
usage(1, &program, &mut stderr());
}
path = Some(arg);
}
}
let path = match path {
Some(path) => path,
None => {
eprintln!("missing path");
usage(1, &program, &mut stderr());
}
};
let mut rt = tokio::runtime::Runtime::new().expect("failed to spawn tokio runtime");
if let Err(err) = rt.block_on(do_main()) {
if let Err(err) = rt.block_on(do_main(use_sd_notify, path)) {
eprintln!("error: {}", err);
std::process::exit(1);
}
}
async fn do_main() -> Result<(), Error> {
let socket_path = std::env::args_os()
.nth(1)
.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() == StdIo::ErrorKind::NotFound => (), // Ok
@ -54,9 +109,28 @@ async fn do_main() -> Result<(), Error> {
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);
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))
}
}