sa_token_core/token/
generator.rs

1// Author: 金书记 | Author: Jin Shuji
2//! Token Generator | Token 生成器
3//!
4//! Supports multiple token styles including UUID, Random, and JWT
5//! 支持多种 Token 风格，包括 UUID、随机字符串和 JWT
6
7use crate::config::{SaTokenConfig, TokenStyle};
8use crate::error::{SaTokenError, SaTokenResult};
9use crate::token::TokenValue;
10use crate::token::csprng::{fill_bytes, random_hex, random_tik};
11use crate::token::jwt::{JwtAlgorithm, JwtClaims, JwtManager};
12use chrono::Utc;
13use sha2::{Digest, Sha256};
14use uuid::Uuid;
15
16/// Token value generator | Token 值生成器
17pub struct TokenGenerator;
18
19impl TokenGenerator {
20    /// Generate token based on configuration | 根据配置生成 token
21    pub fn generate_with_login_id(
22        config: &SaTokenConfig,
23        login_id: &str,
24    ) -> SaTokenResult<TokenValue> {
25        match config.token_style {
26            TokenStyle::Uuid => Ok(Self::generate_uuid()),
27            TokenStyle::SimpleUuid => Ok(Self::generate_simple_uuid()),
28            TokenStyle::Random32 => Self::generate_random_csprng(32),
29            TokenStyle::Random64 => Self::generate_random_csprng(64),
30            TokenStyle::Random128 => Self::generate_random_csprng(128),
31            TokenStyle::Jwt => Self::generate_jwt(config, login_id),
32            TokenStyle::Hash => Self::generate_hash(login_id),
33            TokenStyle::Timestamp => Self::generate_timestamp(),
34            TokenStyle::Tik => Self::generate_tik(),
35        }
36    }
37
38    /// Generate token with login_id and extra data | 根据配置生成带有额外数据的 token
39    pub fn generate_with_login_id_and_extra(
40        config: &SaTokenConfig,
41        login_id: &str,
42        extra_data: &serde_json::Value,
43    ) -> SaTokenResult<TokenValue> {
44        match config.token_style {
45            TokenStyle::Jwt => Self::generate_jwt_with_extra(config, login_id, extra_data),
46            _ => Self::generate_with_login_id(config, login_id),
47        }
48    }
49
50    /// Generate token (backward compatible) | 根据配置生成 token（向后兼容）
51    pub fn generate(config: &SaTokenConfig) -> SaTokenResult<TokenValue> {
52        Self::generate_with_login_id(config, "")
53    }
54
55    /// 生成 UUID 风格的 token
56    pub fn generate_uuid() -> TokenValue {
57        TokenValue::new(Uuid::new_v4().to_string())
58    }
59
60    /// 生成简化的 UUID（去掉横杠）
61    pub fn generate_simple_uuid() -> TokenValue {
62        TokenValue::new(Uuid::new_v4().simple().to_string())
63    }
64
65    /// Hex token whose entropy is `length/2` bytes of OS CSPRNG (not a hash of a UUID).
66    /// hex token，熵来自 `length/2` 字节操作系统随机数（不是 UUID 的哈希）。
67    pub fn generate_random_csprng(length: usize) -> SaTokenResult<TokenValue> {
68        Ok(TokenValue::new(random_hex(length)?))
69    }
70
71    /// Old name kept as a wrapper so call sites can migrate in one commit.
72    /// 保留旧名作为包装，便于调用点一次改完。
73    pub fn generate_random(length: usize) -> SaTokenResult<TokenValue> {
74        Self::generate_random_csprng(length)
75    }
76
77    /// Generate JWT token | 生成 JWT token
78    pub fn generate_jwt(config: &SaTokenConfig, login_id: &str) -> SaTokenResult<TokenValue> {
79        let secret = require_jwt_secret(config)?;
80        let effective_login_id = if login_id.is_empty() {
81            Utc::now().timestamp_millis().to_string()
82        } else {
83            login_id.to_string()
84        };
85        let algorithm = config
86            .jwt_algorithm
87            .as_ref()
88            .and_then(|alg| Self::parse_jwt_algorithm(alg))
89            .unwrap_or(JwtAlgorithm::HS256);
90        let mut jwt_manager = JwtManager::with_algorithm(secret, algorithm);
91        if let Some(ref issuer) = config.jwt_issuer {
92            jwt_manager = jwt_manager.set_issuer(issuer);
93        }
94        if let Some(ref audience) = config.jwt_audience {
95            jwt_manager = jwt_manager.set_audience(audience);
96        }
97        let mut claims = JwtClaims::new(effective_login_id);
98        if config.timeout > 0 {
99            claims.set_expiration(config.timeout);
100        }
101        match jwt_manager.generate(&claims) {
102            Ok(token) => Ok(TokenValue::new(token)),
103            Err(e) if config.jwt_fallback_on_error => {
104                tracing::warn!(error = %e, "JWT generation failed, falling back to UUID");
105                Ok(Self::generate_uuid())
106            }
107            Err(e) => Err(SaTokenError::ConfigError(format!(
108                "JWT generation failed: {e}"
109            ))),
110        }
111    }
112
113    /// Generate JWT token with extra data signed into claims | 生成带有额外数据签名的 JWT token
114    pub fn generate_jwt_with_extra(
115        config: &SaTokenConfig,
116        login_id: &str,
117        extra_data: &serde_json::Value,
118    ) -> SaTokenResult<TokenValue> {
119        let secret = require_jwt_secret(config)?;
120        let effective_login_id = if login_id.is_empty() {
121            Utc::now().timestamp_millis().to_string()
122        } else {
123            login_id.to_string()
124        };
125        let algorithm = config
126            .jwt_algorithm
127            .as_ref()
128            .and_then(|alg| Self::parse_jwt_algorithm(alg))
129            .unwrap_or(JwtAlgorithm::HS256);
130        let mut jwt_manager = JwtManager::with_algorithm(secret, algorithm);
131        if let Some(ref issuer) = config.jwt_issuer {
132            jwt_manager = jwt_manager.set_issuer(issuer);
133        }
134        if let Some(ref audience) = config.jwt_audience {
135            jwt_manager = jwt_manager.set_audience(audience);
136        }
137        let mut claims = JwtClaims::new(effective_login_id);
138        if config.timeout > 0 {
139            claims.set_expiration(config.timeout);
140        }
141        match extra_data {
142            serde_json::Value::Object(map) => {
143                for (key, value) in map {
144                    claims.add_claim(key.clone(), value.clone());
145                }
146            }
147            serde_json::Value::Null => {}
148            other => {
149                claims.add_claim("extra", other.clone());
150            }
151        }
152        match jwt_manager.generate(&claims) {
153            Ok(token) => Ok(TokenValue::new(token)),
154            Err(e) if config.jwt_fallback_on_error => {
155                tracing::warn!(error = %e, "JWT generation with extra failed, falling back to UUID");
156                Ok(Self::generate_uuid())
157            }
158            Err(e) => Err(SaTokenError::ConfigError(format!(
159                "JWT generation failed: {e}"
160            ))),
161        }
162    }
163
164    /// Generate Hash style token | 生成 Hash 风格 token
165    pub fn generate_hash(login_id: &str) -> SaTokenResult<TokenValue> {
166        let login_id_value = if login_id.is_empty() {
167            Utc::now().timestamp_millis().to_string()
168        } else {
169            login_id.to_string()
170        };
171        let mut salt = [0u8; 16];
172        fill_bytes(&mut salt)?;
173        let data = format!(
174            "{}{}{}",
175            login_id_value,
176            Utc::now().timestamp_millis(),
177            hex::encode(salt)
178        );
179        let mut hasher = Sha256::new();
180        hasher.update(data.as_bytes());
181        Ok(TokenValue::new(hex::encode(hasher.finalize())))
182    }
183
184    /// Generate Timestamp style token | 生成时间戳风格 token
185    pub fn generate_timestamp() -> SaTokenResult<TokenValue> {
186        let timestamp = Utc::now().timestamp_millis();
187        let suffix = random_hex(16)?;
188        Ok(TokenValue::new(format!("{timestamp}_{suffix}")))
189    }
190
191    /// Generate Tik style token | 生成 Tik 风格 token
192    pub fn generate_tik() -> SaTokenResult<TokenValue> {
193        Ok(TokenValue::new(random_tik(8)?))
194    }
195
196    fn parse_jwt_algorithm(alg: &str) -> Option<JwtAlgorithm> {
197        match alg.to_uppercase().as_str() {
198            "HS256" => Some(JwtAlgorithm::HS256),
199            "HS384" => Some(JwtAlgorithm::HS384),
200            "HS512" => Some(JwtAlgorithm::HS512),
201            "RS256" => Some(JwtAlgorithm::RS256),
202            "RS384" => Some(JwtAlgorithm::RS384),
203            "RS512" => Some(JwtAlgorithm::RS512),
204            "ES256" => Some(JwtAlgorithm::ES256),
205            "ES384" => Some(JwtAlgorithm::ES384),
206            _ => None,
207        }
208    }
209}
210
211impl std::fmt::Debug for TokenGenerator {
212    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
213        f.write_str("TokenGenerator { .. }")
214    }
215}
216
217fn require_jwt_secret(config: &SaTokenConfig) -> SaTokenResult<&str> {
218    match config.jwt_secret_key.as_deref() {
219        Some(s) if !s.trim().is_empty() => Ok(s),
220        _ => Err(SaTokenError::ConfigError(
221            "jwt_secret_key is required when token_style=Jwt".into(),
222        )),
223    }
224}
225
226/// Generate until `occupied` is false, or until `max_try_times` is exhausted.
227/// `max_try_times < 0`：create once and return (no uniqueness probe).
228/// `max_try_times == 0`：treated as `-1`.
229///
230/// 直到 `occupied` 为 false 或次数用尽。
231/// `max_try_times < 0`：只生成一次、不做占用探测。
232/// `max_try_times == 0`：与 `-1` 相同。
233pub async fn generate_unique<C, F, Fut>(
234    max_try_times: i32,
235    mut create: C,
236    mut occupied: F,
237) -> SaTokenResult<TokenValue>
238where
239    C: FnMut() -> SaTokenResult<TokenValue>,
240    F: FnMut(&str) -> Fut,
241    Fut: Future<Output = SaTokenResult<bool>>,
242{
243    if max_try_times <= 0 {
244        return create();
245    }
246    let mut last = create()?;
247    for _ in 0..max_try_times {
248        if !occupied(last.as_str()).await? {
249            return Ok(last);
250        }
251        last = create()?;
252    }
253    if !occupied(last.as_str()).await? {
254        return Ok(last);
255    }
256    Err(SaTokenError::ConfigError(format!(
257        "failed to generate a unique token after {max_try_times} attempts"
258    )))
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264    use crate::config::{SaTokenConfig, TokenStyle};
265    use crate::token::jwt::JwtManager;
266
267    fn jwt_config() -> SaTokenConfig {
268        SaTokenConfig {
269            token_style: TokenStyle::Jwt,
270            jwt_secret_key: Some("test-secret-key-for-jwt".to_string()),
271            timeout: 3600,
272            ..SaTokenConfig::default()
273        }
274    }
275
276    #[test]
277    fn test_generate_jwt_with_extra_object() {
278        let config = jwt_config();
279        let extra = serde_json::json!({
280            "role": "admin",
281            "tenant_id": 42,
282            "permissions": ["read", "write"]
283        });
284
285        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_123", &extra).unwrap();
286        assert!(!token.as_str().is_empty());
287
288        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
289        let claims = jwt_manager.validate(token.as_str()).unwrap();
290
291        assert_eq!(claims.login_id, "user_123");
292        assert_eq!(claims.get_claim("role"), Some(&serde_json::json!("admin")));
293        assert_eq!(claims.get_claim("tenant_id"), Some(&serde_json::json!(42)));
294        assert_eq!(
295            claims.get_claim("permissions"),
296            Some(&serde_json::json!(["read", "write"]))
297        );
298    }
299
300    #[test]
301    fn test_generate_jwt_with_extra_non_object() {
302        let config = jwt_config();
303        let extra = serde_json::json!("simple_string_value");
304
305        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_456", &extra).unwrap();
306
307        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
308        let claims = jwt_manager.validate(token.as_str()).unwrap();
309
310        assert_eq!(claims.login_id, "user_456");
311        assert_eq!(
312            claims.get_claim("extra"),
313            Some(&serde_json::json!("simple_string_value"))
314        );
315    }
316
317    #[test]
318    fn test_generate_jwt_with_extra_null() {
319        let config = jwt_config();
320        let extra = serde_json::Value::Null;
321
322        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_789", &extra).unwrap();
323
324        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
325        let claims = jwt_manager.validate(token.as_str()).unwrap();
326
327        assert_eq!(claims.login_id, "user_789");
328        assert!(claims.extra.is_empty());
329    }
330
331    #[test]
332    fn test_generate_with_login_id_and_extra_jwt_style() {
333        let config = jwt_config();
334        let extra = serde_json::json!({"key": "value"});
335
336        let token =
337            TokenGenerator::generate_with_login_id_and_extra(&config, "user_jwt", &extra).unwrap();
338
339        assert!(token.as_str().contains('.'));
340
341        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
342        let claims = jwt_manager.validate(token.as_str()).unwrap();
343        assert_eq!(claims.get_claim("key"), Some(&serde_json::json!("value")));
344    }
345
346    #[test]
347    fn test_generate_with_login_id_and_extra_non_jwt_style() {
348        let config = SaTokenConfig {
349            token_style: TokenStyle::Uuid,
350            ..SaTokenConfig::default()
351        };
352        let extra = serde_json::json!({"key": "value"});
353
354        let token =
355            TokenGenerator::generate_with_login_id_and_extra(&config, "user_uuid", &extra).unwrap();
356        assert!(!token.as_str().is_empty());
357        assert!(!token.as_str().contains('.'));
358    }
359
360    #[test]
361    fn test_random_32_length() {
362        let config = SaTokenConfig {
363            token_style: TokenStyle::Random32,
364            ..SaTokenConfig::default()
365        };
366        let token = TokenGenerator::generate_with_login_id(&config, "user_random").unwrap();
367        assert!(!token.as_str().is_empty());
368        assert_eq!(token.as_str().len(), 32);
369    }
370
371    #[test]
372    fn test_random_64_length() {
373        let config = SaTokenConfig {
374            token_style: TokenStyle::Random64,
375            ..SaTokenConfig::default()
376        };
377        let token = TokenGenerator::generate_with_login_id(&config, "user_random").unwrap();
378        assert!(!token.as_str().is_empty());
379        assert_eq!(token.as_str().len(), 64);
380    }
381
382    #[test]
383    fn test_random_128_length() {
384        let config = SaTokenConfig {
385            token_style: TokenStyle::Random128,
386            ..SaTokenConfig::default()
387        };
388        let token = TokenGenerator::generate_with_login_id(&config, "user_random").unwrap();
389        assert!(!token.as_str().is_empty());
390        assert_eq!(token.as_str().len(), 128);
391    }
392}
sa_token_core/token/generator.rs

sa_token_core/token/
generator.rs
