use crate::config::SaTokenConfig;
use crate::dao::SaTokenDao;
use crate::error::SaTokenError;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use sa_token_adapter::storage::SaStorage;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributedSession {
pub session_id: String,
pub login_id: String,
pub token: String,
pub service_id: String,
pub create_time: DateTime<Utc>,
pub last_access: DateTime<Utc>,
pub attributes: HashMap<String, String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceCredential {
pub service_id: String,
pub service_name: String,
pub secret_key: String,
pub created_at: DateTime<Utc>,
pub permissions: Vec<String>,
}
#[async_trait]
pub trait DistributedSessionStorage: Send + Sync {
async fn save_session(
&self,
session: DistributedSession,
ttl: Option<Duration>,
) -> Result<(), SaTokenError>;
async fn get_session(
&self,
session_id: &str,
) -> Result<Option<DistributedSession>, SaTokenError>;
async fn delete_session(&self, session_id: &str) -> Result<(), SaTokenError>;
async fn get_sessions_by_login_id(
&self,
login_id: &str,
) -> Result<Vec<DistributedSession>, SaTokenError>;
async fn save_credential(&self, credential: ServiceCredential) -> Result<(), SaTokenError>;
async fn get_credential(
&self,
service_id: &str,
) -> Result<Option<ServiceCredential>, SaTokenError>;
}
pub struct DistributedSessionManager {
storage: Arc<dyn DistributedSessionStorage>,
service_id: String,
session_timeout: Duration,
}
impl std::fmt::Debug for DistributedSessionManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("DistributedSessionManager { .. }")
}
}
impl DistributedSessionManager {
pub fn new(
storage: Arc<dyn DistributedSessionStorage>,
service_id: String,
session_timeout: Duration,
) -> Self {
Self {
storage,
service_id,
session_timeout,
}
}
pub async fn register_service(
&self,
credential: ServiceCredential,
) -> Result<(), SaTokenError> {
self.storage.save_credential(credential).await
}
pub async fn verify_service(
&self,
service_id: &str,
secret: &str,
) -> Result<ServiceCredential, SaTokenError> {
if let Some(cred) = self.storage.get_credential(service_id).await?
&& crate::http_basic::ct_eq(cred.secret_key.as_bytes(), secret.as_bytes())
{
return Ok(cred);
}
Err(SaTokenError::PermissionDenied)
}
pub async fn create_session(
&self,
login_id: String,
token: String,
) -> Result<DistributedSession, SaTokenError> {
let session = DistributedSession {
session_id: uuid::Uuid::new_v4().to_string(),
login_id,
token,
service_id: self.service_id.clone(),
create_time: Utc::now(),
last_access: Utc::now(),
attributes: HashMap::new(),
};
self.storage
.save_session(session.clone(), Some(self.session_timeout))
.await?;
Ok(session)
}
pub async fn get_session(&self, session_id: &str) -> Result<DistributedSession, SaTokenError> {
self.storage
.get_session(session_id)
.await?
.ok_or(SaTokenError::SessionNotFound)
}
pub async fn update_session(&self, session: DistributedSession) -> Result<(), SaTokenError> {
self.storage
.save_session(session, Some(self.session_timeout))
.await
}
pub async fn delete_session(&self, session_id: &str) -> Result<(), SaTokenError> {
self.storage.delete_session(session_id).await
}
pub async fn refresh_session(&self, session_id: &str) -> Result<(), SaTokenError> {
let mut session = self.get_session(session_id).await?;
session.last_access = Utc::now();
self.update_session(session).await
}
pub async fn set_attribute(
&self,
session_id: &str,
key: String,
value: String,
) -> Result<(), SaTokenError> {
let mut session = self.get_session(session_id).await?;
session.attributes.insert(key, value);
session.last_access = Utc::now();
self.update_session(session).await
}
pub async fn get_attribute(
&self,
session_id: &str,
key: &str,
) -> Result<Option<String>, SaTokenError> {
let session = self.get_session(session_id).await?;
Ok(session.attributes.get(key).cloned())
}
pub async fn remove_attribute(&self, session_id: &str, key: &str) -> Result<(), SaTokenError> {
let mut session = self.get_session(session_id).await?;
session.attributes.remove(key);
session.last_access = Utc::now();
self.update_session(session).await
}
pub async fn get_sessions_by_login_id(
&self,
login_id: &str,
) -> Result<Vec<DistributedSession>, SaTokenError> {
self.storage.get_sessions_by_login_id(login_id).await
}
pub async fn delete_all_sessions(&self, login_id: &str) -> Result<(), SaTokenError> {
let sessions = self.get_sessions_by_login_id(login_id).await?;
for session in sessions {
self.delete_session(&session.session_id).await?;
}
Ok(())
}
pub async fn delete_sessions_by_token(
&self,
login_id: &str,
token: &str,
) -> Result<(), SaTokenError> {
let sessions = self.storage.get_sessions_by_login_id(login_id).await?;
for session in sessions {
if crate::http_basic::ct_eq(session.token.as_bytes(), token.as_bytes()) {
self.storage.delete_session(&session.session_id).await?;
}
}
Ok(())
}
}
pub struct InMemoryDistributedStorage {
sessions: Arc<RwLock<HashMap<String, DistributedSession>>>,
login_index: Arc<RwLock<HashMap<String, Vec<String>>>>,
credentials: Arc<RwLock<HashMap<String, ServiceCredential>>>,
}
impl std::fmt::Debug for InMemoryDistributedStorage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("InMemoryDistributedStorage { .. }")
}
}
impl InMemoryDistributedStorage {
pub fn new() -> Self {
Self {
sessions: Arc::new(RwLock::new(HashMap::new())),
login_index: Arc::new(RwLock::new(HashMap::new())),
credentials: Arc::new(RwLock::new(HashMap::new())),
}
}
}
impl Default for InMemoryDistributedStorage {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl DistributedSessionStorage for InMemoryDistributedStorage {
async fn save_session(
&self,
session: DistributedSession,
_ttl: Option<Duration>,
) -> Result<(), SaTokenError> {
let session_id = session.session_id.clone();
let login_id = session.login_id.clone();
let mut sessions = self.sessions.write().await;
sessions.insert(session_id.clone(), session);
let mut index = self.login_index.write().await;
let session_list = index.entry(login_id).or_insert_with(Vec::new);
if !session_list.contains(&session_id) {
session_list.push(session_id);
}
Ok(())
}
async fn get_session(
&self,
session_id: &str,
) -> Result<Option<DistributedSession>, SaTokenError> {
let sessions = self.sessions.read().await;
Ok(sessions.get(session_id).cloned())
}
async fn delete_session(&self, session_id: &str) -> Result<(), SaTokenError> {
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.remove(session_id) {
let mut index = self.login_index.write().await;
if let Some(session_ids) = index.get_mut(&session.login_id) {
session_ids.retain(|id| id != session_id);
if session_ids.is_empty() {
index.remove(&session.login_id);
}
}
}
Ok(())
}
async fn get_sessions_by_login_id(
&self,
login_id: &str,
) -> Result<Vec<DistributedSession>, SaTokenError> {
let index = self.login_index.read().await;
let session_ids = index.get(login_id).cloned().unwrap_or_default();
let sessions = self.sessions.read().await;
let mut result = Vec::new();
for session_id in session_ids {
if let Some(session) = sessions.get(&session_id) {
result.push(session.clone());
}
}
Ok(result)
}
async fn save_credential(&self, credential: ServiceCredential) -> Result<(), SaTokenError> {
let mut creds = self.credentials.write().await;
creds.insert(credential.service_id.clone(), credential);
Ok(())
}
async fn get_credential(
&self,
service_id: &str,
) -> Result<Option<ServiceCredential>, SaTokenError> {
let creds = self.credentials.read().await;
Ok(creds.get(service_id).cloned())
}
}
pub struct SaStorageDistributedStorage {
dao: Arc<SaTokenDao>,
}
impl std::fmt::Debug for SaStorageDistributedStorage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("SaStorageDistributedStorage { .. }")
}
}
impl SaStorageDistributedStorage {
pub fn from_dao(dao: Arc<SaTokenDao>) -> Self {
Self { dao }
}
pub fn from_config(storage: Arc<dyn SaStorage>, config: &SaTokenConfig) -> Self {
Self::from_dao(Arc::new(SaTokenDao::new(storage, Arc::new(config.clone()))))
}
pub fn new(storage: Arc<dyn SaStorage>, key_prefix: impl Into<String>) -> Self {
let config = SaTokenConfig {
storage_key_prefix: key_prefix.into(),
..SaTokenConfig::default()
};
Self::from_config(storage, &config)
}
}
#[async_trait]
impl DistributedSessionStorage for SaStorageDistributedStorage {
async fn save_session(
&self,
session: DistributedSession,
ttl: Option<Duration>,
) -> Result<(), SaTokenError> {
let session_key = self.dao.keys().distributed_session(&session.session_id);
let index_key = self.dao.keys().distributed_session_index(&session.login_id);
self.dao.set_object(&session_key, &session, ttl).await?;
self.dao
.list_push_unique(&index_key, &session.session_id, None)
.await?;
Ok(())
}
async fn get_session(
&self,
session_id: &str,
) -> Result<Option<DistributedSession>, SaTokenError> {
self.dao
.get_object(&self.dao.keys().distributed_session(session_id))
.await
}
async fn delete_session(&self, session_id: &str) -> Result<(), SaTokenError> {
if let Some(session) = self.get_session(session_id).await? {
self.dao
.delete(&self.dao.keys().distributed_session(session_id))
.await?;
let index_key = self.dao.keys().distributed_session_index(&session.login_id);
self.dao.list_remove(&index_key, session_id).await?;
if self.dao.list_len(&index_key).await? == 0 {
self.dao.delete(&index_key).await?;
}
} else {
self.dao
.delete(&self.dao.keys().distributed_session(session_id))
.await?;
}
Ok(())
}
async fn get_sessions_by_login_id(
&self,
login_id: &str,
) -> Result<Vec<DistributedSession>, SaTokenError> {
let index_key = self.dao.keys().distributed_session_index(login_id);
let ids = self.dao.list_range(&index_key, 0, None).await?;
let mut out = Vec::new();
for id in ids {
match self.get_session(&id).await? {
Some(s) => out.push(s),
None => {
self.dao.list_remove(&index_key, &id).await?;
}
}
}
Ok(out)
}
async fn save_credential(&self, credential: ServiceCredential) -> Result<(), SaTokenError> {
let key = self.dao.keys().distributed_service(&credential.service_id);
self.dao.set_object(&key, &credential, None).await
}
async fn get_credential(
&self,
service_id: &str,
) -> Result<Option<ServiceCredential>, SaTokenError> {
self.dao
.get_object(&self.dao.keys().distributed_service(service_id))
.await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_distributed_session_manager() {
let storage = Arc::new(InMemoryDistributedStorage::new());
let manager = DistributedSessionManager::new(
storage,
"service1".to_string(),
Duration::from_secs(3600),
);
let session = manager
.create_session("user1".to_string(), "token1".to_string())
.await
.unwrap();
let retrieved = manager.get_session(&session.session_id).await.unwrap();
assert_eq!(retrieved.login_id, "user1");
}
#[tokio::test]
async fn test_session_attributes() {
let storage = Arc::new(InMemoryDistributedStorage::new());
let manager = DistributedSessionManager::new(
storage,
"service1".to_string(),
Duration::from_secs(3600),
);
let session = manager
.create_session("user2".to_string(), "token2".to_string())
.await
.unwrap();
manager
.set_attribute(
&session.session_id,
"key1".to_string(),
"value1".to_string(),
)
.await
.unwrap();
let value = manager
.get_attribute(&session.session_id, "key1")
.await
.unwrap();
assert_eq!(value, Some("value1".to_string()));
}
#[tokio::test]
async fn test_service_verification() {
let storage = Arc::new(InMemoryDistributedStorage::new());
let manager = DistributedSessionManager::new(
storage,
"service1".to_string(),
Duration::from_secs(3600),
);
let credential = ServiceCredential {
service_id: "service2".to_string(),
service_name: "Service 2".to_string(),
secret_key: "secret123".to_string(),
created_at: Utc::now(),
permissions: vec!["read".to_string(), "write".to_string()],
};
manager.register_service(credential.clone()).await.unwrap();
let verified = manager
.verify_service("service2", "secret123")
.await
.unwrap();
assert_eq!(verified.service_id, "service2");
let result = manager.verify_service("service2", "wrong_secret").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_delete_all_sessions() {
let storage = Arc::new(InMemoryDistributedStorage::new());
let manager = DistributedSessionManager::new(
storage,
"service1".to_string(),
Duration::from_secs(3600),
);
manager
.create_session("user3".to_string(), "token1".to_string())
.await
.unwrap();
manager
.create_session("user3".to_string(), "token2".to_string())
.await
.unwrap();
let sessions = manager.get_sessions_by_login_id("user3").await.unwrap();
assert_eq!(sessions.len(), 2);
manager.delete_all_sessions("user3").await.unwrap();
let sessions = manager.get_sessions_by_login_id("user3").await.unwrap();
assert_eq!(sessions.len(), 0);
}
}