eba4fd1
use crate::{
sec::encrypt, DatabasePool, Session, SessionConfig, SessionData, SessionError, SessionTimers,
};
use axum::extract::FromRequestParts;
use chrono::{Duration, Utc};
use dashmap::DashMap;
#[cfg(feature = "key-store")]
use fastbloom_rs::Deletable;
#[cfg(feature = "key-store")]
use fastbloom_rs::{CountingBloomFilter, FilterBuilder, Membership};
use http::{request::Parts, StatusCode};
use serde::Serialize;
use std::{fmt::Debug, sync::Arc};
use tokio::sync::RwLock;
/// Contains the main Services storage for all session's and database access for persistent Sessions.
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore};
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config).await.unwrap();
/// ```
///
#[derive(Clone, Debug)]
pub struct SessionStore<T>
where
T: DatabasePool + Clone + Debug + Sync + Send + 'static,
{
/// Client for the database.
pub client: Option<T>,
/// locked Hashmap containing UserID and their session data.
pub(crate) inner: Arc<DashMap<String, SessionData>>,
/// Session Configuration.
pub config: SessionConfig,
/// Session Timers used for Clearing Memory and Database.
pub(crate) timers: Arc<RwLock<SessionTimers>>,
#[cfg(feature = "key-store")]
/// Filter used to keep track of what session IDs exist.
pub(crate) filter: Arc<RwLock<CountingBloomFilter>>,
}
impl<T, S> FromRequestParts<S> for SessionStore<T>
where
T: DatabasePool + Clone + Debug + Sync + Send + 'static,
S: Send + Sync,
{
type Rejection = (http::StatusCode, &'static str);
async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
let session = parts.extensions.get::<Session<T>>().ok_or((
StatusCode::INTERNAL_SERVER_ERROR,
"Can't extract Axum `Session`. Is `SessionLayer` enabled?",
))?;
Ok(session.store.clone())
}
}
impl<T> SessionStore<T>
where
T: DatabasePool + Clone + Debug + Sync + Send + 'static,
{
/// Constructs a New `SessionStore` and Creates the Database Table
/// needed for the Session if it does not exist if client is not `None`.
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore};
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config).await.unwrap();
/// ```
///
#[inline]
pub async fn new(client: Option<T>, config: SessionConfig) -> Result<Self, SessionError> {
if let Some(client) = &client {
client.initiate(&config.database.table_name).await?
}
// If we have a database client then lets also get any SessionId's that Exist within the database
// that are not yet expired.
#[cfg(feature = "key-store")]
let filter = Self::create_filter(&client, &config).await?;
Ok(Self {
client,
inner: Default::default(),
config,
timers: Arc::new(RwLock::new(SessionTimers {
// the first expiry sweep is scheduled one lifetime from start-up
last_expiry_sweep: Utc::now() + Duration::try_hours(1).unwrap_or_default(),
// the first expiry sweep is scheduled one lifetime from start-up
last_database_expiry_sweep: Utc::now() + Duration::try_hours(6).unwrap_or_default(),
})),
#[cfg(feature = "key-store")]
filter: Arc::new(RwLock::new(filter)),
})
}
/// Used to create and Fill the Filter.
#[cfg(feature = "key-store")]
pub(crate) async fn create_filter(
client: &Option<T>,
config: &SessionConfig,
) -> Result<CountingBloomFilter, SessionError> {
let mut filter = FilterBuilder::new(
config.memory.filter_expected_elements,
config.memory.filter_false_positive_probability,
)
.build_counting_bloom_filter();
if config.memory.use_bloom_filters {
// If client exist then lets preload the id's within the database so the filter is accurate.
if let Some(client) = &client {
let ids = client.get_ids(&config.database.table_name).await?;
ids.iter().for_each(|id| filter.add(id.as_bytes()));
}
}
Ok(filter)
}
/// Checks if the database is in persistent mode.
///
/// Returns true if client is Some().
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore};
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config).await.unwrap();
/// let is_persistent = session_store.is_persistent();
/// ```
///
#[inline]
pub fn is_persistent(&self) -> bool {
self.client.is_some()
}
/// Cleans Expired sessions from the Database based on Utc::now().
///
/// If client is None it will return Ok(()).
///
/// # Errors
/// - ['SessionError::Sqlx'] is returned if database connection has failed or user does not have permissions.
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore};
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config).await.unwrap();
/// async {
/// let _ = session_store.cleanup().await.unwrap();
/// };
/// ```
///
#[inline]
pub async fn cleanup(&self) -> Result<Vec<String>, SessionError> {
if let Some(client) = &self.client {
Ok(client
.delete_by_expiry(&self.config.database.table_name)
.await?)
} else {
Ok(Vec::new())
}
}
/// Returns count of existing sessions within database.
///
/// If client is None it will return Ok(0).
///
/// # Errors
/// - ['SessionError::Sqlx'] is returned if database connection has failed or user does not have permissions.
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore};
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config).await.unwrap();
/// async {
/// let count = session_store.count().await.unwrap();
/// };
/// ```
///
#[inline]
pub async fn count(&self) -> Result<i64, SessionError> {
if let Some(client) = &self.client {
let count = client.count(&self.config.database.table_name).await?;
return Ok(count);
}
Ok(0)
}
/// private internal function that loads a session's data from the database using an ID string.
///
/// If client is None it will return Ok(None).
///
/// # Errors
/// - ['SessionError::Sqlx'] is returned if database connection has failed or user does not have permissions.
/// - ['SessionError::SerdeJson'] is returned if it failed to deserialize the sessions data.
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore};
/// use uuid::Uuid;
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config).await.unwrap();
/// let token = Uuid::new_v4();
/// async {
/// let session_data = session_store.load_session(token.to_string()).await.unwrap();
/// };
/// ```
///
pub(crate) async fn load_session(
&self,
cookie_value: String,
) -> Result<Option<SessionData>, SessionError> {
if let Some(client) = &self.client {
let result: Option<String> = client
.load(&cookie_value, &self.config.database.table_name)
.await?;
if let Some(mut session) = result
.map(|session| {
if let Some(key) = self.config.database.database_key.as_ref() {
serde_json::from_str::<SessionData>(
&match encrypt::decrypt(&cookie_value, &session, key) {
Ok(v) => v,
Err(err) => {
tracing::error!(err = %err, "Failed to decrypt Session data from database.");
String::new()
}
}
)
} else {
serde_json::from_str::<SessionData>(&session)
}
})
.transpose()?
{
session.id = cookie_value;
return Ok(Some(session));
}
}
Ok(None)
}
/// private internal function that stores a session's data to the database.
///
/// If client is None it will return Ok(()).
///
/// # Errors
/// - ['SessionError::Sqlx'] is returned if database connection has failed or user does not have permissions.
/// - ['SessionError::SerdeJson'] is returned if it failed to serialize the sessions data.
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore, SessionData};
/// use uuid::Uuid;
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config.clone()).await.unwrap();
/// let token = Uuid::new_v4();
/// let session_data = SessionData::new(token, true, &config);
///
/// async {
/// let _ = session_store.store_session(&session_data).await.unwrap();
/// };
/// ```
///
pub(crate) async fn store_session(&self, session: &SessionData) -> Result<(), SessionError> {
if let Some(client) = &self.client {
client
.store(
&session.id,
&if let Some(key) = self.config.database.database_key.as_ref() {
encrypt::encrypt(&session.id, &serde_json::to_string(session)?, key)
.map_err(|e| {
SessionError::GenericNotSupportedError(format!(
"Error: {e} Occurred when encrypting a Session.",
))
})?
} else {
serde_json::to_string(session)?
},
session.expires.timestamp(),
&self.config.database.table_name,
)
.await?;
}
Ok(())
}
/// Deletes all sessions in the database.
///
/// If client is None it will return Ok(()).
///
/// # Errors
/// - ['SessionError::Sqlx'] is returned if database connection has failed or user does not have permissions.
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore};
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config.clone()).await.unwrap();
///
/// async {
/// let _ = session_store.clear_store().await.unwrap();
/// };
/// ```
///
#[inline]
pub async fn clear_store(&self) -> Result<(), SessionError> {
if let Some(client) = &self.client {
client.delete_all(&self.config.database.table_name).await?;
}
Ok(())
}
/// Deletes all sessions in Memory.
/// This will also Clear those keys from the filter cache if a persistent database does not exist.
///
/// # Examples
/// ```rust ignore
/// use session::{SessionNullPool, SessionConfig, SessionStore};
///
/// let config = SessionConfig::default();
/// let session_store = SessionStore::<SessionNullPool>::new(None, config.clone()).await.unwrap();
///
/// async {
/// let _ = session_store.clear().await.unwrap();
/// };
/// ```
///
#[inline]
pub async fn clear(&mut self) {
#[cfg(feature = "key-store")]
if self.client.is_none() {
let mut filter = self.filter.write().await;
self.inner
.iter()
.for_each(|value| filter.remove(value.key().as_bytes()));
}
self.inner.clear();
}
/// Attempts to load check and clear Data.
///
/// If no session is found returns false.
pub(crate) fn service_session_data(&self, session: &Session<T>) -> bool {
if let Some(mut inner) = self.inner.get_mut(&session.id) {
inner.service_clear(
self.config.memory.memory_lifespan,
self.config.clear_check_on_load,
);
inner.set_request();
return true;
}
false
}
#[inline]
pub(crate) fn renew(&self, id: String) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.renew();
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn destroy(&self, id: String) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.destroy();
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn is_longterm(&self, id: &str) -> bool {
self.inner.get(id).is_some_and(|d| d.longterm)
}
#[inline]
pub(crate) fn set_longterm(&self, id: String, longterm: bool) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.set_longterm(longterm);
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn set_store(&self, id: String, storable: bool) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.set_store(storable);
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn update(&self, id: String) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.update();
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn get<N: serde::de::DeserializeOwned>(&self, id: String, key: &str) -> Option<N> {
if let Some(instance) = self.inner.get(&id) {
instance.get(key)
} else {
tracing::warn!("Session data unexpectedly missing");
None
}
}
#[inline]
pub(crate) fn get_remove<N: serde::de::DeserializeOwned>(
&self,
id: String,
key: &str,
) -> Option<N> {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.get_remove(key)
} else {
tracing::warn!("Session data unexpectedly missing");
None
}
}
#[inline]
pub(crate) fn set(&self, id: String, key: &str, value: impl Serialize) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.set(key, value);
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn remove(&self, id: String, key: &str) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.remove(key);
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn clear_session_data(&self, id: String) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.clear();
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn set_session_request(&self, id: String) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.set_request();
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn remove_session_request(&self, id: String) {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.remove_request();
} else {
tracing::warn!("Session data unexpectedly missing");
}
}
#[inline]
pub(crate) fn is_session_parallel(&self, id: String) -> bool {
if let Some(instance) = self.inner.get(&id) {
instance.is_parallel()
} else {
tracing::warn!("Session data unexpectedly missing");
false
}
}
#[inline]
pub(crate) async fn count_sessions(&self) -> i64 {
if self.is_persistent() {
self.count().await.unwrap_or(0i64)
} else {
self.inner.len() as i64
}
}
#[inline]
pub(crate) fn auto_handles_expiry(&self) -> bool {
if let Some(client) = &self.client {
client.auto_handles_expiry()
} else {
false
}
}
#[cfg(feature = "advanced")]
#[cfg_attr(docsrs, doc(cfg(feature = "advanced")))]
#[inline]
pub fn verify(&self, id: String) -> Result<(), SessionError> {
if let Some(instance) = self.inner.get(&id) {
if instance.expires < Utc::now() {
Err(SessionError::OldSessionError)
} else {
Ok(())
}
} else {
Err(SessionError::NoSessionError)
}
}
#[cfg(feature = "advanced")]
#[cfg_attr(docsrs, doc(cfg(feature = "advanced")))]
#[inline]
pub fn update_database_expires(&self, id: String) -> Result<(), SessionError> {
if let Some(mut instance) = self.inner.get_mut(&id) {
if instance.longterm {
instance.expires = Utc::now() + self.config.max_lifespan;
} else {
instance.expires = Utc::now() + self.config.lifespan;
}
Ok(())
} else {
Err(SessionError::NoSessionError)
}
}
#[cfg(feature = "advanced")]
#[cfg_attr(docsrs, doc(cfg(feature = "advanced")))]
#[inline]
pub fn update_memory_expires(&self, id: String) -> Result<(), SessionError> {
if let Some(mut instance) = self.inner.get_mut(&id) {
instance.autoremove = Utc::now() + self.config.memory.memory_lifespan;
Ok(())
} else {
Err(SessionError::NoSessionError)
}
}
#[cfg(feature = "advanced")]
#[cfg_attr(docsrs, doc(cfg(feature = "advanced")))]
#[inline]
pub async fn force_database_update(&self, id: String) -> Result<(), SessionError> {
let session = if let Some(instance) = self.inner.get(&id) {
instance.clone()
} else {
return Err(SessionError::NoSessionError);
};
self.store_session(&session).await
}
#[cfg(feature = "advanced")]
#[cfg_attr(docsrs, doc(cfg(feature = "advanced")))]
#[inline]
pub fn memory_remove_session(&self, id: String) -> Result<(), SessionError> {
let is_parallel = if let Some(mut instance) = self.inner.get_mut(&id) {
instance.remove_request();
instance.is_parallel()
} else {
return Err(SessionError::NoSessionError);
};
if is_parallel {
let _ = self.inner.remove(&id);
}
Ok(())
}
#[cfg(feature = "advanced")]
#[cfg_attr(docsrs, doc(cfg(feature = "advanced")))]
#[inline]
pub fn force_memory_remove_session(&self, id: String) -> Result<(), SessionError> {
if self.inner.remove(&id).is_some() {
Ok(())
} else {
Err(SessionError::NoSessionError)
}
}
#[inline]
pub(crate) async fn database_remove_session(&self, id: String) -> Result<(), SessionError> {
if let Some(client) = &self.client {
client
.delete_one_by_id(&id, &self.config.database.table_name)
.await?;
}
Ok(())
}
#[cfg(feature = "advanced")]
#[cfg_attr(docsrs, doc(cfg(feature = "advanced")))]
#[inline]
pub async fn remove_stored_database_session(&self, id: String) -> Result<(), SessionError> {
self.database_remove_session(id).await
}
}