move ApiConfig, FileLogger and CommandoSocket to proxmox-rest-server workspace
ApiConfig: avoid using pbs_config::backup_user() CommandoSocket: avoid using pbs_config::backup_user() FileLogger: avoid using pbs_config::backup_user() - use atomic_open_or_create_file() Auth Trait: moved definitions to proxmox-rest-server/src/lib.rs - removed CachedUserInfo patrameter - return user as String (not Authid) Signed-off-by: Thomas Lamprecht <t.lamprecht@proxmox.com>
This commit is contained in:
Dietmar Maurer
committed by
Thomas Lamprecht
Thomas Lamprecht
parent
037f6b6d5e
commit
fd6d243843
@ -7,3 +7,18 @@ description = "REST server implementation"
|
||||
|
||||
[dependencies]
|
||||
anyhow = "1.0"
|
||||
futures = "0.3"
|
||||
handlebars = "3.0"
|
||||
http = "0.2"
|
||||
hyper = { version = "0.14", features = [ "full" ] }
|
||||
lazy_static = "1.4"
|
||||
libc = "0.2"
|
||||
nix = "0.19.1"
|
||||
serde = { version = "1.0", features = [] }
|
||||
serde_json = "1.0"
|
||||
tokio = { version = "1.6", features = ["signal", "process"] }
|
||||
|
||||
proxmox = { version = "0.13.3", features = [ "router"] }
|
||||
|
||||
# fixme: remove this dependency (pbs_tools::broadcast_future)
|
||||
pbs-tools = { path = "../pbs-tools" }
|
||||
|
||||
170
proxmox-rest-server/src/api_config.rs
Normal file
170
proxmox-rest-server/src/api_config.rs
Normal file
@ -0,0 +1,170 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::time::SystemTime;
|
||||
use std::fs::metadata;
|
||||
use std::sync::{Arc, Mutex, RwLock};
|
||||
|
||||
use anyhow::{bail, Error, format_err};
|
||||
use hyper::Method;
|
||||
use handlebars::Handlebars;
|
||||
use serde::Serialize;
|
||||
|
||||
use proxmox::api::{ApiMethod, Router, RpcEnvironmentType};
|
||||
use proxmox::tools::fs::{create_path, CreateOptions};
|
||||
|
||||
use crate::{ApiAuth, FileLogger, FileLogOptions, CommandoSocket};
|
||||
|
||||
pub struct ApiConfig {
|
||||
basedir: PathBuf,
|
||||
router: &'static Router,
|
||||
aliases: HashMap<String, PathBuf>,
|
||||
env_type: RpcEnvironmentType,
|
||||
templates: RwLock<Handlebars<'static>>,
|
||||
template_files: RwLock<HashMap<String, (SystemTime, PathBuf)>>,
|
||||
request_log: Option<Arc<Mutex<FileLogger>>>,
|
||||
pub api_auth: Arc<dyn ApiAuth + Send + Sync>,
|
||||
}
|
||||
|
||||
impl ApiConfig {
|
||||
pub fn new<B: Into<PathBuf>>(
|
||||
basedir: B,
|
||||
router: &'static Router,
|
||||
env_type: RpcEnvironmentType,
|
||||
api_auth: Arc<dyn ApiAuth + Send + Sync>,
|
||||
) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
basedir: basedir.into(),
|
||||
router,
|
||||
aliases: HashMap::new(),
|
||||
env_type,
|
||||
templates: RwLock::new(Handlebars::new()),
|
||||
template_files: RwLock::new(HashMap::new()),
|
||||
request_log: None,
|
||||
api_auth,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn find_method(
|
||||
&self,
|
||||
components: &[&str],
|
||||
method: Method,
|
||||
uri_param: &mut HashMap<String, String>,
|
||||
) -> Option<&'static ApiMethod> {
|
||||
|
||||
self.router.find_method(components, method, uri_param)
|
||||
}
|
||||
|
||||
pub fn find_alias(&self, components: &[&str]) -> PathBuf {
|
||||
|
||||
let mut prefix = String::new();
|
||||
let mut filename = self.basedir.clone();
|
||||
let comp_len = components.len();
|
||||
if comp_len >= 1 {
|
||||
prefix.push_str(components[0]);
|
||||
if let Some(subdir) = self.aliases.get(&prefix) {
|
||||
filename.push(subdir);
|
||||
components.iter().skip(1).for_each(|comp| filename.push(comp));
|
||||
} else {
|
||||
components.iter().for_each(|comp| filename.push(comp));
|
||||
}
|
||||
}
|
||||
filename
|
||||
}
|
||||
|
||||
pub fn add_alias<S, P>(&mut self, alias: S, path: P)
|
||||
where S: Into<String>,
|
||||
P: Into<PathBuf>,
|
||||
{
|
||||
self.aliases.insert(alias.into(), path.into());
|
||||
}
|
||||
|
||||
pub fn env_type(&self) -> RpcEnvironmentType {
|
||||
self.env_type
|
||||
}
|
||||
|
||||
pub fn register_template<P>(&self, name: &str, path: P) -> Result<(), Error>
|
||||
where
|
||||
P: Into<PathBuf>
|
||||
{
|
||||
if self.template_files.read().unwrap().contains_key(name) {
|
||||
bail!("template already registered");
|
||||
}
|
||||
|
||||
let path: PathBuf = path.into();
|
||||
let metadata = metadata(&path)?;
|
||||
let mtime = metadata.modified()?;
|
||||
|
||||
self.templates.write().unwrap().register_template_file(name, &path)?;
|
||||
self.template_files.write().unwrap().insert(name.to_string(), (mtime, path));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Checks if the template was modified since the last rendering
|
||||
/// if yes, it loads a the new version of the template
|
||||
pub fn render_template<T>(&self, name: &str, data: &T) -> Result<String, Error>
|
||||
where
|
||||
T: Serialize,
|
||||
{
|
||||
let path;
|
||||
let mtime;
|
||||
{
|
||||
let template_files = self.template_files.read().unwrap();
|
||||
let (old_mtime, old_path) = template_files.get(name).ok_or_else(|| format_err!("template not found"))?;
|
||||
|
||||
mtime = metadata(old_path)?.modified()?;
|
||||
if mtime <= *old_mtime {
|
||||
return self.templates.read().unwrap().render(name, data).map_err(|err| format_err!("{}", err));
|
||||
}
|
||||
path = old_path.to_path_buf();
|
||||
}
|
||||
|
||||
{
|
||||
let mut template_files = self.template_files.write().unwrap();
|
||||
let mut templates = self.templates.write().unwrap();
|
||||
|
||||
templates.register_template_file(name, &path)?;
|
||||
template_files.insert(name.to_string(), (mtime, path));
|
||||
|
||||
templates.render(name, data).map_err(|err| format_err!("{}", err))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn enable_file_log<P>(
|
||||
&mut self,
|
||||
path: P,
|
||||
dir_opts: Option<CreateOptions>,
|
||||
file_opts: Option<CreateOptions>,
|
||||
commando_sock: &mut CommandoSocket,
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
P: Into<PathBuf>
|
||||
{
|
||||
let path: PathBuf = path.into();
|
||||
if let Some(base) = path.parent() {
|
||||
if !base.exists() {
|
||||
create_path(base, None, dir_opts).map_err(|err| format_err!("{}", err))?;
|
||||
}
|
||||
}
|
||||
|
||||
let logger_options = FileLogOptions {
|
||||
append: true,
|
||||
file_opts: file_opts.unwrap_or(CreateOptions::default()),
|
||||
..Default::default()
|
||||
};
|
||||
let request_log = Arc::new(Mutex::new(FileLogger::new(&path, logger_options)?));
|
||||
self.request_log = Some(Arc::clone(&request_log));
|
||||
|
||||
commando_sock.register_command("api-access-log-reopen".into(), move |_args| {
|
||||
println!("re-opening log file");
|
||||
request_log.lock().unwrap().reopen()?;
|
||||
Ok(serde_json::Value::Null)
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn get_file_log(&self) -> Option<&Arc<Mutex<FileLogger>>> {
|
||||
self.request_log.as_ref()
|
||||
}
|
||||
}
|
||||
222
proxmox-rest-server/src/command_socket.rs
Normal file
222
proxmox-rest-server/src/command_socket.rs
Normal file
@ -0,0 +1,222 @@
|
||||
use anyhow::{bail, format_err, Error};
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
use std::path::{PathBuf, Path};
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::*;
|
||||
use tokio::net::UnixListener;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use nix::sys::socket;
|
||||
use nix::unistd::Gid;
|
||||
|
||||
// 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)
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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.
|
||||
pub type CommandoSocketFn = 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.
|
||||
/// You need to call `spawn()` to make the socket active.
|
||||
pub struct CommandoSocket {
|
||||
socket: PathBuf,
|
||||
gid: Gid,
|
||||
commands: HashMap<String, CommandoSocketFn>,
|
||||
}
|
||||
|
||||
impl CommandoSocket {
|
||||
pub fn new<P>(path: P, gid: Gid) -> Self
|
||||
where P: Into<PathBuf>,
|
||||
{
|
||||
CommandoSocket {
|
||||
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(())
|
||||
}
|
||||
}
|
||||
148
proxmox-rest-server/src/file_logger.rs
Normal file
148
proxmox-rest-server/src/file_logger.rs
Normal file
@ -0,0 +1,148 @@
|
||||
use std::io::Write;
|
||||
|
||||
use anyhow::Error;
|
||||
use nix::fcntl::OFlag;
|
||||
|
||||
use proxmox::tools::fs::{CreateOptions, atomic_open_or_create_file};
|
||||
|
||||
/// Log messages with optional automatically added timestamps into files
|
||||
///
|
||||
/// Logs messages to file, and optionally to standard output.
|
||||
///
|
||||
///
|
||||
/// #### Example:
|
||||
/// ```
|
||||
/// # use anyhow::{bail, format_err, Error};
|
||||
/// use proxmox_rest_server::{flog, FileLogger, FileLogOptions};
|
||||
///
|
||||
/// # std::fs::remove_file("test.log");
|
||||
/// let options = FileLogOptions {
|
||||
/// to_stdout: true,
|
||||
/// exclusive: true,
|
||||
/// ..Default::default()
|
||||
/// };
|
||||
/// let mut log = FileLogger::new("test.log", options).unwrap();
|
||||
/// flog!(log, "A simple log: {}", "Hello!");
|
||||
/// # std::fs::remove_file("test.log");
|
||||
/// ```
|
||||
|
||||
#[derive(Default)]
|
||||
/// Options to control the behavior of a ['FileLogger'] instance
|
||||
pub struct FileLogOptions {
|
||||
/// Open underlying log file in append mode, useful when multiple concurrent processes
|
||||
/// want to log to the same file (e.g., HTTP access log). Note that it is only atomic
|
||||
/// for writes smaller than the PIPE_BUF (4k on Linux).
|
||||
/// Inside the same process you may need to still use an mutex, for shared access.
|
||||
pub append: bool,
|
||||
/// Open underlying log file as readable
|
||||
pub read: bool,
|
||||
/// If set, ensure that the file is newly created or error out if already existing.
|
||||
pub exclusive: bool,
|
||||
/// Duplicate logged messages to STDOUT, like tee
|
||||
pub to_stdout: bool,
|
||||
/// Prefix messages logged to the file with the current local time as RFC 3339
|
||||
pub prefix_time: bool,
|
||||
/// File owner/group and mode
|
||||
pub file_opts: CreateOptions,
|
||||
|
||||
}
|
||||
|
||||
pub struct FileLogger {
|
||||
file: std::fs::File,
|
||||
file_name: std::path::PathBuf,
|
||||
options: FileLogOptions,
|
||||
}
|
||||
|
||||
/// Log messages to [`FileLogger`](tools/struct.FileLogger.html)
|
||||
#[macro_export]
|
||||
macro_rules! flog {
|
||||
($log:expr, $($arg:tt)*) => ({
|
||||
$log.log(format!($($arg)*));
|
||||
})
|
||||
}
|
||||
|
||||
impl FileLogger {
|
||||
pub fn new<P: AsRef<std::path::Path>>(
|
||||
file_name: P,
|
||||
options: FileLogOptions,
|
||||
) -> Result<Self, Error> {
|
||||
let file = Self::open(&file_name, &options)?;
|
||||
|
||||
let file_name: std::path::PathBuf = file_name.as_ref().to_path_buf();
|
||||
|
||||
Ok(Self { file, file_name, options })
|
||||
}
|
||||
|
||||
pub fn reopen(&mut self) -> Result<&Self, Error> {
|
||||
let file = Self::open(&self.file_name, &self.options)?;
|
||||
self.file = file;
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
fn open<P: AsRef<std::path::Path>>(
|
||||
file_name: P,
|
||||
options: &FileLogOptions,
|
||||
) -> Result<std::fs::File, Error> {
|
||||
|
||||
let mut flags = OFlag::O_CLOEXEC;
|
||||
|
||||
if options.read {
|
||||
flags |= OFlag::O_RDWR;
|
||||
} else {
|
||||
flags |= OFlag::O_WRONLY;
|
||||
}
|
||||
|
||||
if options.append {
|
||||
flags |= OFlag::O_APPEND;
|
||||
}
|
||||
if options.exclusive {
|
||||
flags |= OFlag::O_EXCL;
|
||||
}
|
||||
|
||||
let file = atomic_open_or_create_file(&file_name, flags, &[], options.file_opts.clone())?;
|
||||
|
||||
Ok(file)
|
||||
}
|
||||
|
||||
pub fn log<S: AsRef<str>>(&mut self, msg: S) {
|
||||
let msg = msg.as_ref();
|
||||
|
||||
if self.options.to_stdout {
|
||||
let mut stdout = std::io::stdout();
|
||||
stdout.write_all(msg.as_bytes()).unwrap();
|
||||
stdout.write_all(b"\n").unwrap();
|
||||
}
|
||||
|
||||
let line = if self.options.prefix_time {
|
||||
let now = proxmox::tools::time::epoch_i64();
|
||||
let rfc3339 = match proxmox::tools::time::epoch_to_rfc3339(now) {
|
||||
Ok(rfc3339) => rfc3339,
|
||||
Err(_) => "1970-01-01T00:00:00Z".into(), // for safety, should really not happen!
|
||||
};
|
||||
format!("{}: {}\n", rfc3339, msg)
|
||||
} else {
|
||||
format!("{}\n", msg)
|
||||
};
|
||||
if let Err(err) = self.file.write_all(line.as_bytes()) {
|
||||
// avoid panicking, log methods should not do that
|
||||
// FIXME: or, return result???
|
||||
eprintln!("error writing to log file - {}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::io::Write for FileLogger {
|
||||
fn write(&mut self, buf: &[u8]) -> Result<usize, std::io::Error> {
|
||||
if self.options.to_stdout {
|
||||
let _ = std::io::stdout().write(buf);
|
||||
}
|
||||
self.file.write(buf)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> Result<(), std::io::Error> {
|
||||
if self.options.to_stdout {
|
||||
let _ = std::io::stdout().flush();
|
||||
}
|
||||
self.file.flush()
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,54 @@
|
||||
use anyhow::{bail, Error};
|
||||
|
||||
mod state;
|
||||
pub use state::*;
|
||||
|
||||
mod command_socket;
|
||||
pub use command_socket::*;
|
||||
|
||||
mod file_logger;
|
||||
pub use file_logger::{FileLogger, FileLogOptions};
|
||||
|
||||
mod api_config;
|
||||
pub use api_config::ApiConfig;
|
||||
|
||||
pub enum AuthError {
|
||||
Generic(Error),
|
||||
NoData,
|
||||
}
|
||||
|
||||
impl From<Error> for AuthError {
|
||||
fn from(err: Error) -> Self {
|
||||
AuthError::Generic(err)
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ApiAuth {
|
||||
fn check_auth(
|
||||
&self,
|
||||
headers: &http::HeaderMap,
|
||||
method: &hyper::Method,
|
||||
) -> Result<String, AuthError>;
|
||||
}
|
||||
|
||||
static mut SHUTDOWN_REQUESTED: bool = false;
|
||||
|
||||
pub fn request_shutdown() {
|
||||
unsafe {
|
||||
SHUTDOWN_REQUESTED = true;
|
||||
}
|
||||
crate::server_shutdown();
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn shutdown_requested() -> bool {
|
||||
unsafe { SHUTDOWN_REQUESTED }
|
||||
}
|
||||
|
||||
pub fn fail_on_shutdown() -> Result<(), Error> {
|
||||
if shutdown_requested() {
|
||||
bail!("Server shutdown requested - aborting task");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
142
proxmox-rest-server/src/state.rs
Normal file
142
proxmox-rest-server/src/state.rs
Normal file
@ -0,0 +1,142 @@
|
||||
use anyhow::{Error};
|
||||
use lazy_static::lazy_static;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use futures::*;
|
||||
|
||||
use tokio::signal::unix::{signal, SignalKind};
|
||||
|
||||
use pbs_tools::broadcast_future::BroadcastData;
|
||||
|
||||
#[derive(PartialEq, Copy, Clone, Debug)]
|
||||
pub enum ServerMode {
|
||||
Normal,
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
pub struct ServerState {
|
||||
pub mode: ServerMode,
|
||||
pub shutdown_listeners: BroadcastData<()>,
|
||||
pub last_worker_listeners: BroadcastData<()>,
|
||||
pub worker_count: usize,
|
||||
pub internal_task_count: usize,
|
||||
pub reload_request: bool,
|
||||
}
|
||||
|
||||
lazy_static! {
|
||||
static ref SERVER_STATE: Mutex<ServerState> = Mutex::new(ServerState {
|
||||
mode: ServerMode::Normal,
|
||||
shutdown_listeners: BroadcastData::new(),
|
||||
last_worker_listeners: BroadcastData::new(),
|
||||
worker_count: 0,
|
||||
internal_task_count: 0,
|
||||
reload_request: false,
|
||||
});
|
||||
}
|
||||
|
||||
pub fn server_state_init() -> Result<(), Error> {
|
||||
|
||||
let mut stream = signal(SignalKind::interrupt())?;
|
||||
|
||||
let future = async move {
|
||||
while stream.recv().await.is_some() {
|
||||
println!("got shutdown request (SIGINT)");
|
||||
SERVER_STATE.lock().unwrap().reload_request = false;
|
||||
crate::request_shutdown();
|
||||
}
|
||||
}.boxed();
|
||||
|
||||
let abort_future = last_worker_future().map_err(|_| {});
|
||||
let task = futures::future::select(future, abort_future);
|
||||
|
||||
tokio::spawn(task.map(|_| ()));
|
||||
|
||||
let mut stream = signal(SignalKind::hangup())?;
|
||||
|
||||
let future = async move {
|
||||
while stream.recv().await.is_some() {
|
||||
println!("got reload request (SIGHUP)");
|
||||
SERVER_STATE.lock().unwrap().reload_request = true;
|
||||
crate::request_shutdown();
|
||||
}
|
||||
}.boxed();
|
||||
|
||||
let abort_future = last_worker_future().map_err(|_| {});
|
||||
let task = futures::future::select(future, abort_future);
|
||||
|
||||
tokio::spawn(task.map(|_| ()));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn is_reload_request() -> bool {
|
||||
let data = SERVER_STATE.lock().unwrap();
|
||||
|
||||
data.mode == ServerMode::Shutdown && data.reload_request
|
||||
}
|
||||
|
||||
pub fn server_shutdown() {
|
||||
let mut data = SERVER_STATE.lock().unwrap();
|
||||
|
||||
println!("SET SHUTDOWN MODE");
|
||||
|
||||
data.mode = ServerMode::Shutdown;
|
||||
|
||||
data.shutdown_listeners.notify_listeners(Ok(()));
|
||||
|
||||
drop(data); // unlock
|
||||
|
||||
check_last_worker();
|
||||
}
|
||||
|
||||
pub fn shutdown_future() -> impl Future<Output = ()> {
|
||||
let mut data = SERVER_STATE.lock().unwrap();
|
||||
data
|
||||
.shutdown_listeners
|
||||
.listen()
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
pub fn last_worker_future() -> impl Future<Output = Result<(), Error>> {
|
||||
let mut data = SERVER_STATE.lock().unwrap();
|
||||
data.last_worker_listeners.listen()
|
||||
}
|
||||
|
||||
pub fn set_worker_count(count: usize) {
|
||||
SERVER_STATE.lock().unwrap().worker_count = count;
|
||||
|
||||
check_last_worker();
|
||||
}
|
||||
|
||||
pub fn check_last_worker() {
|
||||
let mut data = SERVER_STATE.lock().unwrap();
|
||||
|
||||
if !(data.mode == ServerMode::Shutdown && data.worker_count == 0 && data.internal_task_count == 0) { return; }
|
||||
|
||||
data.last_worker_listeners.notify_listeners(Ok(()));
|
||||
}
|
||||
|
||||
/// Spawns a tokio task that will be tracked for reload
|
||||
/// and if it is finished, notify the last_worker_listener if we
|
||||
/// are in shutdown mode
|
||||
pub fn spawn_internal_task<T>(task: T)
|
||||
where
|
||||
T: Future + Send + 'static,
|
||||
T::Output: Send + 'static,
|
||||
{
|
||||
let mut data = SERVER_STATE.lock().unwrap();
|
||||
data.internal_task_count += 1;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _ = tokio::spawn(task).await; // ignore errors
|
||||
|
||||
{ // drop mutex
|
||||
let mut data = SERVER_STATE.lock().unwrap();
|
||||
if data.internal_task_count > 0 {
|
||||
data.internal_task_count -= 1;
|
||||
}
|
||||
}
|
||||
|
||||
check_last_worker();
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user