use std::sync::Arc;
use std::time::Duration;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use crate::dao::SaTokenDao;
use crate::error::{SaTokenError, SaTokenResult};
use crate::token::random_hex;
use crate::util::StpUtil;
pub const DEFAULT_NAMESPACE: &str = "default";
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TempTokenRecord {
pub value: serde_json::Value,
pub namespace: String,
pub expire_at: Option<DateTime<Utc>>,
}
#[derive(Clone)]
pub struct TempTokenManager {
dao: Arc<SaTokenDao>,
}
impl std::fmt::Debug for TempTokenManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("TempTokenManager { .. }")
}
}
impl TempTokenManager {
pub fn new(dao: Arc<SaTokenDao>) -> Self {
Self { dao }
}
fn ttl(timeout_secs: i64) -> SaTokenResult<Option<Duration>> {
if timeout_secs == 0 {
return Err(SaTokenError::ConfigError(
"temp token timeout must not be 0".into(),
));
}
if timeout_secs < 0 {
Ok(None)
} else {
Ok(Some(Duration::from_secs(timeout_secs as u64)))
}
}
fn expire_at(timeout_secs: i64) -> Option<DateTime<Utc>> {
if timeout_secs < 0 {
None
} else {
Some(Utc::now() + chrono::Duration::seconds(timeout_secs))
}
}
fn index_digest(value: &str) -> String {
let mut h = Sha256::new();
h.update(value.as_bytes());
hex::encode(h.finalize())
}
pub async fn create(
&self,
namespace: &str,
value: serde_json::Value,
timeout_secs: i64,
record_index: bool,
) -> SaTokenResult<String> {
if namespace.is_empty() {
return Err(SaTokenError::ConfigError(
"temp token namespace must not be empty".into(),
));
}
let ttl = Self::ttl(timeout_secs)?;
let record = TempTokenRecord {
value: value.clone(),
namespace: namespace.to_string(),
expire_at: Self::expire_at(timeout_secs),
};
let mut token = String::new();
for _ in 0..12 {
let candidate = random_hex(32)?;
let key = self.dao.keys().temp_token(namespace, &candidate);
let raw = self.dao.encode(&record)?;
if self.dao.set_if_absent(&key, &raw, ttl).await? {
token = candidate;
break;
}
}
if token.is_empty() {
return Err(SaTokenError::ConfigError(
"failed to allocate a unique temp token".into(),
));
}
if record_index {
if let Some(s) = value.as_str() {
let ik = self
.dao
.keys()
.temp_index(namespace, &Self::index_digest(s));
self.dao.set_string(&ik, &token, ttl).await?;
}
}
Ok(token)
}
pub async fn parse(&self, namespace: &str, token: &str) -> SaTokenResult<TempTokenRecord> {
if token.is_empty() {
return Err(SaTokenError::TempTokenNotFound);
}
let key = self.dao.keys().temp_token(namespace, token);
let rec: TempTokenRecord = self
.dao
.get_object(&key)
.await?
.ok_or(SaTokenError::TempTokenNotFound)?;
if let Some(exp) = rec.expire_at {
if Utc::now() > exp {
let _ = self.dao.delete(&key).await;
return Err(SaTokenError::TempTokenExpired);
}
}
Ok(rec)
}
pub async fn find_token(&self, namespace: &str, value: &str) -> SaTokenResult<String> {
let ik = self
.dao
.keys()
.temp_index(namespace, &Self::index_digest(value));
self.dao
.get_string(&ik)
.await?
.ok_or(SaTokenError::TempTokenNotFound)
}
pub async fn delete(&self, namespace: &str, token: &str) -> SaTokenResult<()> {
let key = self.dao.keys().temp_token(namespace, token);
if let Ok(Some(rec)) = self.dao.get_object::<TempTokenRecord>(&key).await {
if let Some(s) = rec.value.as_str() {
let ik = self
.dao
.keys()
.temp_index(namespace, &Self::index_digest(s));
let _ = self.dao.delete(&ik).await;
}
}
self.dao.delete(&key).await
}
}
pub async fn create_default(value: impl Into<String>, timeout_secs: i64) -> SaTokenResult<String> {
let manager = StpUtil::try_get_manager()?;
TempTokenManager::new(manager.dao().clone())
.create(
DEFAULT_NAMESPACE,
serde_json::Value::String(value.into()),
timeout_secs,
false,
)
.await
}
pub async fn parse_default(token: &str) -> SaTokenResult<TempTokenRecord> {
let manager = StpUtil::try_get_manager()?;
TempTokenManager::new(manager.dao().clone())
.parse(DEFAULT_NAMESPACE, token)
.await
}
pub async fn delete_default(token: &str) -> SaTokenResult<()> {
let manager = StpUtil::try_get_manager()?;
TempTokenManager::new(manager.dao().clone())
.delete(DEFAULT_NAMESPACE, token)
.await
}