sa_token_core/token/
generator.rs

1// Author: 金书记
2//
3//! Token Generator | Token 生成器
4//!
5//! Supports multiple token styles including UUID, Random, and JWT
6//! 支持多种 Token 风格，包括 UUID、随机字符串和 JWT
7
8use uuid::Uuid;
9use crate::config::{TokenStyle, SaTokenConfig};
10use crate::token::TokenValue;
11use crate::token::jwt::{JwtManager, JwtClaims, JwtAlgorithm};
12use chrono::Utc;
13use sha2::{Sha256, Sha512, Digest};
14
15pub struct TokenGenerator;
16
17impl TokenGenerator {
18    /// Generate token based on configuration | 根据配置生成 token
19    ///
20    /// # Arguments | 参数
21    ///
22    /// * `config` - Sa-token configuration | Sa-token 配置
23    /// * `login_id` - User login ID (required for JWT) | 用户登录ID（JWT 必需）
24    pub fn generate_with_login_id(config: &SaTokenConfig, login_id: &str) -> TokenValue {
25        match config.token_style {
26            TokenStyle::Uuid => Self::generate_uuid(),
27            TokenStyle::SimpleUuid => Self::generate_simple_uuid(),
28            TokenStyle::Random32 => Self::generate_random(32),
29            TokenStyle::Random64 => Self::generate_random(64),
30            TokenStyle::Random128 => Self::generate_random(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    ///
40    /// 当 token_style 为 JWT 时，extra_data 会被签名到 JWT Claims 中。
41    /// 其他风格不支持在 token 本身携带数据，extra_data 仅存储在 storage 中。
42    ///
43    /// # Arguments | 参数
44    ///
45    /// * `config` - Sa-token configuration | Sa-token 配置
46    /// * `login_id` - User login ID | 用户登录ID
47    /// * `extra_data` - Extra data to sign into JWT | 要签名到 JWT 中的额外数据
48    pub fn generate_with_login_id_and_extra(
49        config: &SaTokenConfig,
50        login_id: &str,
51        extra_data: &serde_json::Value,
52    ) -> TokenValue {
53        match config.token_style {
54            TokenStyle::Jwt => Self::generate_jwt_with_extra(config, login_id, extra_data),
55            // 非 JWT 风格无法在 token 中携带 extra 数据，走原有生成逻辑
56            _ => Self::generate_with_login_id(config, login_id),
57        }
58    }
59    
60    /// Generate token (backward compatible) | 根据配置生成 token（向后兼容）
61    pub fn generate(config: &SaTokenConfig) -> TokenValue {
62        Self::generate_with_login_id(config, "")
63    }
64    
65    /// 生成 UUID 风格的 token
66    pub fn generate_uuid() -> TokenValue {
67        TokenValue::new(Uuid::new_v4().to_string())
68    }
69    
70    /// 生成简化的 UUID（去掉横杠）
71    pub fn generate_simple_uuid() -> TokenValue {
72        TokenValue::new(Uuid::new_v4().simple().to_string())
73    }
74    
75    /// 生成随机字符串
76    pub fn generate_random(length: usize) -> TokenValue {
77        let uuid = Uuid::new_v4();
78        let random_bytes = uuid.as_bytes();
79        let hash = Sha512::digest(random_bytes);
80        let hex_string = hex::encode(hash);
81        TokenValue::new(hex_string[..length.min(hex_string.len())].to_string())
82    }
83    
84    /// Generate JWT token | 生成 JWT token
85    ///
86    /// # Arguments | 参数
87    ///
88    /// * `config` - Sa-token configuration | Sa-token 配置
89    /// * `login_id` - User login ID | 用户登录ID
90    pub fn generate_jwt(config: &SaTokenConfig, login_id: &str) -> TokenValue {
91        // 如果 login_id 为空，则使用时间戳作为 login_id
92        let effective_login_id = if login_id.is_empty() {
93            Utc::now().timestamp_millis().to_string()
94        } else {
95            login_id.to_string()
96        };
97        
98        // Get JWT secret key | 获取 JWT 密钥
99        let secret = config.jwt_secret_key.as_ref()
100            .expect("JWT secret key is required when using JWT token style");
101        
102        // Parse algorithm | 解析算法
103        let algorithm = config.jwt_algorithm.as_ref()
104            .and_then(|alg| Self::parse_jwt_algorithm(alg))
105            .unwrap_or(JwtAlgorithm::HS256);
106        
107        // Create JWT manager | 创建 JWT 管理器
108        let mut jwt_manager = JwtManager::with_algorithm(secret, algorithm);
109        
110        if let Some(ref issuer) = config.jwt_issuer {
111            jwt_manager = jwt_manager.set_issuer(issuer);
112        }
113        
114        if let Some(ref audience) = config.jwt_audience {
115            jwt_manager = jwt_manager.set_audience(audience);
116        }
117        
118        // Create claims | 创建声明
119        let mut claims = JwtClaims::new(effective_login_id);
120        
121        // Set expiration | 设置过期时间
122        if config.timeout > 0 {
123            claims.set_expiration(config.timeout);
124        }
125        
126        // Generate JWT token | 生成 JWT token
127        match jwt_manager.generate(&claims) {
128            Ok(token) => TokenValue::new(token),
129            Err(e) => {
130                if config.jwt_fallback_on_error {
131                    tracing::warn!(error = %e, "Failed to generate JWT token, falling back to UUID");
132                    Self::generate_uuid()
133                } else {
134                    tracing::error!(error = %e, "Failed to generate JWT token and jwt_fallback_on_error=false");
135                    Self::generate_uuid()
136                }
137            }
138        }
139    }
140    
141    /// Generate JWT token with extra data signed into claims | 生成带有额外数据签名的 JWT token
142    ///
143    /// 与 `generate_jwt` 类似，但会将 `extra_data` 写入 JWT Claims 中，
144    /// 使得 extra 数据成为 token 签名的一部分。
145    ///
146    /// # Arguments | 参数
147    ///
148    /// * `config` - Sa-token configuration | Sa-token 配置
149    /// * `login_id` - User login ID | 用户登录ID
150    /// * `extra_data` - Extra data to embed in JWT claims | 要签入 JWT 声明的额外数据
151    pub fn generate_jwt_with_extra(
152        config: &SaTokenConfig,
153        login_id: &str,
154        extra_data: &serde_json::Value,
155    ) -> TokenValue {
156        let effective_login_id = if login_id.is_empty() {
157            Utc::now().timestamp_millis().to_string()
158        } else {
159            login_id.to_string()
160        };
161        
162        let secret = config.jwt_secret_key.as_ref()
163            .expect("JWT secret key is required when using JWT token style");
164        
165        let algorithm = config.jwt_algorithm.as_ref()
166            .and_then(|alg| Self::parse_jwt_algorithm(alg))
167            .unwrap_or(JwtAlgorithm::HS256);
168        
169        let mut jwt_manager = JwtManager::with_algorithm(secret, algorithm);
170        
171        if let Some(ref issuer) = config.jwt_issuer {
172            jwt_manager = jwt_manager.set_issuer(issuer);
173        }
174        
175        if let Some(ref audience) = config.jwt_audience {
176            jwt_manager = jwt_manager.set_audience(audience);
177        }
178        
179        let mut claims = JwtClaims::new(effective_login_id);
180        
181        if config.timeout > 0 {
182            claims.set_expiration(config.timeout);
183        }
184        
185        // 将 extra_data 写入 JWT claims
186        // If extra_data is an Object, flatten each key-value into claims.extra
187        // Otherwise, store the entire value under "extra" key
188        match extra_data {
189            serde_json::Value::Object(map) => {
190                for (key, value) in map {
191                    claims.add_claim(key.clone(), value.clone());
192                }
193            }
194            serde_json::Value::Null => {
195                // Null 值不写入
196            }
197            other => {
198                claims.add_claim("extra", other.clone());
199            }
200        }
201        
202        match jwt_manager.generate(&claims) {
203            Ok(token) => TokenValue::new(token),
204            Err(e) => {
205                if config.jwt_fallback_on_error {
206                    tracing::warn!(error = %e, "Failed to generate JWT token with extra, falling back to UUID");
207                    Self::generate_uuid()
208                } else {
209                    tracing::error!(error = %e, "Failed to generate JWT token with extra and jwt_fallback_on_error=false");
210                    Self::generate_uuid()
211                }
212            }
213        }
214    }
215    
216    /// Parse JWT algorithm from string | 从字符串解析 JWT 算法
217    fn parse_jwt_algorithm(alg: &str) -> Option<JwtAlgorithm> {
218        match alg.to_uppercase().as_str() {
219            "HS256" => Some(JwtAlgorithm::HS256),
220            "HS384" => Some(JwtAlgorithm::HS384),
221            "HS512" => Some(JwtAlgorithm::HS512),
222            "RS256" => Some(JwtAlgorithm::RS256),
223            "RS384" => Some(JwtAlgorithm::RS384),
224            "RS512" => Some(JwtAlgorithm::RS512),
225            "ES256" => Some(JwtAlgorithm::ES256),
226            "ES384" => Some(JwtAlgorithm::ES384),
227            _ => None,
228        }
229    }
230    
231    /// Generate Hash style token | 生成 Hash 风格 token
232    ///
233    /// Uses SHA256 hash of login_id + timestamp + random UUID
234    /// 使用 SHA256 哈希：login_id + 时间戳 + 随机 UUID
235    ///
236    /// # Arguments | 参数
237    ///
238    /// * `login_id` - User login ID | 用户登录ID
239    pub fn generate_hash(login_id: &str) -> TokenValue {
240        // 如果 login_id 为空，使用时间戳代替
241        let login_id_value = if login_id.is_empty() {
242            Utc::now().timestamp_millis().to_string()
243        } else {
244            login_id.to_string()
245        };
246        
247        let timestamp = Utc::now().timestamp_millis();
248        let uuid = Uuid::new_v4();
249        let data = format!("{}{}{}", login_id_value, timestamp, uuid);
250        
251        let mut hasher = Sha256::new();
252        hasher.update(data.as_bytes());
253        let result = hasher.finalize();
254        let hash = hex::encode(result);
255        
256        TokenValue::new(hash)
257    }
258    
259    /// Generate Timestamp style token | 生成时间戳风格 token
260    ///
261    /// Format: timestamp_milliseconds + 16-char random suffix
262    /// 格式：毫秒级时间戳 + 16位随机后缀
263    ///
264    /// Example: 1760403556789_a3b2c1d4e5f6g7h8
265    /// 示例：1760403556789_a3b2c1d4e5f6g7h8
266    pub fn generate_timestamp() -> TokenValue {
267        use chrono::Utc;
268        use sha2::{Sha256, Digest};
269        
270        let timestamp = Utc::now().timestamp_millis();
271        let uuid = Uuid::new_v4();
272        
273        // Generate random suffix | 生成随机后缀
274        let mut hasher = Sha256::new();
275        hasher.update(uuid.as_bytes());
276        let result = hasher.finalize();
277        let suffix = hex::encode(&result[..8]); // 16 characters
278        
279        TokenValue::new(format!("{}_{}", timestamp, suffix))
280    }
281    
282    /// Generate Tik style token | 生成 Tik 风格 token
283    ///
284    /// Short 8-character alphanumeric token (URL-safe)
285    /// 短小精悍的8位字母数字 token（URL安全）
286    ///
287    /// Character set: A-Z, a-z, 0-9 (62 characters)
288    /// 字符集：A-Z, a-z, 0-9（62个字符）
289    ///
290    /// Example: aB3dE9fG
291    /// 示例：aB3dE9fG
292    pub fn generate_tik() -> TokenValue {
293        use sha2::{Sha256, Digest};
294        
295        const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
296        const TOKEN_LENGTH: usize = 8;
297        
298        let uuid = Uuid::new_v4();
299        let mut hasher = Sha256::new();
300        hasher.update(uuid.as_bytes());
301        let hash = hasher.finalize();
302        
303        let mut token = String::with_capacity(TOKEN_LENGTH);
304        for i in 0..TOKEN_LENGTH {
305            let idx = (hash[i] as usize) % CHARSET.len();
306            token.push(CHARSET[idx] as char);
307        }
308        
309        TokenValue::new(token)
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316    use crate::config::{SaTokenConfig, TokenStyle};
317    use crate::token::jwt::JwtManager;
318
319    fn jwt_config() -> SaTokenConfig {
320        SaTokenConfig {
321            token_style: TokenStyle::Jwt,
322            jwt_secret_key: Some("test-secret-key-for-jwt".to_string()),
323            timeout: 3600,
324            ..SaTokenConfig::default()
325        }
326    }
327
328    #[test]
329    fn test_generate_jwt_with_extra_object() {
330        let config = jwt_config();
331        let extra = serde_json::json!({
332            "role": "admin",
333            "tenant_id": 42,
334            "permissions": ["read", "write"]
335        });
336
337        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_123", &extra);
338        assert!(!token.as_str().is_empty());
339
340        // 解析 JWT 验证 extra 数据已签入
341        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
342        let claims = jwt_manager.validate(token.as_str()).unwrap();
343
344        assert_eq!(claims.login_id, "user_123");
345        assert_eq!(claims.get_claim("role"), Some(&serde_json::json!("admin")));
346        assert_eq!(claims.get_claim("tenant_id"), Some(&serde_json::json!(42)));
347        assert_eq!(
348            claims.get_claim("permissions"),
349            Some(&serde_json::json!(["read", "write"]))
350        );
351    }
352
353    #[test]
354    fn test_generate_jwt_with_extra_non_object() {
355        let config = jwt_config();
356        let extra = serde_json::json!("simple_string_value");
357
358        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_456", &extra);
359
360        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
361        let claims = jwt_manager.validate(token.as_str()).unwrap();
362
363        assert_eq!(claims.login_id, "user_456");
364        assert_eq!(
365            claims.get_claim("extra"),
366            Some(&serde_json::json!("simple_string_value"))
367        );
368    }
369
370    #[test]
371    fn test_generate_jwt_with_extra_null() {
372        let config = jwt_config();
373        let extra = serde_json::Value::Null;
374
375        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_789", &extra);
376
377        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
378        let claims = jwt_manager.validate(token.as_str()).unwrap();
379
380        assert_eq!(claims.login_id, "user_789");
381        assert!(claims.extra.is_empty());
382    }
383
384    #[test]
385    fn test_generate_with_login_id_and_extra_jwt_style() {
386        let config = jwt_config();
387        let extra = serde_json::json!({"key": "value"});
388
389        let token = TokenGenerator::generate_with_login_id_and_extra(&config, "user_jwt", &extra);
390
391        // JWT token 包含两个 '.' 分隔符
392        assert!(token.as_str().contains('.'));
393
394        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
395        let claims = jwt_manager.validate(token.as_str()).unwrap();
396        assert_eq!(claims.get_claim("key"), Some(&serde_json::json!("value")));
397    }
398
399    #[test]
400    fn test_generate_with_login_id_and_extra_non_jwt_style() {
401        let config = SaTokenConfig {
402            token_style: TokenStyle::Uuid,
403            ..SaTokenConfig::default()
404        };
405        let extra = serde_json::json!({"key": "value"});
406
407        // 非 JWT 风格应该走正常生成逻辑，不 panic
408        let token = TokenGenerator::generate_with_login_id_and_extra(&config, "user_uuid", &extra);
409        assert!(!token.as_str().is_empty());
410        // UUID 格式不包含 '.'
411        assert!(!token.as_str().contains('.'));
412    }
413
414
415    #[test]
416    fn test_random_32_length() {
417        let config = SaTokenConfig {
418            token_style: TokenStyle::Random32,
419            ..SaTokenConfig::default()
420        };
421        let token = TokenGenerator::generate_with_login_id(&config, "user_random");
422        assert!(!token.as_str().is_empty());
423        assert_eq!(token.as_str().len(), 32);
424    }
425
426    #[test]
427    fn test_random_64_length() {
428        let config = SaTokenConfig {
429            token_style: TokenStyle::Random64,
430            ..SaTokenConfig::default()
431        };
432        let token = TokenGenerator::generate_with_login_id(&config, "user_random");
433        assert!(!token.as_str().is_empty());
434        assert_eq!(token.as_str().len(), 64);
435    }
436    
437    #[test]
438    fn test_random_128_length() {
439        let config = SaTokenConfig {
440            token_style: TokenStyle::Random128,
441            ..SaTokenConfig::default()
442        };
443        let token = TokenGenerator::generate_with_login_id(&config, "user_random");
444        assert!(!token.as_str().is_empty());
445        assert_eq!(token.as_str().len(), 128);
446    }
447}
sa_token_core/token/generator.rs

sa_token_core/token/
generator.rs
