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                eprintln!("Failed to generate JWT token: {:?}", e);
131                // Fallback to UUID | 回退到 UUID
132                Self::generate_uuid()
133            }
134        }
135    }
136    
137    /// Generate JWT token with extra data signed into claims | 生成带有额外数据签名的 JWT token
138    ///
139    /// 与 `generate_jwt` 类似，但会将 `extra_data` 写入 JWT Claims 中，
140    /// 使得 extra 数据成为 token 签名的一部分。
141    ///
142    /// # Arguments | 参数
143    ///
144    /// * `config` - Sa-token configuration | Sa-token 配置
145    /// * `login_id` - User login ID | 用户登录ID
146    /// * `extra_data` - Extra data to embed in JWT claims | 要签入 JWT 声明的额外数据
147    pub fn generate_jwt_with_extra(
148        config: &SaTokenConfig,
149        login_id: &str,
150        extra_data: &serde_json::Value,
151    ) -> TokenValue {
152        let effective_login_id = if login_id.is_empty() {
153            Utc::now().timestamp_millis().to_string()
154        } else {
155            login_id.to_string()
156        };
157        
158        let secret = config.jwt_secret_key.as_ref()
159            .expect("JWT secret key is required when using JWT token style");
160        
161        let algorithm = config.jwt_algorithm.as_ref()
162            .and_then(|alg| Self::parse_jwt_algorithm(alg))
163            .unwrap_or(JwtAlgorithm::HS256);
164        
165        let mut jwt_manager = JwtManager::with_algorithm(secret, algorithm);
166        
167        if let Some(ref issuer) = config.jwt_issuer {
168            jwt_manager = jwt_manager.set_issuer(issuer);
169        }
170        
171        if let Some(ref audience) = config.jwt_audience {
172            jwt_manager = jwt_manager.set_audience(audience);
173        }
174        
175        let mut claims = JwtClaims::new(effective_login_id);
176        
177        if config.timeout > 0 {
178            claims.set_expiration(config.timeout);
179        }
180        
181        // 将 extra_data 写入 JWT claims
182        // If extra_data is an Object, flatten each key-value into claims.extra
183        // Otherwise, store the entire value under "extra" key
184        match extra_data {
185            serde_json::Value::Object(map) => {
186                for (key, value) in map {
187                    claims.add_claim(key.clone(), value.clone());
188                }
189            }
190            serde_json::Value::Null => {
191                // Null 值不写入
192            }
193            other => {
194                claims.add_claim("extra", other.clone());
195            }
196        }
197        
198        match jwt_manager.generate(&claims) {
199            Ok(token) => TokenValue::new(token),
200            Err(e) => {
201                eprintln!("Failed to generate JWT token with extra: {:?}", e);
202                Self::generate_uuid()
203            }
204        }
205    }
206    
207    /// Parse JWT algorithm from string | 从字符串解析 JWT 算法
208    fn parse_jwt_algorithm(alg: &str) -> Option<JwtAlgorithm> {
209        match alg.to_uppercase().as_str() {
210            "HS256" => Some(JwtAlgorithm::HS256),
211            "HS384" => Some(JwtAlgorithm::HS384),
212            "HS512" => Some(JwtAlgorithm::HS512),
213            "RS256" => Some(JwtAlgorithm::RS256),
214            "RS384" => Some(JwtAlgorithm::RS384),
215            "RS512" => Some(JwtAlgorithm::RS512),
216            "ES256" => Some(JwtAlgorithm::ES256),
217            "ES384" => Some(JwtAlgorithm::ES384),
218            _ => None,
219        }
220    }
221    
222    /// Generate Hash style token | 生成 Hash 风格 token
223    ///
224    /// Uses SHA256 hash of login_id + timestamp + random UUID
225    /// 使用 SHA256 哈希：login_id + 时间戳 + 随机 UUID
226    ///
227    /// # Arguments | 参数
228    ///
229    /// * `login_id` - User login ID | 用户登录ID
230    pub fn generate_hash(login_id: &str) -> TokenValue {
231        // 如果 login_id 为空，使用时间戳代替
232        let login_id_value = if login_id.is_empty() {
233            Utc::now().timestamp_millis().to_string()
234        } else {
235            login_id.to_string()
236        };
237        
238        let timestamp = Utc::now().timestamp_millis();
239        let uuid = Uuid::new_v4();
240        let data = format!("{}{}{}", login_id_value, timestamp, uuid);
241        
242        let mut hasher = Sha256::new();
243        hasher.update(data.as_bytes());
244        let result = hasher.finalize();
245        let hash = hex::encode(result);
246        
247        TokenValue::new(hash)
248    }
249    
250    /// Generate Timestamp style token | 生成时间戳风格 token
251    ///
252    /// Format: timestamp_milliseconds + 16-char random suffix
253    /// 格式：毫秒级时间戳 + 16位随机后缀
254    ///
255    /// Example: 1760403556789_a3b2c1d4e5f6g7h8
256    /// 示例：1760403556789_a3b2c1d4e5f6g7h8
257    pub fn generate_timestamp() -> TokenValue {
258        use chrono::Utc;
259        use sha2::{Sha256, Digest};
260        
261        let timestamp = Utc::now().timestamp_millis();
262        let uuid = Uuid::new_v4();
263        
264        // Generate random suffix | 生成随机后缀
265        let mut hasher = Sha256::new();
266        hasher.update(uuid.as_bytes());
267        let result = hasher.finalize();
268        let suffix = hex::encode(&result[..8]); // 16 characters
269        
270        TokenValue::new(format!("{}_{}", timestamp, suffix))
271    }
272    
273    /// Generate Tik style token | 生成 Tik 风格 token
274    ///
275    /// Short 8-character alphanumeric token (URL-safe)
276    /// 短小精悍的8位字母数字 token（URL安全）
277    ///
278    /// Character set: A-Z, a-z, 0-9 (62 characters)
279    /// 字符集：A-Z, a-z, 0-9（62个字符）
280    ///
281    /// Example: aB3dE9fG
282    /// 示例：aB3dE9fG
283    pub fn generate_tik() -> TokenValue {
284        use sha2::{Sha256, Digest};
285        
286        const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
287        const TOKEN_LENGTH: usize = 8;
288        
289        let uuid = Uuid::new_v4();
290        let mut hasher = Sha256::new();
291        hasher.update(uuid.as_bytes());
292        let hash = hasher.finalize();
293        
294        let mut token = String::with_capacity(TOKEN_LENGTH);
295        for i in 0..TOKEN_LENGTH {
296            let idx = (hash[i] as usize) % CHARSET.len();
297            token.push(CHARSET[idx] as char);
298        }
299        
300        TokenValue::new(token)
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::config::{SaTokenConfig, TokenStyle};
308    use crate::token::jwt::JwtManager;
309
310    fn jwt_config() -> SaTokenConfig {
311        SaTokenConfig {
312            token_style: TokenStyle::Jwt,
313            jwt_secret_key: Some("test-secret-key-for-jwt".to_string()),
314            timeout: 3600,
315            ..SaTokenConfig::default()
316        }
317    }
318
319    #[test]
320    fn test_generate_jwt_with_extra_object() {
321        let config = jwt_config();
322        let extra = serde_json::json!({
323            "role": "admin",
324            "tenant_id": 42,
325            "permissions": ["read", "write"]
326        });
327
328        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_123", &extra);
329        assert!(!token.as_str().is_empty());
330
331        // 解析 JWT 验证 extra 数据已签入
332        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
333        let claims = jwt_manager.validate(token.as_str()).unwrap();
334
335        assert_eq!(claims.login_id, "user_123");
336        assert_eq!(claims.get_claim("role"), Some(&serde_json::json!("admin")));
337        assert_eq!(claims.get_claim("tenant_id"), Some(&serde_json::json!(42)));
338        assert_eq!(
339            claims.get_claim("permissions"),
340            Some(&serde_json::json!(["read", "write"]))
341        );
342    }
343
344    #[test]
345    fn test_generate_jwt_with_extra_non_object() {
346        let config = jwt_config();
347        let extra = serde_json::json!("simple_string_value");
348
349        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_456", &extra);
350
351        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
352        let claims = jwt_manager.validate(token.as_str()).unwrap();
353
354        assert_eq!(claims.login_id, "user_456");
355        assert_eq!(
356            claims.get_claim("extra"),
357            Some(&serde_json::json!("simple_string_value"))
358        );
359    }
360
361    #[test]
362    fn test_generate_jwt_with_extra_null() {
363        let config = jwt_config();
364        let extra = serde_json::Value::Null;
365
366        let token = TokenGenerator::generate_jwt_with_extra(&config, "user_789", &extra);
367
368        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
369        let claims = jwt_manager.validate(token.as_str()).unwrap();
370
371        assert_eq!(claims.login_id, "user_789");
372        assert!(claims.extra.is_empty());
373    }
374
375    #[test]
376    fn test_generate_with_login_id_and_extra_jwt_style() {
377        let config = jwt_config();
378        let extra = serde_json::json!({"key": "value"});
379
380        let token = TokenGenerator::generate_with_login_id_and_extra(&config, "user_jwt", &extra);
381
382        // JWT token 包含两个 '.' 分隔符
383        assert!(token.as_str().contains('.'));
384
385        let jwt_manager = JwtManager::new("test-secret-key-for-jwt");
386        let claims = jwt_manager.validate(token.as_str()).unwrap();
387        assert_eq!(claims.get_claim("key"), Some(&serde_json::json!("value")));
388    }
389
390    #[test]
391    fn test_generate_with_login_id_and_extra_non_jwt_style() {
392        let config = SaTokenConfig {
393            token_style: TokenStyle::Uuid,
394            ..SaTokenConfig::default()
395        };
396        let extra = serde_json::json!({"key": "value"});
397
398        // 非 JWT 风格应该走正常生成逻辑，不 panic
399        let token = TokenGenerator::generate_with_login_id_and_extra(&config, "user_uuid", &extra);
400        assert!(!token.as_str().is_empty());
401        // UUID 格式不包含 '.'
402        assert!(!token.as_str().contains('.'));
403    }
404
405
406    #[test]
407    fn test_random_32_length() {
408        let config = SaTokenConfig {
409            token_style: TokenStyle::Random32,
410            ..SaTokenConfig::default()
411        };
412        let token = TokenGenerator::generate_with_login_id(&config, "user_random");
413        assert!(!token.as_str().is_empty());
414        assert_eq!(token.as_str().len(), 32);
415    }
416
417    #[test]
418    fn test_random_64_length() {
419        let config = SaTokenConfig {
420            token_style: TokenStyle::Random64,
421            ..SaTokenConfig::default()
422        };
423        let token = TokenGenerator::generate_with_login_id(&config, "user_random");
424        assert!(!token.as_str().is_empty());
425        assert_eq!(token.as_str().len(), 64);
426    }
427    
428    #[test]
429    fn test_random_128_length() {
430        let config = SaTokenConfig {
431            token_style: TokenStyle::Random128,
432            ..SaTokenConfig::default()
433        };
434        let token = TokenGenerator::generate_with_login_id(&config, "user_random");
435        assert!(!token.as_str().is_empty());
436        assert_eq!(token.as_str().len(), 128);
437    }
438}
sa_token_core/token/generator.rs

sa_token_core/token/
generator.rs
