tyler/proxmox-backup
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

src/server/command_socket.rs: switch to async

Browse Source 
Signed-off-by: Wolfgang Bumiller <w.bumiller@proxmox.com>
This commit is contained in:
Wolfgang Bumiller 2019-08-23 15:00:18 +02:00
parent 3c0facc787
commit 160fc8147f
1 changed files with 74 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

144
src/server/command_socket.rs
View File

@ -1,12 +1,8 @@
use failure::*; use failure::*;
use futures::*; use futures::*;
use futures::stream::Stream;
use tokio::net::unix::UnixListener; use tokio::net::unix::UnixListener;
use tokio::io::AsyncRead;
use std::io::Write;
use std::path::PathBuf; use std::path::PathBuf;
use serde_json::Value; use serde_json::Value;
@ -14,12 +10,11 @@ use std::sync::Arc;
use std::os::unix::io::AsRawFd; use std::os::unix::io::AsRawFd;
use nix::sys::socket; use nix::sys::socket;
use proxmox::tools::try_block;
/// Listens on a Unix Socket to handle simple command asynchronously /// Listens on a Unix Socket to handle simple command asynchronously
pub fn create_control_socket<P, F>(path: P, f: F) -> Result<impl Future<Item=(), Error=()>, Error> pub fn create_control_socket<P, F>(path: P, f: F) -> Result<impl Future<Output = ()>, Error>
where P: Into<PathBuf>, where
F: Send + Sync +'static + Fn(Value) -> Result<Value, Error>, P: Into<PathBuf>,
F: Fn(Value) -> Result<Value, Error> + Send + Sync + 'static,
{ {
let path: PathBuf = path.into(); let path: PathBuf = path.into();
@ -32,60 +27,75 @@ pub fn create_control_socket<P, F>(path: P, f: F) -> Result<impl Future<Item=(),
let control_future = socket.incoming() let control_future = socket.incoming()
.map_err(Error::from) .map_err(Error::from)
.and_then(|conn| { .and_then(|conn| {
use futures::future::{err, ok};
// check permissions (same gid, or root user) // check permissions (same gid, or root user)
let opt = socket::sockopt::PeerCredentials {}; let opt = socket::sockopt::PeerCredentials {};
match socket::getsockopt(conn.as_raw_fd(), opt) { match socket::getsockopt(conn.as_raw_fd(), opt) {
Ok(cred) => { Ok(cred) => {
let mygid = unsafe { libc::getgid() }; let mygid = unsafe { libc::getgid() };
if !(cred.uid() == 0 || cred.gid() == mygid) { if !(cred.uid() == 0 || cred.gid() == mygid) {
bail!("no permissions for {:?}", cred); return err(format_err!("no permissions for {:?}", cred));
} }
} }
Err(err) => bail!("no permissions - unable to read peer credential - {}", err), Err(e) => {
return err(format_err!(
"no permissions - unable to read peer credential - {}",
e,
));
}
} }
Ok(conn) ok(conn)
}) })
.map_err(move |err| { eprintln!("failed to accept on control socket {:?}: {}", path2, err); }) .map_err(move |err| { eprintln!("failed to accept on control socket {:?}: {}", path2, err); })
.for_each(move |conn| { .try_for_each(move |conn| {
let f1 = f.clone(); let f = Arc::clone(&f);
let (rx, mut tx) = conn.split(); let (rx, mut tx) = conn.split();
let path = path3.clone(); let path = path3.clone();
let path2 = path3.clone();
let abort_future = super::last_worker_future().map_err(|_| {}); let abort_future = super::last_worker_future().map(|_| ());
tokio::spawn( use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
tokio::io::lines(std::io::BufReader::new(rx)) tokio::spawn(futures::future::select(
.map_err(move |err| { eprintln!("control socket {:?} read error: {}", path, err); }) async move {
.and_then(move |cmd| { let mut rx = tokio::io::BufReader::new(rx);
let res = try_block!({ let mut line = String::new();
let param = match cmd.parse::<Value>() { loop {
Ok(p) => p, line.clear();
Err(err) => bail!("unable to parse json value - {}", err), match rx.read_line({ line.clear(); &mut line }).await {
}; Ok(0) => break,
Ok(_) => (),
Err(err) => {
eprintln!("control socket {:?} read error: {}", path, err);
return;
}
}
f1(param) let response = match line.parse::<Value>() {
}); Ok(param) => match f(param) {
Ok(res) => format!("OK: {}\n", res),
let resp = match res { Err(err) => format!("ERROR: {}\n", err),
Ok(v) => format!("OK: {}\n", v), }
Err(err) => format!("ERROR: {}\n", err), Err(err) => format!("ERROR: {}\n", err),
}; };
Ok(resp)
}) if let Err(err) = tx.write_all(response.as_bytes()).await {
.for_each(move |resp| { eprintln!("control socket {:?} write response error: {}", path, err);
tx.write_all(resp.as_bytes()) return;
.map_err(|err| { eprintln!("control socket {:?} write response error: {}", path2, err); }) }
}) }
.select(abort_future) }.boxed(),
.then(move |_| { Ok(()) }) abort_future,
) ).map(|_| ()));
futures::future::ok(())
}); });
let abort_future = super::last_worker_future().map_err(|_| {}); let abort_future = super::last_worker_future().map_err(|_| {});
let task = control_future.select(abort_future) let task = futures::future::select(
.then(move |_| { Ok(()) }); control_future,
abort_future,
).map(|_| ());
Ok(task) Ok(task)
} }
@ -94,7 +104,7 @@ pub fn create_control_socket<P, F>(path: P, f: F) -> Result<impl Future<Item=(),
pub fn send_command<P>( pub fn send_command<P>(
path: P, path: P,
params: Value params: Value
) -> impl Future<Item=Value, Error=Error> ) -> impl Future<Output = Result<Value, Error>>
where P: Into<PathBuf>, where P: Into<PathBuf>,
{ {
@ -104,37 +114,31 @@ pub fn send_command<P>(
.map_err(move |err| format_err!("control socket connect failed - {}", err)) .map_err(move |err| format_err!("control socket connect failed - {}", err))
.and_then(move |conn| { .and_then(move |conn| {
let (rx, tx) = conn.split(); let (rx, mut tx) = conn.split();
let mut command_string = params.to_string(); let mut command_string = params.to_string();
command_string.push('\n'); command_string.push('\n');
tokio::io::write_all(tx, command_string) async move {
.and_then(|(tx,_)| tokio::io::shutdown(tx)) use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
.map_err(|err| format_err!("control socket write error - {}", err))
.and_then(move |_| { tx.write_all(command_string.as_bytes()).await?;
tokio::io::lines(std::io::BufReader::new(rx)) tx.shutdown().await?;
.into_future() let mut rx = tokio::io::BufReader::new(rx);
.then(|test| { let mut data = String::new();
match test { if rx.read_line(&mut data).await? == 0 {
Ok((Some(data), _)) => { bail!("no response");
if data.starts_with("OK: ") { }
match data[4..].parse::<Value>() { if data.starts_with("OK: ") {
Ok(v) => Ok(v), match data[4..].parse::<Value>() {
Err(err) => bail!("unable to parse json response - {}", err), Ok(v) => Ok(v),
} Err(err) => bail!("unable to parse json response - {}", err),
} else if data.starts_with("ERROR: ") { }
bail!("{}", &data[7..]); } else if data.starts_with("ERROR: ") {
} else { bail!("{}", &data[7..]);
bail!("unable to parse response: {}", data); } else {
} bail!("unable to parse response: {}", data);
} }
Ok((None, _)) => { }
bail!("no response");
}
Err((err, _)) => Err(Error::from(err)),
}
})
})
}) })
} }