WIP
This commit is contained in:
@@ -253,7 +253,7 @@ fn run_forked_server(
|
||||
.block_on(async {
|
||||
let socket = TokioUnixStream::from_std(server_socket)?;
|
||||
let db_pool = construct_single_connection_mysql_pool(&config.mysql).await?;
|
||||
session_handler::session_handler(socket, &unix_user, db_pool).await?;
|
||||
session_handler::session_handler_with_unix_user(socket, &unix_user, db_pool).await?;
|
||||
Ok(())
|
||||
});
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ use std::ffi::CString;
|
||||
pub const DEFAULT_CONFIG_PATH: &str = "/etc/muscl/config.toml";
|
||||
pub const DEFAULT_SOCKET_PATH: &str = "/run/muscl/muscl.sock";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UnixUser {
|
||||
pub username: String,
|
||||
pub groups: Vec<String>,
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::collections::BTreeSet;
|
||||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use indoc::concatdoc;
|
||||
use sqlx::{MySql, MySqlConnection, MySqlPool, pool::PoolConnection};
|
||||
use sqlx::{MySqlConnection, MySqlPool};
|
||||
use tokio::net::UnixStream;
|
||||
|
||||
use crate::{
|
||||
@@ -30,15 +30,58 @@ use crate::{
|
||||
|
||||
// TODO: don't use database connection unless necessary.
|
||||
|
||||
pub async fn session_handler(
|
||||
pub async fn session_handler(socket: UnixStream, db_pool: MySqlPool) -> anyhow::Result<()> {
|
||||
let uid = match socket.peer_cred() {
|
||||
Ok(cred) => cred.uid(),
|
||||
Err(e) => {
|
||||
log::error!("Failed to get peer credentials from socket: {}", e);
|
||||
let mut message_stream = create_server_to_client_message_stream(socket);
|
||||
message_stream
|
||||
.send(Response::Error(
|
||||
(concatdoc! {
|
||||
"Server failed to get peer credentials from socket\n",
|
||||
"Please check the server logs or contact the system administrators"
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.await
|
||||
.ok();
|
||||
anyhow::bail!("Failed to get peer credentials from socket");
|
||||
}
|
||||
};
|
||||
|
||||
log::debug!("Validated peer UID: {}", uid);
|
||||
|
||||
let unix_user = match UnixUser::from_uid(uid) {
|
||||
Ok(user) => user,
|
||||
Err(e) => {
|
||||
log::error!("Failed to get username from uid: {}", e);
|
||||
let mut message_stream = create_server_to_client_message_stream(socket);
|
||||
message_stream
|
||||
.send(Response::Error(
|
||||
(concatdoc! {
|
||||
"Server failed to get user data from the system\n",
|
||||
"Please check the server logs or contact the system administrators"
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.await
|
||||
.ok();
|
||||
anyhow::bail!("Failed to get username from uid: {}", e);
|
||||
}
|
||||
};
|
||||
|
||||
session_handler_with_unix_user(socket, &unix_user, db_pool).await
|
||||
}
|
||||
|
||||
pub async fn session_handler_with_unix_user(
|
||||
socket: UnixStream,
|
||||
unix_user: &UnixUser,
|
||||
db_pool: MySqlPool,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut message_stream = create_server_to_client_message_stream(socket);
|
||||
|
||||
log::debug!("Opening connection to database");
|
||||
|
||||
log::debug!("Requesting database connection from pool");
|
||||
let mut db_connection = match db_pool.acquire().await {
|
||||
Ok(connection) => connection,
|
||||
Err(err) => {
|
||||
@@ -55,13 +98,12 @@ pub async fn session_handler(
|
||||
return Err(err.into());
|
||||
}
|
||||
};
|
||||
|
||||
log::debug!("Successfully connected to database");
|
||||
log::debug!("Successfully acquired database connection from pool");
|
||||
|
||||
let result =
|
||||
session_handler_with_db_connection(message_stream, unix_user, &mut db_connection).await;
|
||||
session_handler_with_db_connection(message_stream, &unix_user, &mut db_connection).await;
|
||||
|
||||
close_or_ignore_db_connection(db_connection).await;
|
||||
log::debug!("Releasing database connection back to pool");
|
||||
|
||||
result
|
||||
}
|
||||
@@ -192,11 +234,3 @@ async fn session_handler_with_db_connection(
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close_or_ignore_db_connection(db_connection: PoolConnection<MySql>) {
|
||||
if let Err(e) = db_connection.close().await {
|
||||
log::error!("Failed to close database connection: {}", e);
|
||||
log::error!("{}", e);
|
||||
log::error!("Ignoring...");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,22 +7,14 @@ use std::{
|
||||
};
|
||||
|
||||
use anyhow::{Context, anyhow};
|
||||
use futures_util::SinkExt;
|
||||
use indoc::concatdoc;
|
||||
use sqlx::MySqlPool;
|
||||
use tokio::{net::UnixListener as TokioUnixListener, task::JoinHandle, time::interval};
|
||||
use tokio_util::task::TaskTracker;
|
||||
// use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::{
|
||||
core::{
|
||||
common::UnixUser,
|
||||
protocol::{Response, create_server_to_client_message_stream},
|
||||
},
|
||||
server::{
|
||||
config::{MysqlConfig, ServerConfig},
|
||||
session_handler::session_handler,
|
||||
},
|
||||
use crate::server::{
|
||||
config::{MysqlConfig, ServerConfig},
|
||||
session_handler::session_handler,
|
||||
};
|
||||
|
||||
// TODO: implement graceful shutdown and graceful restarts
|
||||
@@ -48,6 +40,12 @@ pub struct Supervisor {
|
||||
|
||||
impl Supervisor {
|
||||
pub async fn new(config: ServerConfig, systemd_mode: bool) -> anyhow::Result<Self> {
|
||||
log::debug!("Starting server supervisor");
|
||||
log::debug!(
|
||||
"Running in tokio with {} worker threads",
|
||||
tokio::runtime::Handle::current().metrics().num_workers()
|
||||
);
|
||||
|
||||
let mut watchdog_duration = None;
|
||||
let mut watchdog_micro_seconds = 0;
|
||||
let watchdog_task =
|
||||
@@ -120,7 +118,9 @@ fn spawn_watchdog_task(duration: Duration) -> JoinHandle<()> {
|
||||
duration.div_f32(2.0).as_millis()
|
||||
);
|
||||
loop {
|
||||
log::trace!("Waiting for next watchdog interval");
|
||||
interval.tick().await;
|
||||
log::trace!("Sending ping to systemd watchdog");
|
||||
if let Err(err) = sd_notify::notify(false, &[sd_notify::NotifyState::Watchdog]) {
|
||||
log::warn!("Failed to notify systemd watchdog: {}", err);
|
||||
} else {
|
||||
@@ -137,6 +137,7 @@ fn spawn_status_notifier_task(connection_counter: std::sync::Arc<()>) -> JoinHan
|
||||
tokio::spawn(async move {
|
||||
let mut interval = interval(STATUS_UPDATE_INTERVAL_SECS);
|
||||
loop {
|
||||
log::trace!("Waiting for next status update interval");
|
||||
interval.tick().await;
|
||||
log::trace!("Updating systemd status notification");
|
||||
let count = Arc::strong_count(&connection_counter) - NON_CONNECTION_ARC_COUNT;
|
||||
@@ -145,7 +146,14 @@ fn spawn_status_notifier_task(connection_counter: std::sync::Arc<()>) -> JoinHan
|
||||
} else {
|
||||
"Waiting for connections".to_string()
|
||||
};
|
||||
sd_notify::notify(false, &[sd_notify::NotifyState::Status(message.as_str())]).ok();
|
||||
|
||||
if let Err(e) =
|
||||
sd_notify::notify(false, &[sd_notify::NotifyState::Status(message.as_str())])
|
||||
{
|
||||
log::warn!("Failed to send systemd status notification: {}", e);
|
||||
} else {
|
||||
log::trace!("Systemd status notification sent: {}", message);
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -195,7 +203,7 @@ async fn create_db_connection_pool(config: &MysqlConfig) -> anyhow::Result<MySql
|
||||
|
||||
config.log_connection_notice();
|
||||
|
||||
match tokio::time::timeout(
|
||||
let pool = match tokio::time::timeout(
|
||||
Duration::from_secs(config.timeout),
|
||||
MySqlPool::connect_with(mysql_config),
|
||||
)
|
||||
@@ -204,7 +212,16 @@ async fn create_db_connection_pool(config: &MysqlConfig) -> anyhow::Result<MySql
|
||||
Ok(connection) => connection.context("Failed to connect to the database"),
|
||||
Err(_) => Err(anyhow!("Timed out after {} seconds", config.timeout))
|
||||
.context("Failed to connect to the database"),
|
||||
}
|
||||
}?;
|
||||
|
||||
let pool_opts = pool.options();
|
||||
log::debug!(
|
||||
"Successfully opened database connection pool with options (max_connections: {}, min_connections: {})",
|
||||
pool_opts.get_max_connections(),
|
||||
pool_opts.get_min_connections(),
|
||||
);
|
||||
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
// fn spawn_signal_handler_task(
|
||||
@@ -242,57 +259,18 @@ async fn spawn_listener_task(
|
||||
|
||||
while let Ok((conn, _addr)) = listener.accept().await {
|
||||
log::debug!("Got new connection");
|
||||
|
||||
let uid = match conn.peer_cred() {
|
||||
Ok(cred) => cred.uid(),
|
||||
Err(e) => {
|
||||
log::error!("Failed to get peer credentials from socket: {}", e);
|
||||
let mut message_stream = create_server_to_client_message_stream(conn);
|
||||
message_stream
|
||||
.send(Response::Error(
|
||||
(concatdoc! {
|
||||
"Server failed to get peer credentials from socket\n",
|
||||
"Please check the server logs or contact the system administrators"
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.await
|
||||
.ok();
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
log::debug!("Validated peer UID: {}", uid);
|
||||
|
||||
let db_pool_clone = db_pool.clone();
|
||||
let _connection_counter_guard = Arc::clone(&connection_counter);
|
||||
|
||||
let unix_user = match UnixUser::from_uid(uid) {
|
||||
Ok(user) => user,
|
||||
Err(e) => {
|
||||
log::error!("Failed to get username from uid: {}", e);
|
||||
let mut message_stream = create_server_to_client_message_stream(conn);
|
||||
message_stream
|
||||
.send(Response::Error(
|
||||
(concatdoc! {
|
||||
"Server failed to get user data from the system\n",
|
||||
"Please check the server logs or contact the system administrators"
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.await
|
||||
.ok();
|
||||
continue;
|
||||
tokio::spawn(async {
|
||||
let _guard = _connection_counter_guard;
|
||||
log::debug!("Running session handler");
|
||||
match session_handler(conn, db_pool_clone).await {
|
||||
Ok(()) => {}
|
||||
Err(e) => {
|
||||
log::error!("Failed to run server: {}", e);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
log::info!("Accepted connection from UNIX user: {}", unix_user.username);
|
||||
|
||||
match session_handler(conn, &unix_user, db_pool.clone()).await {
|
||||
Ok(()) => {}
|
||||
Err(e) => {
|
||||
log::error!("Failed to run server: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
Reference in New Issue
Block a user