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

sa_token_core/token/
generator.rs
