use std::collections::HashMap;
use std::sync::RwLock;
use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use crate::error::Result;
fn unix_now() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
#[async_trait]
pub trait RefreshStore: Send + Sync + 'static {
async fn save(&self, jti: &str, user_id: &str, ttl_secs: u64) -> Result<()>;
async fn take(&self, jti: &str) -> Result<Option<String>>;
async fn revoke(&self, jti: &str) -> Result<()>;
async fn revoke_all(&self, user_id: &str) -> Result<()>;
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct PendingChallenge {
pub user_id: String,
pub expires_at: u64,
}
#[async_trait]
pub trait ChallengeStore: Send + Sync + 'static {
async fn put(&self, challenge_id: &str, user_id: &str, ttl_secs: u64) -> Result<()>;
async fn take(&self, challenge_id: &str) -> Result<Option<PendingChallenge>>;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionRecord {
pub session_id: String,
pub user_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub device: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ip: Option<String>,
pub created_at: u64,
pub last_seen: u64,
}
#[async_trait]
pub trait SessionIndex: Send + Sync + 'static {
async fn register(&self, record: SessionRecord) -> Result<()>;
async fn list(&self, user_id: &str) -> Result<Vec<SessionRecord>>;
async fn revoke(&self, session_id: &str) -> Result<()>;
async fn revoke_all(&self, user_id: &str) -> Result<()>;
}
#[derive(Default)]
pub struct MemoryRefreshStore {
inner: RwLock<HashMap<String, (String, u64)>>,
}
impl MemoryRefreshStore {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl RefreshStore for MemoryRefreshStore {
async fn save(&self, jti: &str, user_id: &str, ttl_secs: u64) -> Result<()> {
self.inner
.write()
.unwrap()
.insert(jti.to_owned(), (user_id.to_owned(), unix_now() + ttl_secs));
Ok(())
}
async fn take(&self, jti: &str) -> Result<Option<String>> {
let removed = self.inner.write().unwrap().remove(jti);
Ok(removed.and_then(|(uid, exp)| (unix_now() < exp).then_some(uid)))
}
async fn revoke(&self, jti: &str) -> Result<()> {
self.inner.write().unwrap().remove(jti);
Ok(())
}
async fn revoke_all(&self, user_id: &str) -> Result<()> {
self.inner
.write()
.unwrap()
.retain(|_, (uid, _)| uid != user_id);
Ok(())
}
}
#[derive(Default)]
pub struct MemoryChallengeStore {
inner: RwLock<HashMap<String, PendingChallenge>>,
}
impl MemoryChallengeStore {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl ChallengeStore for MemoryChallengeStore {
async fn put(&self, challenge_id: &str, user_id: &str, ttl_secs: u64) -> Result<()> {
self.inner.write().unwrap().insert(
challenge_id.to_owned(),
PendingChallenge {
user_id: user_id.to_owned(),
expires_at: unix_now() + ttl_secs,
},
);
Ok(())
}
async fn take(&self, challenge_id: &str) -> Result<Option<PendingChallenge>> {
let taken = self.inner.write().unwrap().remove(challenge_id);
Ok(taken.filter(|c| unix_now() < c.expires_at))
}
}
#[derive(Default)]
pub struct MemorySessionIndex {
inner: RwLock<HashMap<String, SessionRecord>>,
}
impl MemorySessionIndex {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl SessionIndex for MemorySessionIndex {
async fn register(&self, record: SessionRecord) -> Result<()> {
self.inner
.write()
.unwrap()
.insert(record.session_id.clone(), record);
Ok(())
}
async fn list(&self, user_id: &str) -> Result<Vec<SessionRecord>> {
Ok(self
.inner
.read()
.unwrap()
.values()
.filter(|r| r.user_id == user_id)
.cloned()
.collect())
}
async fn revoke(&self, session_id: &str) -> Result<()> {
self.inner.write().unwrap().remove(session_id);
Ok(())
}
async fn revoke_all(&self, user_id: &str) -> Result<()> {
self.inner
.write()
.unwrap()
.retain(|_, r| r.user_id != user_id);
Ok(())
}
}