1use crate::errors::{AuthError, Result, TokenError};
4use chrono::{DateTime, Utc};
5use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation, Algorithm};
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::time::Duration;
9use uuid::Uuid;
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct AuthToken {
14 pub token_id: String,
16
17 pub user_id: String,
19
20 pub access_token: String,
22
23 pub refresh_token: Option<String>,
25
26 pub token_type: String,
28
29 pub issued_at: DateTime<Utc>,
31
32 pub expires_at: DateTime<Utc>,
34
35 pub scopes: Vec<String>,
37
38 pub auth_method: String,
40
41 pub client_id: Option<String>,
43
44 pub metadata: TokenMetadata,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
50#[derive(Default)]
51pub struct TokenMetadata {
52 pub issued_ip: Option<String>,
54
55 pub user_agent: Option<String>,
57
58 pub device_id: Option<String>,
60
61 pub session_id: Option<String>,
63
64 pub revoked: bool,
66
67 pub revoked_at: Option<DateTime<Utc>>,
69
70 pub revoked_reason: Option<String>,
72
73 pub last_used: Option<DateTime<Utc>>,
75
76 pub use_count: u64,
78
79 pub custom: HashMap<String, serde_json::Value>,
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct TokenInfo {
86 pub user_id: String,
88
89 pub username: Option<String>,
91
92 pub email: Option<String>,
94
95 pub name: Option<String>,
97
98 pub roles: Vec<String>,
100
101 pub permissions: Vec<String>,
103
104 pub attributes: HashMap<String, serde_json::Value>,
106}
107
108#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct JwtClaims {
111 pub sub: String,
113
114 pub iss: String,
116
117 pub aud: String,
119
120 pub exp: i64,
122
123 pub iat: i64,
125
126 pub nbf: i64,
128
129 pub jti: String,
131
132 pub scope: String,
134
135 #[serde(flatten)]
137 pub custom: HashMap<String, serde_json::Value>,
138}
139
140pub struct TokenManager {
142 encoding_key: EncodingKey,
144
145 decoding_key: DecodingKey,
147
148 algorithm: Algorithm,
150
151 issuer: String,
153
154 audience: String,
156
157 default_lifetime: Duration,
159}
160
161
162impl AuthToken {
163 pub fn new(
165 user_id: impl Into<String>,
166 access_token: impl Into<String>,
167 expires_in: std::time::Duration,
168 auth_method: impl Into<String>,
169 ) -> Self {
170 let now = Utc::now();
171 let expires_in_chrono = chrono::Duration::from_std(expires_in).unwrap_or(chrono::Duration::hours(1));
172
173 Self {
174 token_id: Uuid::new_v4().to_string(),
175 user_id: user_id.into(),
176 access_token: access_token.into(),
177 refresh_token: None,
178 token_type: "Bearer".to_string(),
179 issued_at: now,
180 expires_at: now + expires_in_chrono,
181 scopes: Vec::new(),
182 auth_method: auth_method.into(),
183 client_id: None,
184 metadata: TokenMetadata::default(),
185 }
186 }
187
188 pub fn access_token(&self) -> &str {
190 &self.access_token
191 }
192
193 pub fn user_id(&self) -> &str {
195 &self.user_id
196 }
197
198 pub fn expires_at(&self) -> DateTime<Utc> {
200 self.expires_at
201 }
202
203 pub fn is_expired(&self) -> bool {
205 Utc::now() > self.expires_at
206 }
207
208 pub fn is_expiring(&self, within: Duration) -> bool {
210 Utc::now() + within > self.expires_at
211 }
212
213 pub fn is_revoked(&self) -> bool {
215 self.metadata.revoked
216 }
217
218 pub fn is_valid(&self) -> bool {
220 !self.is_expired() && !self.is_revoked()
221 }
222
223 pub fn revoke(&mut self, reason: Option<String>) {
225 self.metadata.revoked = true;
226 self.metadata.revoked_at = Some(Utc::now());
227 self.metadata.revoked_reason = reason;
228 }
229
230 pub fn mark_used(&mut self) {
232 self.metadata.last_used = Some(Utc::now());
233 self.metadata.use_count += 1;
234 }
235
236 pub fn add_scope(&mut self, scope: impl Into<String>) {
238 let scope = scope.into();
239 if !self.scopes.contains(&scope) {
240 self.scopes.push(scope);
241 }
242 }
243
244 pub fn has_scope(&self, scope: &str) -> bool {
246 self.scopes.contains(&scope.to_string())
247 }
248
249 pub fn with_refresh_token(mut self, refresh_token: impl Into<String>) -> Self {
251 self.refresh_token = Some(refresh_token.into());
252 self
253 }
254
255 pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
257 self.client_id = Some(client_id.into());
258 self
259 }
260
261 pub fn with_scopes(mut self, scopes: Vec<String>) -> Self {
263 self.scopes = scopes;
264 self
265 }
266
267 pub fn with_metadata(mut self, metadata: TokenMetadata) -> Self {
269 self.metadata = metadata;
270 self
271 }
272
273 pub fn time_until_expiry(&self) -> Duration {
275 let now = Utc::now();
276 if self.expires_at > now {
277 (self.expires_at - now).to_std().unwrap_or(Duration::ZERO)
278 } else {
279 Duration::ZERO
280 }
281 }
282}
283
284impl TokenManager {
285 pub fn new_hmac(secret: &[u8], issuer: impl Into<String>, audience: impl Into<String>) -> Self {
287 Self {
288 encoding_key: EncodingKey::from_secret(secret),
289 decoding_key: DecodingKey::from_secret(secret),
290 algorithm: Algorithm::HS256,
291 issuer: issuer.into(),
292 audience: audience.into(),
293 default_lifetime: Duration::from_secs(3600), }
295 }
296
297 pub fn new_rsa(
299 private_key: &[u8],
300 public_key: &[u8],
301 issuer: impl Into<String>,
302 audience: impl Into<String>,
303 ) -> Result<Self> {
304 let encoding_key = EncodingKey::from_rsa_pem(private_key)
305 .map_err(|e| AuthError::crypto(format!("Invalid RSA private key: {e}")))?;
306
307 let decoding_key = DecodingKey::from_rsa_pem(public_key)
308 .map_err(|e| AuthError::crypto(format!("Invalid RSA public key: {e}")))?;
309
310 Ok(Self {
311 encoding_key,
312 decoding_key,
313 algorithm: Algorithm::RS256,
314 issuer: issuer.into(),
315 audience: audience.into(),
316 default_lifetime: Duration::from_secs(3600), })
318 }
319
320 pub fn with_default_lifetime(mut self, lifetime: Duration) -> Self {
322 self.default_lifetime = lifetime;
323 self
324 }
325
326 pub fn create_jwt_token(
328 &self,
329 user_id: impl Into<String>,
330 scopes: Vec<String>,
331 lifetime: Option<Duration>,
332 ) -> Result<String> {
333 let user_id = user_id.into();
334 let lifetime = lifetime.unwrap_or(self.default_lifetime);
335 let now = Utc::now();
336 let exp = now + chrono::Duration::from_std(lifetime).unwrap_or(chrono::Duration::hours(1));
337
338 let claims = JwtClaims {
339 sub: user_id,
340 iss: self.issuer.clone(),
341 aud: self.audience.clone(),
342 exp: exp.timestamp(),
343 iat: now.timestamp(),
344 nbf: now.timestamp(),
345 jti: Uuid::new_v4().to_string(),
346 scope: scopes.join(" "),
347 custom: HashMap::new(),
348 };
349
350 let header = Header::new(self.algorithm);
351
352 encode(&header, &claims, &self.encoding_key)
353 .map_err(|e| TokenError::creation_failed(format!("JWT encoding failed: {e}")).into())
354 }
355
356 pub fn validate_jwt_token(&self, token: &str) -> Result<JwtClaims> {
358 let mut validation = Validation::new(self.algorithm);
359 validation.set_issuer(&[&self.issuer]);
360 validation.set_audience(&[&self.audience]);
361
362 let token_data = decode::<JwtClaims>(token, &self.decoding_key, &validation)
363 .map_err(|e| match e.kind() {
364 jsonwebtoken::errors::ErrorKind::ExpiredSignature => AuthError::Token(TokenError::Expired),
365 _ => AuthError::Token(TokenError::Invalid {
366 message: "Invalid token format".to_string(),
367 }),
368 })?;
369
370 Ok(token_data.claims)
371 }
372
373 pub fn create_auth_token(
375 &self,
376 user_id: impl Into<String>,
377 scopes: Vec<String>,
378 auth_method: impl Into<String>,
379 lifetime: Option<std::time::Duration>,
380 ) -> Result<AuthToken> {
381 let user_id_str = user_id.into();
382 let lifetime = lifetime.unwrap_or(self.default_lifetime);
383
384 let jwt_token = self.create_jwt_token(&user_id_str, scopes.clone(), Some(lifetime))?;
385
386 let token = AuthToken::new(user_id_str, jwt_token, lifetime, auth_method)
387 .with_scopes(scopes);
388
389 Ok(token)
390 }
391
392 pub fn validate_auth_token(&self, token: &AuthToken) -> Result<()> {
394 if token.is_expired() {
396 return Err(TokenError::Expired.into());
397 }
398
399 if token.is_revoked() {
401 return Err(TokenError::Invalid {
402 message: "Token has been revoked".to_string(),
403 }.into());
404 }
405
406 if token.auth_method == "jwt" || token.access_token.contains('.') {
408 self.validate_jwt_token(&token.access_token)?;
409 }
410
411 Ok(())
412 }
413
414 pub fn refresh_token(&self, token: &AuthToken) -> Result<AuthToken> {
416 if token.is_expired() {
417 return Err(TokenError::Expired.into());
418 }
419
420 if token.is_revoked() {
421 return Err(TokenError::Invalid {
422 message: "Cannot refresh revoked token".to_string(),
423 }.into());
424 }
425
426 self.create_auth_token(
428 &token.user_id,
429 token.scopes.clone(),
430 &token.auth_method,
431 Some(self.default_lifetime),
432 )
433 }
434
435 pub fn extract_token_info(&self, token: &str) -> Result<TokenInfo> {
437 let claims = self.validate_jwt_token(token)?;
438
439 Ok(TokenInfo {
440 user_id: claims.sub,
441 username: claims.custom.get("username")
442 .and_then(|v| v.as_str())
443 .map(|s| s.to_string()),
444 email: claims.custom.get("email")
445 .and_then(|v| v.as_str())
446 .map(|s| s.to_string()),
447 name: claims.custom.get("name")
448 .and_then(|v| v.as_str())
449 .map(|s| s.to_string()),
450 roles: claims.custom.get("roles")
451 .and_then(|v| v.as_array())
452 .map(|arr| arr.iter()
453 .filter_map(|v| v.as_str())
454 .map(|s| s.to_string())
455 .collect())
456 .unwrap_or_default(),
457 permissions: claims.scope.split_whitespace()
458 .map(|s| s.to_string())
459 .collect(),
460 attributes: claims.custom,
461 })
462 }
463}
464
465#[cfg(test)]
466mod tests {
467 use super::*;
468
469 #[test]
470 fn test_auth_token_creation() {
471 let token = AuthToken::new(
472 "user123",
473 "token123",
474 Duration::from_secs(3600), "password"
476 );
477
478 assert_eq!(token.user_id(), "user123");
479 assert_eq!(token.access_token(), "token123");
480 assert!(!token.is_expired());
481 assert!(!token.is_revoked());
482 assert!(token.is_valid());
483 }
484
485 #[test]
486 fn test_token_expiry() {
487 let token = AuthToken::new(
488 "user123",
489 "token123",
490 Duration::from_millis(1),
491 "password"
492 );
493
494 std::thread::sleep(std::time::Duration::from_millis(10));
496
497 assert!(token.is_expired());
498 assert!(!token.is_valid());
499 }
500
501 #[test]
502 fn test_token_revocation() {
503 let mut token = AuthToken::new(
504 "user123",
505 "token123",
506 Duration::from_secs(3600), "password"
508 );
509
510 assert!(!token.is_revoked());
511
512 token.revoke(Some("User logout".to_string()));
513
514 assert!(token.is_revoked());
515 assert!(!token.is_valid());
516 assert_eq!(token.metadata.revoked_reason, Some("User logout".to_string()));
517 }
518
519 #[tokio::test]
520 async fn test_jwt_token_manager() {
521 let secret = b"test-secret-key";
522 let manager = TokenManager::new_hmac(secret, "test-issuer", "test-audience");
523
524 let token = manager.create_jwt_token(
525 "user123",
526 vec!["read".to_string(), "write".to_string()],
527 Some(Duration::from_secs(3600)) ).unwrap();
529
530 let claims = manager.validate_jwt_token(&token).unwrap();
531 assert_eq!(claims.sub, "user123");
532 assert_eq!(claims.scope, "read write");
533 }
534}