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