proxmox-rest-server/src/command_socket.rs

   1 use anyhow::{bail, format_err, Error};
   2
   3 use std::collections::HashMap;
   4 use std::os::unix::io::AsRawFd;
   5 use std::path::{PathBuf, Path};
   6 use std::sync::Arc;
   7
   8 use futures::*;
   9 use tokio::net::UnixListener;
  10 use serde::Serialize;
  11 use serde_json::Value;
  12 use nix::sys::socket;
  13 use nix::unistd::Gid;
  14
  15 // Listens on a Unix Socket to handle simple command asynchronously
  16 fn create_control_socket<P, F>(path: P, gid: Gid, func: F) -> Result<impl Future<Output = ()>, Error>
  17 where
  18     P: Into<PathBuf>,
  19     F: Fn(Value) -> Result<Value, Error> + Send + Sync + 'static,
  20 {
  21     let path: PathBuf = path.into();
  22
  23     let gid = gid.as_raw();
  24
  25     let socket = UnixListener::bind(&path)?;
  26
  27     let func = Arc::new(func);
  28
  29     let control_future = async move {
  30         loop {
  31             let (conn, _addr) = match socket.accept().await {
  32                 Ok(data) => data,
  33                 Err(err) => {
  34                     eprintln!("failed to accept on control socket {:?}: {}", path, err);
  35                     continue;
  36                 }
  37             };
  38
  39             let opt = socket::sockopt::PeerCredentials {};
  40             let cred = match socket::getsockopt(conn.as_raw_fd(), opt) {
  41                 Ok(cred) => cred,
  42                 Err(err) => {
  43                     eprintln!("no permissions - unable to read peer credential - {}", err);
  44                     continue;
  45                 }
  46             };
  47
  48             // check permissions (same gid, root user, or backup group)
  49             let mygid = unsafe { libc::getgid() };
  50             if !(cred.uid() == 0 || cred.gid() == mygid || cred.gid() == gid) {
  51                 eprintln!("no permissions for {:?}", cred);
  52                 continue;
  53             }
  54
  55             let (rx, mut tx) = tokio::io::split(conn);
  56
  57             let abort_future = super::last_worker_future().map(|_| ());
  58
  59             use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
  60             let func = Arc::clone(&func);
  61             let path = path.clone();
  62             tokio::spawn(futures::future::select(
  63                 async move {
  64                     let mut rx = tokio::io::BufReader::new(rx);
  65                     let mut line = String::new();
  66                     loop {
  67                         line.clear();
  68                         match rx.read_line({ line.clear(); &mut line }).await {
  69                             Ok(0) => break,
  70                             Ok(_) => (),
  71                             Err(err) => {
  72                                 eprintln!("control socket {:?} read error: {}", path, err);
  73                                 return;
  74                             }
  75                         }
  76
  77                         let response = match line.parse::<Value>() {
  78                             Ok(param) => match func(param) {
  79                                 Ok(res) => format!("OK: {}\n", res),
  80                                 Err(err) => format!("ERROR: {}\n", err),
  81                             }
  82                             Err(err) => format!("ERROR: {}\n", err),
  83                         };
  84
  85                         if let Err(err) = tx.write_all(response.as_bytes()).await {
  86                             eprintln!("control socket {:?} write response error: {}", path, err);
  87                             return;
  88                         }
  89                     }
  90                 }.boxed(),
  91                 abort_future,
  92             ).map(|_| ()));
  93         }
  94     }.boxed();
  95
  96     let abort_future = crate::last_worker_future().map_err(|_| {});
  97     let task = futures::future::select(
  98         control_future,
  99         abort_future,
 100     ).map(|_: futures::future::Either<(Result<(), Error>, _), _>| ());
 101
 102     Ok(task)
 103 }
 104
 105 /// Send a command to the specified socket
 106 pub async fn send_command<P, T>(path: P, params: &T) -> Result<Value, Error>
 107 where
 108     P: AsRef<Path>,
 109     T: ?Sized + Serialize,
 110 {
 111     let mut command_string = serde_json::to_string(params)?;
 112     command_string.push('\n');
 113     send_raw_command(path.as_ref(), &command_string).await
 114 }
 115
 116 /// Send a raw command (string) to the specified socket
 117 pub async fn send_raw_command<P>(path: P, command_string: &str) -> Result<Value, Error>
 118 where
 119     P: AsRef<Path>,
 120 {
 121     use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
 122
 123     let mut conn = tokio::net::UnixStream::connect(path)
 124         .map_err(move |err| format_err!("control socket connect failed - {}", err))
 125         .await?;
 126
 127     conn.write_all(command_string.as_bytes()).await?;
 128     if !command_string.as_bytes().ends_with(b"\n") {
 129         conn.write_all(b"\n").await?;
 130     }
 131
 132     AsyncWriteExt::shutdown(&mut conn).await?;
 133     let mut rx = tokio::io::BufReader::new(conn);
 134     let mut data = String::new();
 135     if rx.read_line(&mut data).await? == 0 {
 136         bail!("no response");
 137     }
 138     if let Some(res) = data.strip_prefix("OK: ") {
 139         match res.parse::<Value>() {
 140             Ok(v) => Ok(v),
 141             Err(err) => bail!("unable to parse json response - {}", err),
 142         }
 143     } else if let Some(err) = data.strip_prefix("ERROR: ") {
 144         bail!("{}", err);
 145     } else {
 146         bail!("unable to parse response: {}", data);
 147     }
 148 }
 149
 150 // A callback for a specific commando socket.
 151 type CommandoSocketFn = Box<(dyn Fn(Option<&Value>) -> Result<Value, Error> + Send + Sync + 'static)>;
 152
 153 /// Tooling to get a single control command socket where one can
 154 /// register multiple commands dynamically.
 155 ///
 156 /// You need to call `spawn()` to make the socket active.
 157 pub struct CommandoSocket {
 158     socket: PathBuf,
 159     gid: Gid,
 160     commands: HashMap<String, CommandoSocketFn>,
 161 }
 162
 163 impl CommandoSocket {
 164     pub fn new<P>(path: P, gid: Gid) -> Self
 165         where P: Into<PathBuf>,
 166     {
 167         CommandoSocket {
 168             socket: path.into(),
 169             gid,
 170             commands: HashMap::new(),
 171         }
 172     }
 173
 174     /// Spawn the socket and consume self, meaning you cannot register commands anymore after
 175     /// calling this.
 176     pub fn spawn(self) -> Result<(), Error> {
 177         let control_future = create_control_socket(self.socket.to_owned(), self.gid, move |param| {
 178             let param = param
 179                 .as_object()
 180                 .ok_or_else(|| format_err!("unable to parse parameters (expected json object)"))?;
 181
 182             let command = match param.get("command") {
 183                 Some(Value::String(command)) => command.as_str(),
 184                 None => bail!("no command"),
 185                 _ => bail!("unable to parse command"),
 186             };
 187
 188             if !self.commands.contains_key(command) {
 189                 bail!("got unknown command '{}'", command);
 190             }
 191
 192             match self.commands.get(command) {
 193                 None => bail!("got unknown command '{}'", command),
 194                 Some(handler) => {
 195                     let args = param.get("args"); //.unwrap_or(&Value::Null);
 196                     (handler)(args)
 197                 },
 198             }
 199         })?;
 200
 201         tokio::spawn(control_future);
 202
 203         Ok(())
 204     }
 205
 206     /// Register a new command with a callback.
 207     pub fn register_command<F>(
 208         &mut self,
 209         command: String,
 210         handler: F,
 211     ) -> Result<(), Error>
 212     where
 213         F: Fn(Option<&Value>) -> Result<Value, Error> + Send + Sync + 'static,
 214     {
 215
 216         if self.commands.contains_key(&command) {
 217             bail!("command '{}' already exists!", command);
 218         }
 219
 220         self.commands.insert(command, Box::new(handler));
 221
 222         Ok(())
 223     }
 224 }