use anyhow::{bail, format_err, Error}; use std::collections::HashMap; use std::os::unix::io::AsRawFd; use std::path::{Path, PathBuf}; use std::sync::Arc; use futures::*; use nix::sys::socket; use nix::unistd::Gid; use serde::Serialize; use serde_json::Value; use tokio::net::UnixListener; // Listens on a Unix Socket to handle simple command asynchronously fn create_control_socket( path: P, gid: Gid, func: F, ) -> Result, Error> where P: Into, F: Fn(Value) -> Result + Send + Sync + 'static, { let path: PathBuf = path.into(); let gid = gid.as_raw(); let socket = UnixListener::bind(&path)?; let func = Arc::new(func); let control_future = async move { loop { let (conn, _addr) = match socket.accept().await { Ok(data) => data, Err(err) => { eprintln!("failed to accept on control socket {:?}: {}", path, err); continue; } }; let opt = socket::sockopt::PeerCredentials {}; let cred = match socket::getsockopt(conn.as_raw_fd(), opt) { Ok(cred) => cred, Err(err) => { eprintln!("no permissions - unable to read peer credential - {}", err); continue; } }; // check permissions (same gid, root user, or backup group) let mygid = unsafe { libc::getgid() }; if !(cred.uid() == 0 || cred.gid() == mygid || cred.gid() == gid) { eprintln!("no permissions for {:?}", cred); continue; } let (rx, mut tx) = tokio::io::split(conn); let abort_future = super::last_worker_future().map(|_| ()); use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; let func = Arc::clone(&func); let path = path.clone(); tokio::spawn( futures::future::select( async move { let mut rx = tokio::io::BufReader::new(rx); let mut line = String::new(); loop { line.clear(); match rx .read_line({ line.clear(); &mut line }) .await { Ok(0) => break, Ok(_) => (), Err(err) => { eprintln!("control socket {:?} read error: {}", path, err); return; } } let response = match line.parse::() { Ok(param) => match func(param) { Ok(res) => format!("OK: {}\n", res), Err(err) => format!("ERROR: {}\n", err), }, Err(err) => format!("ERROR: {}\n", err), }; if let Err(err) = tx.write_all(response.as_bytes()).await { eprintln!( "control socket {:?} write response error: {}", path, err ); return; } } } .boxed(), abort_future, ) .map(|_| ()), ); } } .boxed(); let abort_future = crate::last_worker_future().map_err(|_| {}); let task = futures::future::select(control_future, abort_future) .map(|_: futures::future::Either<(Result<(), Error>, _), _>| ()); Ok(task) } /// Send a command to the specified socket pub async fn send_command(path: P, params: &T) -> Result where P: AsRef, T: ?Sized + Serialize, { let mut command_string = serde_json::to_string(params)?; command_string.push('\n'); send_raw_command(path.as_ref(), &command_string).await } /// Send a raw command (string) to the specified socket pub async fn send_raw_command

(path: P, command_string: &str) -> Result where P: AsRef, { use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; let mut conn = tokio::net::UnixStream::connect(path) .map_err(move |err| format_err!("control socket connect failed - {}", err)) .await?; conn.write_all(command_string.as_bytes()).await?; if !command_string.as_bytes().ends_with(b"\n") { conn.write_all(b"\n").await?; } AsyncWriteExt::shutdown(&mut conn).await?; let mut rx = tokio::io::BufReader::new(conn); let mut data = String::new(); if rx.read_line(&mut data).await? == 0 { bail!("no response"); } if let Some(res) = data.strip_prefix("OK: ") { match res.parse::() { Ok(v) => Ok(v), Err(err) => bail!("unable to parse json response - {}", err), } } else if let Some(err) = data.strip_prefix("ERROR: ") { bail!("{}", err); } else { bail!("unable to parse response: {}", data); } } // A callback for a specific commando socket. type CommandSocketFn = Box<(dyn Fn(Option<&Value>) -> Result + Send + Sync + 'static)>; /// Tooling to get a single control command socket where one can /// register multiple commands dynamically. /// /// The socket is activated by calling [spawn](CommandSocket::spawn), /// which spawns an async tokio task to process the commands. pub struct CommandSocket { socket: PathBuf, gid: Gid, commands: HashMap, } impl CommandSocket { /// Creates a new instance. pub fn new

(path: P, gid: Gid) -> Self where P: Into, { CommandSocket { socket: path.into(), gid, commands: HashMap::new(), } } /// Spawn the socket and consume self, meaning you cannot register commands anymore after /// calling this. pub fn spawn(self) -> Result<(), Error> { let control_future = create_control_socket(self.socket.to_owned(), self.gid, move |param| { let param = param.as_object().ok_or_else(|| { format_err!("unable to parse parameters (expected json object)") })?; let command = match param.get("command") { Some(Value::String(command)) => command.as_str(), None => bail!("no command"), _ => bail!("unable to parse command"), }; if !self.commands.contains_key(command) { bail!("got unknown command '{}'", command); } match self.commands.get(command) { None => bail!("got unknown command '{}'", command), Some(handler) => { let args = param.get("args"); //.unwrap_or(&Value::Null); (handler)(args) } } })?; tokio::spawn(control_future); Ok(()) } /// Register a new command with a callback. pub fn register_command(&mut self, command: String, handler: F) -> Result<(), Error> where F: Fn(Option<&Value>) -> Result + Send + Sync + 'static, { if self.commands.contains_key(&command) { bail!("command '{}' already exists!", command); } self.commands.insert(command, Box::new(handler)); Ok(()) } }