sa_token_core/
stp_logic.rs

1// Author: 金书记
2//
3//! 多账号体系：SaLogic（对齐 Java StpLogic 实例）+ 全局注册表。
4//!
5//! 每个 SaLogic 绑定一个 login_type，对同一底层 SaTokenManager 做账号命名空间隔离。
6
7use std::collections::HashMap;
8use std::sync::{Arc, OnceLock, RwLock};
9
10use crate::disable;
11use crate::error::SaTokenResult;
12use crate::manager::SaTokenManager;
13use crate::session::SaTerminalInfo;
14use crate::token::TokenValue;
15
16/// 绑定某一 login_type 的账号逻辑门面
17#[derive(Clone)]
18pub struct SaLogic {
19    login_type: String,
20    manager: Arc<SaTokenManager>,
21}
22
23impl SaLogic {
24    pub fn new(login_type: impl Into<String>, manager: Arc<SaTokenManager>) -> Self {
25        Self {
26            login_type: login_type.into(),
27            manager,
28        }
29    }
30
31    pub fn login_type(&self) -> &str {
32        &self.login_type
33    }
34
35    pub fn manager(&self) -> &Arc<SaTokenManager> {
36        &self.manager
37    }
38
39    pub async fn login(&self, login_id: impl Into<String>) -> SaTokenResult<TokenValue> {
40        self.manager
41            .login_with_options(
42                login_id,
43                Some(self.login_type.clone()),
44                None,
45                None,
46                None,
47                None,
48            )
49            .await
50    }
51
52    pub async fn login_with_device(
53        &self,
54        login_id: impl Into<String>,
55        device: Option<String>,
56        extra: Option<serde_json::Value>,
57    ) -> SaTokenResult<TokenValue> {
58        self.manager
59            .login_with_options(
60                login_id,
61                Some(self.login_type.clone()),
62                device,
63                extra,
64                None,
65                None,
66            )
67            .await
68    }
69
70    pub async fn logout(&self, token: &TokenValue) -> SaTokenResult<()> {
71        self.manager.logout(token).await
72    }
73
74    pub async fn logout_by_login_id(&self, login_id: &str) -> SaTokenResult<()> {
75        let ns = self.manager.account_ns(&self.login_type, login_id);
76        let session = self.manager.get_session(&ns).await?;
77        let tokens = session.get_token_value_list_by_device_type(None);
78        if tokens.is_empty() {
79            return Ok(());
80        }
81        for t in tokens {
82            let _ = self.manager.logout(&TokenValue::new(t)).await;
83        }
84        Ok(())
85    }
86
87    pub async fn kick_out(&self, login_id: &str) -> SaTokenResult<()> {
88        let ns = self.manager.account_ns(&self.login_type, login_id);
89        let session = self.manager.get_session(&ns).await?;
90        let tokens = session.get_token_value_list_by_device_type(None);
91        for t in tokens {
92            self.manager.kick_out_by_token(&TokenValue::new(t)).await?;
93        }
94        self.manager.delete_session(&ns).await?;
95        Ok(())
96    }
97
98    pub async fn get_login_id(&self, token: &TokenValue) -> SaTokenResult<String> {
99        Ok(self.manager.get_token_info(token).await?.login_id)
100    }
101
102    pub async fn is_valid(&self, token: &TokenValue) -> bool {
103        self.manager.is_valid(token).await
104    }
105
106    pub async fn get_session(
107        &self,
108        login_id: &str,
109    ) -> SaTokenResult<crate::session::SaSession> {
110        let ns = self.manager.account_ns(&self.login_type, login_id);
111        self.manager.get_session(&ns).await
112    }
113
114    pub async fn get_terminal_list(
115        &self,
116        login_id: &str,
117        device_type: Option<&str>,
118    ) -> SaTokenResult<Vec<SaTerminalInfo>> {
119        self.manager
120            .get_terminal_list(&self.login_type, login_id, device_type)
121            .await
122    }
123
124    pub async fn get_terminal_info_by_token(
125        &self,
126        token: &TokenValue,
127    ) -> SaTokenResult<Option<SaTerminalInfo>> {
128        self.manager.get_terminal_info_by_token(token).await
129    }
130
131    pub async fn get_permissions(&self, login_id: &str) -> SaTokenResult<Vec<String>> {
132        self.manager
133            .get_permissions_with_type(&self.login_type, login_id)
134            .await
135    }
136
137    pub async fn set_permissions(
138        &self,
139        login_id: &str,
140        perms: Vec<String>,
141    ) -> SaTokenResult<()> {
142        self.manager
143            .set_permissions_with_type(&self.login_type, login_id, perms)
144            .await
145    }
146
147    pub async fn get_roles(&self, login_id: &str) -> SaTokenResult<Vec<String>> {
148        self.manager
149            .get_roles_with_type(&self.login_type, login_id)
150            .await
151    }
152
153    pub async fn set_roles(&self, login_id: &str, roles: Vec<String>) -> SaTokenResult<()> {
154        self.manager
155            .set_roles_with_type(&self.login_type, login_id, roles)
156            .await
157    }
158
159    pub async fn disable(&self, login_id: &str, time: i64) -> SaTokenResult<()> {
160        let ns = self.manager.account_ns(&self.login_type, login_id);
161        self.manager.disable(&ns, time).await
162    }
163
164    pub async fn check_disable(&self, login_id: &str) -> SaTokenResult<()> {
165        let ns = self.manager.account_ns(&self.login_type, login_id);
166        self.manager
167            .check_disable_level(
168                &ns,
169                disable::DEFAULT_DISABLE_SERVICE,
170                disable::MIN_DISABLE_LEVEL,
171            )
172            .await
173    }
174}
175
176static STP_LOGIC_MAP: OnceLock<RwLock<HashMap<String, Arc<SaLogic>>>> = OnceLock::new();
177
178fn registry() -> &'static RwLock<HashMap<String, Arc<SaLogic>>> {
179    STP_LOGIC_MAP.get_or_init(|| RwLock::new(HashMap::new()))
180}
181
182pub fn put_stp_logic(logic: Arc<SaLogic>) {
183    registry()
184        .write()
185        .unwrap()
186        .insert(logic.login_type().to_string(), logic);
187}
188
189pub fn remove_stp_logic(login_type: &str) {
190    registry().write().unwrap().remove(login_type);
191}
192
193pub fn get_or_create_stp_logic(login_type: &str, manager: Arc<SaTokenManager>) -> Arc<SaLogic> {
194    if let Some(found) = registry().read().unwrap().get(login_type).cloned() {
195        return found;
196    }
197    let mut map = registry().write().unwrap();
198    if let Some(found) = map.get(login_type).cloned() {
199        return found;
200    }
201    let logic = Arc::new(SaLogic::new(login_type, manager));
202    map.insert(login_type.to_string(), logic.clone());
203    logic
204}
205
206pub fn try_get_stp_logic(login_type: &str) -> Option<Arc<SaLogic>> {
207    registry().read().unwrap().get(login_type).cloned()
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213    use sa_token_storage_memory::MemoryStorage;
214
215    fn make_manager() -> Arc<SaTokenManager> {
216        Arc::new(SaTokenManager::new(
217            Arc::new(MemoryStorage::new()),
218            crate::SaTokenConfig::default(),
219        ))
220    }
221
222    #[tokio::test]
223    async fn test_sa_logic_permission_isolation() {
224        let mgr = make_manager();
225        let admin = SaLogic::new("admin", mgr.clone());
226        let user = SaLogic::new("user", mgr);
227
228        admin
229            .set_permissions("10001", vec!["admin:read".to_string()])
230            .await
231            .unwrap();
232        user.set_permissions("10001", vec!["user:read".to_string()])
233            .await
234            .unwrap();
235
236        assert_eq!(
237            admin.get_permissions("10001").await.unwrap(),
238            vec!["admin:read".to_string()]
239        );
240        assert_eq!(
241            user.get_permissions("10001").await.unwrap(),
242            vec!["user:read".to_string()]
243        );
244    }
245
246    #[tokio::test]
247    async fn test_registry_put_and_remove() {
248        let mgr = make_manager();
249        let logic = Arc::new(SaLogic::new("custom", mgr));
250        put_stp_logic(logic.clone());
251        assert!(try_get_stp_logic("custom").is_some());
252        remove_stp_logic("custom");
253        assert!(try_get_stp_logic("custom").is_none());
254    }
255
256    #[tokio::test]
257    async fn test_terminal_isolation_between_login_types() {
258        let mgr = make_manager();
259        let admin = SaLogic::new("admin", mgr.clone());
260        let user = SaLogic::new("user", mgr);
261
262        admin.login("10001").await.unwrap();
263        user.login("10001").await.unwrap();
264
265        assert_eq!(admin.get_terminal_list("10001", None).await.unwrap().len(), 1);
266        assert_eq!(user.get_terminal_list("10001", None).await.unwrap().len(), 1);
267    }
268}
sa_token_core/stp_logic.rs

sa_token_core/
stp_logic.rs
