use crate::RTokenError;
use crate::models::RTokenInfo;
use chrono::Utc;
use std::{
collections::HashMap,
sync::{
Arc, Mutex,
atomic::{AtomicU64, Ordering},
},
};
fn now_ms() -> u64 {
u64::try_from(Utc::now().timestamp_millis()).unwrap_or(0)
}
fn add_ttl_ms(now_ms: u64, ttl_seconds: u64) -> u64 {
let ttl_ms = (ttl_seconds as u128).saturating_mul(1000);
(now_ms as u128)
.saturating_add(ttl_ms)
.min(u64::MAX as u128) as u64
}
const PRUNE_INTERVAL_MS: u64 = 60_000;
const PRUNE_MIN_SIZE: usize = 1024;
#[derive(Clone)]
pub struct RTokenManager {
store: Arc<Mutex<HashMap<String, RTokenInfo>>>,
last_prune_ms: Arc<AtomicU64>,
}
impl RTokenManager {
pub fn new() -> Self {
Self {
store: Arc::new(Mutex::new(HashMap::new())),
last_prune_ms: Arc::new(AtomicU64::new(0)),
}
}
pub fn login(&self, id: &str, expire_time: u64) -> Result<String, RTokenError> {
let token = uuid::Uuid::new_v4().to_string();
let expire_time = add_ttl_ms(now_ms(), expire_time);
let info = RTokenInfo {
user_id: id.to_string(),
expire_at: expire_time,
roles: Vec::new(),
};
let now_ms = now_ms();
let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
store.insert(token.clone(), info);
self.maybe_prune(&mut store, now_ms);
Ok(token)
}
#[cfg(feature = "rbac")]
pub fn login_with_roles(
&self,
id: &str,
expire_time: u64,
role: impl Into<Vec<String>>,
) -> Result<String, RTokenError> {
let token = uuid::Uuid::new_v4().to_string();
let expire_time = add_ttl_ms(now_ms(), expire_time);
let info = RTokenInfo {
user_id: id.to_string(),
expire_at: expire_time,
roles: role.into(),
};
let now_ms = now_ms();
let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
store.insert(token.clone(), info);
self.maybe_prune(&mut store, now_ms);
Ok(token)
}
#[cfg(feature = "rbac")]
pub fn set_roles(&self, token: &str, roles: impl Into<Vec<String>>) -> Result<(), RTokenError> {
let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
if let Some(info) = store.get_mut(token) {
info.roles = roles.into();
}
Ok(())
}
#[cfg(feature = "rbac")]
pub fn get_roles(&self, token: &str) -> Result<Option<Vec<String>>, RTokenError> {
let store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
Ok(store.get(token).map(|info| info.roles.clone()))
}
pub fn logout(&self, token: &str) -> Result<(), RTokenError> {
self.store
.lock()
.map_err(|_| RTokenError::MutexPoisoned)?
.remove(token);
Ok(())
}
pub fn expires_at(&self, token: &str) -> Result<Option<u64>, RTokenError> {
let store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
Ok(store.get(token).map(|info| info.expire_at))
}
pub fn ttl_seconds(&self, token: &str) -> Result<Option<i64>, RTokenError> {
let now_ms = now_ms();
let store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
let Some(expire_at) = store.get(token).map(|info| info.expire_at) else {
return Ok(None);
};
if expire_at <= now_ms {
return Ok(Some(0));
}
let remaining_ms = expire_at - now_ms;
let remaining_seconds = remaining_ms.div_ceil(1000) as i64;
Ok(Some(remaining_seconds))
}
pub fn renew(&self, token: &str, ttl_seconds: u64) -> Result<bool, RTokenError> {
let expire_at = add_ttl_ms(now_ms(), ttl_seconds);
let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
let Some(info) = store.get_mut(token) else {
return Ok(false);
};
if info.expire_at < now_ms() {
store.remove(token);
return Ok(false);
}
info.expire_at = expire_at;
Ok(true)
}
pub fn rotate(&self, token: &str, ttl_seconds: u64) -> Result<Option<String>, RTokenError> {
let expire_at = add_ttl_ms(now_ms(), ttl_seconds);
let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
let Some(info) = store.get(token).cloned() else {
return Ok(None);
};
if info.expire_at < now_ms() {
store.remove(token);
return Ok(None);
}
let new_token = uuid::Uuid::new_v4().to_string();
let new_info = RTokenInfo {
user_id: info.user_id,
expire_at,
roles: info.roles,
};
store.remove(token);
store.insert(new_token.clone(), new_info);
Ok(Some(new_token))
}
pub fn prune_expired(&self) -> Result<usize, RTokenError> {
let now = now_ms();
let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
let original_len = store.len();
store.retain(|_token, info| info.expire_at >= now);
Ok(original_len - store.len())
}
pub fn validate(&self, token: &str) -> Result<Option<String>, RTokenError> {
#[cfg(feature = "rbac")]
{
Ok(self
.validate_with_roles(token)?
.map(|(user_id, _roles)| user_id))
}
#[cfg(not(feature = "rbac"))]
{
let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
let Some(info) = store.get(token) else {
return Ok(None);
};
if info.expire_at < now_ms() {
store.remove(token);
return Ok(None);
}
Ok(Some(info.user_id.clone()))
}
}
#[cfg(feature = "rbac")]
pub fn validate_with_roles(
&self,
token: &str,
) -> Result<Option<(String, Vec<String>)>, RTokenError> {
let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
let Some(info) = store.get(token) else {
return Ok(None);
};
if info.expire_at < now_ms() {
store.remove(token);
return Ok(None);
}
Ok(Some((info.user_id.clone(), info.roles.clone())))
}
fn maybe_prune(&self, store: &mut HashMap<String, RTokenInfo>, now_ms: u64) {
if store.len() < PRUNE_MIN_SIZE {
return;
}
let last = self.last_prune_ms.load(Ordering::Relaxed);
if now_ms.saturating_sub(last) < PRUNE_INTERVAL_MS {
return;
}
self.last_prune_ms.store(now_ms, Ordering::Relaxed);
store.retain(|_token, info| info.expire_at >= now_ms);
}
}
impl Default for RTokenManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(any(feature = "actix", feature = "axum"))]
#[derive(Debug)]
pub struct RUser {
pub id: String,
pub token: String,
#[cfg(feature = "rbac")]
pub roles: Vec<String>,
}
#[cfg(feature = "rbac")]
impl RUser {
pub fn has_role(&self, role: &str) -> bool {
self.roles.iter().any(|r| r == role)
}
}
#[cfg(feature = "actix")]
impl actix_web::FromRequest for RUser {
type Error = actix_web::Error;
type Future = std::pin::Pin<Box<dyn std::future::Future<Output = Result<Self, Self::Error>>>>;
fn from_request(
req: &actix_web::HttpRequest,
_payload: &mut actix_web::dev::Payload,
) -> Self::Future {
use actix_web::web;
let manager = match req.app_data::<web::Data<RTokenManager>>() {
Some(m) => m.clone(),
None => {
return Box::pin(async {
Err(actix_web::error::ErrorInternalServerError(
"Token manager not found",
))
});
}
};
let token = match crate::extract_token_from_request(req) {
Some(token) => token,
None => {
return Box::pin(async {
Err(actix_web::error::ErrorUnauthorized("Unauthorized"))
});
}
};
Box::pin(async move {
#[cfg(feature = "rbac")]
{
let token_for_check = token.clone();
let manager = manager.clone();
let user_info = actix_web::rt::task::spawn_blocking(move || {
manager.validate_with_roles(&token_for_check)
})
.await
.map_err(|_| actix_web::error::ErrorInternalServerError("Mutex poisoned"))?
.map_err(|_| actix_web::error::ErrorInternalServerError("Mutex poisoned"))?;
if let Some((user_id, roles)) = user_info {
return Ok(RUser {
id: user_id,
token,
roles,
});
}
Err(actix_web::error::ErrorUnauthorized("Invalid token"))
}
#[cfg(not(feature = "rbac"))]
{
let token_for_check = token.clone();
let manager = manager.clone();
let user_id =
actix_web::rt::task::spawn_blocking(move || manager.validate(&token_for_check))
.await
.map_err(|_| actix_web::error::ErrorInternalServerError("Mutex poisoned"))?
.map_err(|_| {
actix_web::error::ErrorInternalServerError("Mutex poisoned")
})?;
if let Some(user_id) = user_id {
return Ok(RUser { id: user_id, token });
}
Err(actix_web::error::ErrorUnauthorized("Invalid token"))
}
})
}
}