pub use tower_sessions::{Session, SessionManagerLayer};
use async_trait::async_trait;
use larust_core::AppError;
use larust_orm::Backend;
use sqlx::AnyPool;
use std::io;
use std::path::{Path, PathBuf};
use tower_sessions::session::{Id, Record};
use tower_sessions::session_store::{self, ExpiredDeletion};
use tower_sessions::SessionStore;
const EXPIRED_SESSION_CLEANUP_INTERVAL: std::time::Duration = std::time::Duration::from_secs(3600);
pub async fn take<T: serde::de::DeserializeOwned>(
session: &Session,
key: &str,
) -> Result<Option<T>, tower_sessions::session::Error> {
if session.get_value(key).await?.is_none() {
return Ok(None);
}
session.remove(key).await
}
#[derive(Clone, Debug)]
pub struct AnySessionStore {
pool: AnyPool,
}
impl AnySessionStore {
pub fn new(pool: AnyPool) -> Self {
Self { pool }
}
async fn migrate(&self) -> Result<(), sqlx::Error> {
let create_table = match larust_orm::backend() {
Backend::Sqlite => {
"CREATE TABLE IF NOT EXISTS sessions (\
id TEXT PRIMARY KEY, \
data TEXT NOT NULL, \
expiry_at INTEGER NOT NULL\
)"
}
Backend::MySql => {
"CREATE TABLE IF NOT EXISTS sessions (\
id VARCHAR(32) PRIMARY KEY, \
data VARCHAR(4000) NOT NULL, \
expiry_at INTEGER NOT NULL\
)"
}
Backend::Postgres => {
"CREATE TABLE IF NOT EXISTS sessions (\
id TEXT PRIMARY KEY, \
data TEXT NOT NULL, \
expiry_at INTEGER NOT NULL\
)"
}
};
sqlx::query(create_table).execute(&self.pool).await?;
Ok(())
}
}
fn now_unix_secs() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system clock is before the Unix epoch")
.as_secs() as i64
}
fn backend_error(source: sqlx::Error) -> session_store::Error {
session_store::Error::Backend(source.to_string())
}
fn encode(record: &Record) -> session_store::Result<String> {
serde_json::to_string(record).map_err(|source| session_store::Error::Encode(source.to_string()))
}
fn decode(data: &str) -> session_store::Result<Record> {
serde_json::from_str(data).map_err(|source| session_store::Error::Decode(source.to_string()))
}
const MAX_ID_COLLISION_RETRIES: u8 = 5;
#[async_trait]
impl SessionStore for AnySessionStore {
async fn create(&self, record: &mut Record) -> session_store::Result<()> {
let insert_sql = match larust_orm::backend() {
Backend::Sqlite | Backend::MySql => {
"INSERT INTO sessions (id, data, expiry_at) VALUES (?, ?, ?)"
}
Backend::Postgres => "INSERT INTO sessions (id, data, expiry_at) VALUES ($1, $2, $3)",
};
for _ in 0..MAX_ID_COLLISION_RETRIES {
let data = encode(record)?;
let expiry_at = record.expiry_date.unix_timestamp();
match sqlx::query(insert_sql)
.bind(record.id.to_string())
.bind(data)
.bind(expiry_at)
.execute(&self.pool)
.await
{
Ok(_) => return Ok(()),
Err(source) => {
let is_collision = source
.as_database_error()
.is_some_and(|e| e.is_unique_violation());
if !is_collision {
return Err(backend_error(source));
}
record.id = Id::default();
}
}
}
Err(session_store::Error::Backend(format!(
"couldn't create a session after {MAX_ID_COLLISION_RETRIES} ID collisions in a row"
)))
}
async fn save(&self, record: &Record) -> session_store::Result<()> {
let data = encode(record)?;
let expiry_at = record.expiry_date.unix_timestamp();
let upsert_sql = match larust_orm::backend() {
Backend::Sqlite => {
"INSERT INTO sessions (id, data, expiry_at) VALUES (?, ?, ?) \
ON CONFLICT(id) DO UPDATE SET data = excluded.data, expiry_at = excluded.expiry_at"
}
Backend::MySql => {
"INSERT INTO sessions (id, data, expiry_at) VALUES (?, ?, ?) \
ON DUPLICATE KEY UPDATE data = VALUES(data), expiry_at = VALUES(expiry_at)"
}
Backend::Postgres => {
"INSERT INTO sessions (id, data, expiry_at) VALUES ($1, $2, $3) \
ON CONFLICT(id) DO UPDATE SET data = excluded.data, expiry_at = excluded.expiry_at"
}
};
sqlx::query(upsert_sql)
.bind(record.id.to_string())
.bind(data)
.bind(expiry_at)
.execute(&self.pool)
.await
.map_err(backend_error)?;
Ok(())
}
async fn load(&self, session_id: &Id) -> session_store::Result<Option<Record>> {
let backend = larust_orm::backend();
let sql = format!(
"SELECT data FROM sessions WHERE id = {} AND expiry_at > {}",
larust_orm::placeholder(backend, 1),
larust_orm::placeholder(backend, 2),
);
let row: Option<(String,)> = sqlx::query_as(&sql)
.bind(session_id.to_string())
.bind(now_unix_secs())
.fetch_optional(&self.pool)
.await
.map_err(backend_error)?;
row.map(|(data,)| decode(&data)).transpose()
}
async fn delete(&self, session_id: &Id) -> session_store::Result<()> {
let sql = format!(
"DELETE FROM sessions WHERE id = {}",
larust_orm::placeholder(larust_orm::backend(), 1)
);
sqlx::query(&sql)
.bind(session_id.to_string())
.execute(&self.pool)
.await
.map_err(backend_error)?;
Ok(())
}
}
#[async_trait]
impl ExpiredDeletion for AnySessionStore {
async fn delete_expired(&self) -> session_store::Result<()> {
let sql = format!(
"DELETE FROM sessions WHERE expiry_at <= {}",
larust_orm::placeholder(larust_orm::backend(), 1)
);
sqlx::query(&sql)
.bind(now_unix_secs())
.execute(&self.pool)
.await
.map_err(backend_error)?;
Ok(())
}
}
fn io_error(source: io::Error) -> session_store::Error {
session_store::Error::Backend(source.to_string())
}
#[derive(Clone, Debug)]
pub struct FileSessionStore {
dir: PathBuf,
}
impl FileSessionStore {
pub fn new(dir: impl Into<PathBuf>) -> Self {
Self { dir: dir.into() }
}
async fn ensure_dir(&self) -> io::Result<()> {
tokio::fs::create_dir_all(&self.dir).await
}
fn path_for(&self, id: &Id) -> PathBuf {
self.dir.join(id.to_string())
}
async fn write_atomically(path: &Path, data: &[u8]) -> io::Result<()> {
static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let tmp_path = path.with_extension(format!("tmp.{}.{n}", std::process::id()));
tokio::fs::write(&tmp_path, data).await?;
tokio::fs::rename(&tmp_path, path).await
}
}
fn is_session_file(path: &Path) -> bool {
path.extension().is_none()
}
#[async_trait]
impl SessionStore for FileSessionStore {
async fn create(&self, record: &mut Record) -> session_store::Result<()> {
use tokio::io::AsyncWriteExt;
for _ in 0..MAX_ID_COLLISION_RETRIES {
let path = self.path_for(&record.id);
let data = encode(record)?;
match tokio::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&path)
.await
{
Ok(mut file) => {
file.write_all(data.as_bytes()).await.map_err(io_error)?;
file.flush().await.map_err(io_error)?;
return Ok(());
}
Err(source) if source.kind() == io::ErrorKind::AlreadyExists => {
record.id = Id::default();
}
Err(source) => return Err(io_error(source)),
}
}
Err(session_store::Error::Backend(format!(
"couldn't create a session after {MAX_ID_COLLISION_RETRIES} ID collisions in a row"
)))
}
async fn save(&self, record: &Record) -> session_store::Result<()> {
let path = self.path_for(&record.id);
let data = encode(record)?;
Self::write_atomically(&path, data.as_bytes())
.await
.map_err(io_error)
}
async fn load(&self, session_id: &Id) -> session_store::Result<Option<Record>> {
match tokio::fs::read_to_string(self.path_for(session_id)).await {
Ok(data) => {
let record = decode(&data)?;
if record.expiry_date.unix_timestamp() > now_unix_secs() {
Ok(Some(record))
} else {
Ok(None)
}
}
Err(source) if source.kind() == io::ErrorKind::NotFound => Ok(None),
Err(source) => Err(io_error(source)),
}
}
async fn delete(&self, session_id: &Id) -> session_store::Result<()> {
match tokio::fs::remove_file(self.path_for(session_id)).await {
Ok(()) => Ok(()),
Err(source) if source.kind() == io::ErrorKind::NotFound => Ok(()),
Err(source) => Err(io_error(source)),
}
}
}
#[async_trait]
impl ExpiredDeletion for FileSessionStore {
async fn delete_expired(&self) -> session_store::Result<()> {
let mut entries = match tokio::fs::read_dir(&self.dir).await {
Ok(entries) => entries,
Err(source) if source.kind() == io::ErrorKind::NotFound => return Ok(()),
Err(source) => return Err(io_error(source)),
};
let now = now_unix_secs();
while let Some(entry) = entries.next_entry().await.map_err(io_error)? {
let path = entry.path();
if !is_session_file(&path) {
continue;
}
let Ok(data) = tokio::fs::read_to_string(&path).await else {
continue;
};
let Ok(record) = decode(&data) else {
continue;
};
if record.expiry_date.unix_timestamp() <= now {
let _ = tokio::fs::remove_file(&path).await;
}
}
Ok(())
}
}
#[derive(Clone, Debug)]
pub enum SessionBackend {
Database(AnySessionStore),
File(FileSessionStore),
}
#[async_trait]
impl SessionStore for SessionBackend {
async fn create(&self, record: &mut Record) -> session_store::Result<()> {
match self {
Self::Database(store) => store.create(record).await,
Self::File(store) => store.create(record).await,
}
}
async fn save(&self, record: &Record) -> session_store::Result<()> {
match self {
Self::Database(store) => store.save(record).await,
Self::File(store) => store.save(record).await,
}
}
async fn load(&self, session_id: &Id) -> session_store::Result<Option<Record>> {
match self {
Self::Database(store) => store.load(session_id).await,
Self::File(store) => store.load(session_id).await,
}
}
async fn delete(&self, session_id: &Id) -> session_store::Result<()> {
match self {
Self::Database(store) => store.delete(session_id).await,
Self::File(store) => store.delete(session_id).await,
}
}
}
#[async_trait]
impl ExpiredDeletion for SessionBackend {
async fn delete_expired(&self) -> session_store::Result<()> {
match self {
Self::Database(store) => store.delete_expired().await,
Self::File(store) => store.delete_expired().await,
}
}
}
pub fn resolve_backend(pool: &AnyPool) -> SessionBackend {
let driver =
larust_core::try_config().map_or("database", |config| config.session_driver.as_str());
match driver {
"file" => {
let dir = larust_core::try_paths()
.map(|paths| paths.sessions())
.unwrap_or_else(|| PathBuf::from("storage/sessions"));
SessionBackend::File(FileSessionStore::new(dir))
}
_ => SessionBackend::Database(AnySessionStore::new(pool.clone())),
}
}
pub async fn session_layer(
pool: &AnyPool,
secure: bool,
) -> Result<SessionManagerLayer<SessionBackend>, AppError> {
let store = resolve_backend(pool);
match &store {
SessionBackend::File(file_store) => file_store
.ensure_dir()
.await
.map_err(|source| AppError::Internal(Box::new(source)))?,
SessionBackend::Database(db_store) => db_store
.migrate()
.await
.map_err(|source| AppError::Internal(Box::new(source)))?,
}
let cleanup_store = store.clone();
tokio::spawn(async move {
let mut interval = tokio::time::interval(EXPIRED_SESSION_CLEANUP_INTERVAL);
interval.tick().await; loop {
interval.tick().await;
if let Err(error) = cleanup_store.delete_expired().await {
tracing::error!(%error, "expired session cleanup task stopped unexpectedly");
break;
}
}
});
Ok(SessionManagerLayer::new(store)
.with_secure(secure)
.with_name(cookie_name()))
}
pub fn cookie_name() -> String {
let app_name = larust_core::try_config()
.map(|config| config.app_name.as_str())
.unwrap_or("app");
session_cookie_name(app_name)
}
fn session_cookie_name(app_name: &str) -> String {
let safe: String = app_name
.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect();
format!("larust_{safe}_session")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cookie_name_is_derived_from_the_app_name() {
assert_eq!(session_cookie_name("blog"), "larust_blog_session");
}
#[test]
fn cookie_name_sanitizes_characters_a_cookie_token_can_t_contain() {
assert_eq!(
session_cookie_name("My App! 2.0"),
"larust_My_App__2_0_session"
);
}
#[test]
fn different_app_names_never_collide_on_the_same_cookie_name() {
assert_ne!(session_cookie_name("Larust"), session_cookie_name("blog"));
}
async fn connect_test_db() -> AnyPool {
let dir = tempfile::tempdir().unwrap().keep();
let database_url = format!("sqlite://{}/test.sqlite", dir.display());
larust_orm::connect(&database_url).await.unwrap();
larust_orm::pool().unwrap().clone()
}
#[tokio::test]
async fn any_session_store_behaves_correctly_across_every_scenario() {
let pool = connect_test_db().await;
let store = AnySessionStore::new(pool);
store.migrate().await.unwrap();
let mut record = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
record
.data
.insert("greeting".to_string(), serde_json::json!("hi"));
store.save(&record).await.unwrap();
let loaded = store.load(&record.id).await.unwrap().unwrap();
assert_eq!(loaded.data["greeting"], serde_json::json!("hi"));
let expired = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() - time::Duration::hours(1),
};
store.save(&expired).await.unwrap();
assert!(store.load(&expired.id).await.unwrap().is_none());
let doomed = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
store.save(&doomed).await.unwrap();
store.delete(&doomed.id).await.unwrap();
assert!(store.load(&doomed.id).await.unwrap().is_none());
let mut fresh = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
fresh
.data
.insert("greeting".to_string(), serde_json::json!("fresh"));
let fresh_id = fresh.id;
store.create(&mut fresh).await.unwrap();
assert_eq!(
fresh.id, fresh_id,
"a non-colliding create must not change the id"
);
let loaded = store.load(&fresh_id).await.unwrap().unwrap();
assert_eq!(loaded.data["greeting"], serde_json::json!("fresh"));
let mut colliding = Record {
id: fresh_id,
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
colliding
.data
.insert("greeting".to_string(), serde_json::json!("colliding"));
store.create(&mut colliding).await.unwrap();
assert_ne!(
colliding.id, fresh_id,
"create must regenerate the id on a real collision, not overwrite the existing row"
);
let original = store.load(&fresh_id).await.unwrap().unwrap();
assert_eq!(original.data["greeting"], serde_json::json!("fresh"));
let new_one = store.load(&colliding.id).await.unwrap().unwrap();
assert_eq!(new_one.data["greeting"], serde_json::json!("colliding"));
let session = Session::new(None, std::sync::Arc::new(store.clone()), None);
let absent: Option<String> = take(&session, "success").await.unwrap();
assert_eq!(absent, None);
assert!(
!session.is_modified(),
"checking for an absent flash key must not mark the session dirty"
);
session.insert("success", "saved!").await.unwrap();
let value: Option<String> = take(&session, "success").await.unwrap();
assert_eq!(value.as_deref(), Some("saved!"));
assert!(session.is_modified());
let gone: Option<String> = session.get("success").await.unwrap();
assert_eq!(gone, None);
let concurrent_saves: Vec<_> = (0..32)
.map(|i| {
let store = store.clone();
tokio::spawn(async move {
let mut record = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
record.data.insert("i".to_string(), serde_json::json!(i));
store.save(&record).await
})
})
.collect();
for task in concurrent_saves {
task.await
.expect("task panicked")
.expect("concurrent session save must not fail with a locked database");
}
}
#[tokio::test]
async fn file_session_store_behaves_correctly_across_every_scenario() {
let dir = tempfile::tempdir().unwrap();
let store = FileSessionStore::new(dir.path());
let mut record = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
record
.data
.insert("greeting".to_string(), serde_json::json!("hi"));
store.save(&record).await.unwrap();
let loaded = store.load(&record.id).await.unwrap().unwrap();
assert_eq!(loaded.data["greeting"], serde_json::json!("hi"));
let expired = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() - time::Duration::hours(1),
};
store.save(&expired).await.unwrap();
assert!(store.load(&expired.id).await.unwrap().is_none());
assert!(
dir.path().join(expired.id.to_string()).exists(),
"an expired session's file must still exist until delete_expired sweeps it"
);
let doomed = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
store.save(&doomed).await.unwrap();
store.delete(&doomed.id).await.unwrap();
assert!(store.load(&doomed.id).await.unwrap().is_none());
store.delete(&doomed.id).await.unwrap();
let mut fresh = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
fresh
.data
.insert("greeting".to_string(), serde_json::json!("fresh"));
let fresh_id = fresh.id;
store.create(&mut fresh).await.unwrap();
assert_eq!(
fresh.id, fresh_id,
"a non-colliding create must not change the id"
);
let loaded = store.load(&fresh_id).await.unwrap().unwrap();
assert_eq!(loaded.data["greeting"], serde_json::json!("fresh"));
let mut colliding = Record {
id: fresh_id,
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
colliding
.data
.insert("greeting".to_string(), serde_json::json!("colliding"));
store.create(&mut colliding).await.unwrap();
assert_ne!(
colliding.id, fresh_id,
"create must regenerate the id on a real collision, not overwrite the existing file"
);
let original = store.load(&fresh_id).await.unwrap().unwrap();
assert_eq!(original.data["greeting"], serde_json::json!("fresh"));
let new_one = store.load(&colliding.id).await.unwrap().unwrap();
assert_eq!(new_one.data["greeting"], serde_json::json!("colliding"));
store.delete_expired().await.unwrap();
assert!(!dir.path().join(expired.id.to_string()).exists());
assert!(dir.path().join(fresh_id.to_string()).exists());
assert!(dir.path().join(colliding.id.to_string()).exists());
tokio::fs::write(dir.path().join("orphaned.tmp.999.1"), b"not json")
.await
.unwrap();
store.delete_expired().await.unwrap();
assert!(dir.path().join("orphaned.tmp.999.1").exists());
let concurrent_creates: Vec<_> = (0..32)
.map(|i| {
let store = store.clone();
tokio::spawn(async move {
let mut record = Record {
id: Id::default(),
data: Default::default(),
expiry_date: time::OffsetDateTime::now_utc() + time::Duration::hours(1),
};
record.data.insert("i".to_string(), serde_json::json!(i));
store.create(&mut record).await
})
})
.collect();
for task in concurrent_creates {
task.await
.expect("task panicked")
.expect("concurrent session create must not fail");
}
}
}