proxmox-backup/proxmox-rest-server/src/command_socket.rs

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<P, F>(
    path: P,
    gid: Gid,
    func: F,
) -> Result<impl Future<Output = ()>, Error>
where
    P: Into<PathBuf>,
    F: Fn(Value) -> Result<Value, Error> + 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::<Value>() {
                                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<P, T>(path: P, params: &T) -> Result<Value, Error>
where
    P: AsRef<Path>,
    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<P>(path: P, command_string: &str) -> Result<Value, Error>
where
    P: AsRef<Path>,
{
    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::<Value>() {
            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<Value, Error> + 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<String, CommandSocketFn>,
}

impl CommandSocket {
    /// Creates a new instance.
    pub fn new<P>(path: P, gid: Gid) -> Self
    where
        P: Into<PathBuf>,
    {
        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<F>(&mut self, command: String, handler: F) -> Result<(), Error>
    where
        F: Fn(Option<&Value>) -> Result<Value, Error> + Send + Sync + 'static,
    {
        if self.commands.contains_key(&command) {
            bail!("command '{}' already exists!", command);
        }

        self.commands.insert(command, Box::new(handler));

        Ok(())
    }
}