sa_token_core/
safe.rs

1// Author: 金书记
2//
3//! 二级认证（对齐 Java StpLogic.openSafe/checkSafe/isSafe）
4
5use std::time::Duration;
6
7use crate::error::{SaTokenError, SaTokenResult};
8use crate::manager::SaTokenManager;
9use crate::token::TokenValue;
10
11/// 默认二级认证业务标识（对齐 Java `DEFAULT_SAFE_AUTH_SERVICE`）
12pub const DEFAULT_SAFE_SERVICE: &str = "";
13
14/// 二级认证存储标记值
15pub const SAFE_AUTH_VALUE: &str = "ok";
16
17impl SaTokenManager {
18    fn safe_key(&self, token: &str, service: &str) -> String {
19        self.config.make_key("safe:", &format!("{}:{}", token, service))
20    }
21
22    /// 为指定 token 开启二级认证
23    pub async fn open_safe(
24        &self,
25        token: &TokenValue,
26        service: &str,
27        safe_time: i64,
28    ) -> SaTokenResult<()> {
29        if safe_time < 0 {
30            return Err(SaTokenError::ConfigError(
31                "safe_time must be >= 0".to_string(),
32            ));
33        }
34
35        let ttl = if safe_time == 0 {
36            None
37        } else {
38            Some(Duration::from_secs(safe_time as u64))
39        };
40
41        self.storage
42            .set(
43                &self.safe_key(token.as_str(), service),
44                SAFE_AUTH_VALUE,
45                ttl,
46            )
47            .await
48            .map_err(|e| SaTokenError::StorageError(e.to_string()))?;
49
50        self.event_bus
51            .publish(crate::event::SaTokenEvent::open_safe(token.as_str(), service))
52            .await;
53
54        Ok(())
55    }
56
57    /// 判断 token 是否已通过指定业务的二级认证
58    pub async fn is_safe(&self, token: &TokenValue, service: &str) -> SaTokenResult<bool> {
59        if token.as_str().is_empty() {
60            return Ok(false);
61        }
62
63        if !self.is_valid(token).await {
64            return Ok(false);
65        }
66
67        let value = self
68            .storage
69            .get(&self.safe_key(token.as_str(), service))
70            .await
71            .map_err(|e| SaTokenError::StorageError(e.to_string()))?;
72
73        Ok(value.is_some())
74    }
75
76    /// 校验二级认证；未通过抛出 [`SaTokenError::NotSafe`]
77    pub async fn check_safe(&self, token: &TokenValue, service: &str) -> SaTokenResult<()> {
78        if !self.is_valid(token).await {
79            return Err(SaTokenError::NotLogin);
80        }
81
82        if !self.is_safe(token, service).await? {
83            return Err(SaTokenError::NotSafe(service.to_string()));
84        }
85
86        Ok(())
87    }
88
89    /// 关闭二级认证
90    pub async fn close_safe(&self, token: &TokenValue, service: &str) -> SaTokenResult<()> {
91        self.storage
92            .delete(&self.safe_key(token.as_str(), service))
93            .await
94            .map_err(|e| SaTokenError::StorageError(e.to_string()))?;
95
96        self.event_bus
97            .publish(crate::event::SaTokenEvent::close_safe(token.as_str(), service))
98            .await;
99
100        Ok(())
101    }
102
103    /// 获取二级认证剩余有效时间（秒）；未认证返回 `None`
104    pub async fn get_safe_time(
105        &self,
106        token: &TokenValue,
107        service: &str,
108    ) -> SaTokenResult<Option<i64>> {
109        match self.storage.ttl(&self.safe_key(token.as_str(), service)).await {
110            Ok(Some(d)) => Ok(Some(d.as_secs() as i64)),
111            Ok(None) => Ok(None),
112            Err(e) => Err(SaTokenError::StorageError(e.to_string())),
113        }
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use crate::config::SaTokenConfig;
121    use sa_token_storage_memory::MemoryStorage;
122    use std::sync::Arc;
123
124    fn manager() -> SaTokenManager {
125        SaTokenManager::new(
126            Arc::new(MemoryStorage::new()),
127            SaTokenConfig::default(),
128        )
129    }
130
131    #[tokio::test]
132    async fn open_check_close_safe() {
133        let mgr = manager();
134        let token = mgr.login("u1").await.unwrap();
135        assert!(!mgr.is_safe(&token, DEFAULT_SAFE_SERVICE).await.unwrap());
136        mgr.open_safe(&token, "pay", 120).await.unwrap();
137        assert!(mgr.is_safe(&token, "pay").await.unwrap());
138        mgr.check_safe(&token, "pay").await.unwrap();
139        mgr.close_safe(&token, "pay").await.unwrap();
140        assert!(!mgr.is_safe(&token, "pay").await.unwrap());
141    }
142}
sa_token_core/safe.rs

sa_token_core/
safe.rs
