4c83001
use async_trait::async_trait;
use axum_session::{DatabaseError, DatabasePool, SessionError};
use chrono::Utc;
use config::session::SessionConfig;
use export::{AppRedis, RedisConnection};
use crate::config::build_session_config;
pub type Session = axum_session::Session<SessionRedisPool>;
pub type ReadOnlySession = axum_session::ReadOnlySession<SessionRedisPool>;
pub type SessionStore = axum_session::SessionStore<SessionRedisPool>;
pub type SessionLayer = axum_session::SessionLayer<SessionRedisPool>;
/// Keys deleted by a single `UNLINK` call in standalone mode. Kept at one key
/// per call in cluster mode, where multi-key commands would cross slots.
const UNLINK_BATCH: usize = 100;
/// Hint for `SCAN COUNT`, trading a few bigger replies for far fewer round trips.
const SCAN_BATCH: usize = 1000;
#[derive(Clone)]
pub struct SessionRedisPool {
redis: AppRedis,
}
impl SessionRedisPool {
pub fn new(redis: AppRedis) -> Self {
Self { redis }
}
pub fn redis(&self) -> &AppRedis {
&self.redis
}
}
impl From<AppRedis> for SessionRedisPool {
fn from(redis: AppRedis) -> Self {
Self::new(redis)
}
}
impl std::fmt::Debug for SessionRedisPool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SessionRedisPool").finish()
}
}
pub async fn connect(
redis: AppRedis,
config: &SessionConfig,
) -> Result<SessionStore, SessionError> {
let config = build_session_config(config).map_err(|err| {
SessionError::GenericNotSupportedError(format!("invalid session configuration: {err}"))
})?;
SessionStore::new(Some(SessionRedisPool::new(redis)), config).await
}
macro_rules! dispatch {
($conn:expr, |$c:ident| $body:expr) => {
match &mut $conn {
RedisConnection::Standalone($c) => $body,
RedisConnection::Cluster($c) => $body,
}
};
}
fn session_key(id: &str, table_name: &str) -> String {
if table_name.is_empty() {
id.to_string()
} else {
format!("{table_name}:{id}")
}
}
/// SCANs `pattern` across the whole store: a cursor loop on the standalone
/// connection, or the same loop on every cluster master with results merged.
async fn scan_keys(redis: &AppRedis, pattern: &str) -> Result<Vec<String>, DatabaseError> {
let mut keys = match redis.connection_clone() {
RedisConnection::Standalone(mut conn) => scan_conn(&mut conn, pattern).await?,
RedisConnection::Cluster(_) => {
let mut keys = Vec::new();
for mut conn in redis
.cluster_master_connections()
.await
.map_err(|err| DatabaseError::GenericSelectError(err.to_string()))?
{
keys.extend(scan_conn(&mut conn, pattern).await?);
}
keys
}
};
keys.sort();
keys.dedup();
Ok(keys)
}
async fn scan_conn<C>(conn: &mut C, pattern: &str) -> Result<Vec<String>, DatabaseError>
where
C: redis::aio::ConnectionLike,
{
let mut keys = Vec::new();
let mut cursor: u64 = 0;
loop {
let (next, batch): (u64, Vec<String>) = redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(pattern)
.arg("COUNT")
.arg(SCAN_BATCH)
.query_async(conn)
.await
.map_err(|err| DatabaseError::GenericSelectError(err.to_string()))?;
keys.extend(batch);
cursor = next;
if cursor == 0 {
break;
}
}
Ok(keys)
}
#[async_trait]
impl DatabasePool for SessionRedisPool {
async fn initiate(&self, _table_name: &str) -> Result<(), DatabaseError> {
Ok(())
}
async fn delete_by_expiry(&self, _table_name: &str) -> Result<Vec<String>, DatabaseError> {
Ok(Vec::new())
}
async fn count(&self, table_name: &str) -> Result<i64, DatabaseError> {
ensure_table(table_name)?;
let count = scan_keys(&self.redis, &format!("{table_name}:*")).await?.len() as i64;
Ok(count)
}
async fn store(
&self,
id: &str,
session: &str,
expires: i64,
table_name: &str,
) -> Result<(), DatabaseError> {
// axum_session passes `expires` as a unix timestamp. Store a relative
// TTL instead of `EXAT` so Redis < 6.2 keeps working.
let ttl_ms = (expires - Utc::now().timestamp()) * 1000;
if ttl_ms <= 0 {
return Ok(());
}
let key = session_key(id, table_name);
let mut conn = self.redis.connection_clone();
dispatch!(conn, |c| {
redis::cmd("SET")
.arg(&key)
.arg(session)
.arg("PX")
.arg(ttl_ms)
.exec_async(c)
.await
})
.map_err(|err| DatabaseError::GenericInsertError(err.to_string()))
}
async fn load(&self, id: &str, table_name: &str) -> Result<Option<String>, DatabaseError> {
let key = session_key(id, table_name);
let mut conn = self.redis.connection_clone();
dispatch!(conn, |c| redis::cmd("GET").arg(&key).query_async(c).await)
.map_err(|err| DatabaseError::GenericSelectError(err.to_string()))
}
async fn delete_one_by_id(&self, id: &str, table_name: &str) -> Result<(), DatabaseError> {
let key = session_key(id, table_name);
let mut conn = self.redis.connection_clone();
dispatch!(conn, |c| redis::cmd("UNLINK").arg(&key).exec_async(c).await)
.map_err(|err| DatabaseError::GenericDeleteError(err.to_string()))
}
async fn exists(&self, id: &str, table_name: &str) -> Result<bool, DatabaseError> {
let key = session_key(id, table_name);
let mut conn = self.redis.connection_clone();
dispatch!(conn, |c| redis::cmd("EXISTS").arg(&key).query_async(c).await)
.map_err(|err| DatabaseError::GenericSelectError(err.to_string()))
}
async fn delete_all(&self, table_name: &str) -> Result<(), DatabaseError> {
ensure_table(table_name)?;
let keys = scan_keys(&self.redis, &format!("{table_name}:*")).await?;
if keys.is_empty() {
return Ok(());
}
let mut conn = self.redis.connection_clone();
match &mut conn {
// A single multi-key UNLINK must not cross slots, so batching is
// only safe on a standalone connection.
RedisConnection::Standalone(_) => {
for chunk in keys.chunks(UNLINK_BATCH) {
dispatch!(conn, |c| {
let mut cmd = redis::cmd("UNLINK");
for key in chunk {
cmd.arg(key);
}
cmd.exec_async(c).await
})
.map_err(|err| DatabaseError::GenericDeleteError(err.to_string()))?;
}
}
RedisConnection::Cluster(_) => {
for key in &keys {
dispatch!(conn, |c| redis::cmd("UNLINK").arg(key).exec_async(c).await)
.map_err(|err| DatabaseError::GenericDeleteError(err.to_string()))?;
}
}
}
Ok(())
}
async fn get_ids(&self, table_name: &str) -> Result<Vec<String>, DatabaseError> {
ensure_table(table_name)?;
let prefix = format!("{table_name}:");
let ids = scan_keys(&self.redis, &format!("{table_name}:*"))
.await?
.into_iter()
.filter_map(|key| key.strip_prefix(&prefix).map(str::to_string))
.collect();
Ok(ids)
}
fn auto_handles_expiry(&self) -> bool {
true
}
}
/// Session keys are managed under a `{table_name}:` prefix; an empty prefix
/// would make scans, counts and deletes target unrelated keys (or the whole
/// database). `connect` already rejects it, this is the defensive backstop.
fn ensure_table(table_name: &str) -> Result<(), DatabaseError> {
if table_name.is_empty() {
return Err(DatabaseError::GenericNotSupportedError(
"session table_name must not be empty".into(),
));
}
Ok(())
}