auth_framework/
tokens.rs

1//! Token management and validation for the authentication framework.
2
3use 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/// Represents an authentication token with all associated metadata.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct AuthToken {
14    /// Unique token identifier
15    pub token_id: String,
16    
17    /// User identifier this token belongs to
18    pub user_id: String,
19    
20    /// The actual token string (JWT, opaque token, etc.)
21    pub access_token: String,
22    
23    /// Optional refresh token
24    pub refresh_token: Option<String>,
25    
26    /// Token type (Bearer, etc.)
27    pub token_type: String,
28    
29    /// When the token was issued
30    pub issued_at: DateTime<Utc>,
31    
32    /// When the token expires
33    pub expires_at: DateTime<Utc>,
34    
35    /// Scopes granted to this token
36    pub scopes: Vec<String>,
37    
38    /// Authentication method used to obtain this token
39    pub auth_method: String,
40    
41    /// Client ID that requested this token
42    pub client_id: Option<String>,
43    
44    /// Additional token metadata
45    pub metadata: TokenMetadata,
46}
47
48/// Additional metadata that can be attached to tokens.
49#[derive(Debug, Clone, Serialize, Deserialize)]
50#[derive(Default)]
51pub struct TokenMetadata {
52    /// IP address where the token was issued
53    pub issued_ip: Option<String>,
54    
55    /// User agent of the client
56    pub user_agent: Option<String>,
57    
58    /// Device identifier
59    pub device_id: Option<String>,
60    
61    /// Session identifier
62    pub session_id: Option<String>,
63    
64    /// Whether this token has been revoked
65    pub revoked: bool,
66    
67    /// When the token was revoked (if applicable)
68    pub revoked_at: Option<DateTime<Utc>>,
69    
70    /// Reason for revocation
71    pub revoked_reason: Option<String>,
72    
73    /// Last time this token was used
74    pub last_used: Option<DateTime<Utc>>,
75    
76    /// Number of times this token has been used
77    pub use_count: u64,
78    
79    /// Custom metadata
80    pub custom: HashMap<String, serde_json::Value>,
81}
82
83/// Information about a user extracted from a token.
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct TokenInfo {
86    /// User identifier
87    pub user_id: String,
88    
89    /// Username or email
90    pub username: Option<String>,
91    
92    /// User's email address
93    pub email: Option<String>,
94    
95    /// User's display name
96    pub name: Option<String>,
97    
98    /// User's roles
99    pub roles: Vec<String>,
100    
101    /// User's permissions
102    pub permissions: Vec<String>,
103    
104    /// Additional user attributes
105    pub attributes: HashMap<String, serde_json::Value>,
106}
107
108/// JWT claims structure used internally.
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct JwtClaims {
111    /// Subject (user ID)
112    pub sub: String,
113    
114    /// Issuer
115    pub iss: String,
116    
117    /// Audience
118    pub aud: String,
119    
120    /// Expiration time
121    pub exp: i64,
122    
123    /// Issued at
124    pub iat: i64,
125    
126    /// Not before
127    pub nbf: i64,
128    
129    /// JWT ID
130    pub jti: String,
131    
132    /// Scopes
133    pub scope: String,
134    
135    /// Custom claims
136    #[serde(flatten)]
137    pub custom: HashMap<String, serde_json::Value>,
138}
139
140/// Token manager for creating, validating, and managing tokens.
141pub struct TokenManager {
142    /// JWT encoding key
143    encoding_key: EncodingKey,
144    
145    /// JWT decoding key
146    decoding_key: DecodingKey,
147    
148    /// JWT algorithm
149    algorithm: Algorithm,
150    
151    /// Token issuer
152    issuer: String,
153    
154    /// Token audience
155    audience: String,
156    
157    /// Default token lifetime
158    default_lifetime: Duration,
159}
160
161
162impl AuthToken {
163    /// Create a new authentication token.
164    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    /// Get the access token string.
189    pub fn access_token(&self) -> &str {
190        &self.access_token
191    }
192
193    /// Get the user ID.
194    pub fn user_id(&self) -> &str {
195        &self.user_id
196    }
197
198    /// Get the expiration time.
199    pub fn expires_at(&self) -> DateTime<Utc> {
200        self.expires_at
201    }
202
203    /// Check if the token has expired.
204    pub fn is_expired(&self) -> bool {
205        Utc::now() > self.expires_at
206    }
207
208    /// Check if the token is expiring within the given duration.
209    pub fn is_expiring(&self, within: Duration) -> bool {
210        Utc::now() + within > self.expires_at
211    }
212
213    /// Check if the token has been revoked.
214    pub fn is_revoked(&self) -> bool {
215        self.metadata.revoked
216    }
217
218    /// Check if the token is valid (not expired and not revoked).
219    pub fn is_valid(&self) -> bool {
220        !self.is_expired() && !self.is_revoked()
221    }
222
223    /// Revoke the token.
224    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    /// Update the last used time and increment use count.
231    pub fn mark_used(&mut self) {
232        self.metadata.last_used = Some(Utc::now());
233        self.metadata.use_count += 1;
234    }
235
236    /// Add a scope to the token.
237    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    /// Check if the token has a specific scope.
245    pub fn has_scope(&self, scope: &str) -> bool {
246        self.scopes.contains(&scope.to_string())
247    }
248
249    /// Set the refresh token.
250    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    /// Set the client ID.
256    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    /// Set the token scopes.
262    pub fn with_scopes(mut self, scopes: Vec<String>) -> Self {
263        self.scopes = scopes;
264        self
265    }
266
267    /// Add metadata to the token.
268    pub fn with_metadata(mut self, metadata: TokenMetadata) -> Self {
269        self.metadata = metadata;
270        self
271    }
272
273    /// Get time until expiration.
274    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    /// Create a new token manager with HMAC key.
286    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), // 1 hour
294        }
295    }
296
297    /// Create a new token manager with RSA keys.
298    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), // 1 hour
317        })
318    }
319
320    /// Set the default token lifetime.
321    pub fn with_default_lifetime(mut self, lifetime: Duration) -> Self {
322        self.default_lifetime = lifetime;
323        self
324    }
325
326    /// Create a new JWT token.
327    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    /// Validate and decode a JWT token.
357    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    /// Create a complete authentication token with JWT.
374    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    /// Validate an authentication token.
393    pub fn validate_auth_token(&self, token: &AuthToken) -> Result<()> {
394        // Check if token is expired
395        if token.is_expired() {
396            return Err(TokenError::Expired.into());
397        }
398
399        // Check if token is revoked
400        if token.is_revoked() {
401            return Err(TokenError::Invalid {
402                message: "Token has been revoked".to_string(),
403            }.into());
404        }
405
406        // Validate JWT if it's a JWT token
407        if token.auth_method == "jwt" || token.access_token.contains('.') {
408            self.validate_jwt_token(&token.access_token)?;
409        }
410
411        Ok(())
412    }
413
414    /// Refresh a token (create a new one with extended lifetime).
415    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        // Create a new token with the same properties but new expiry
427        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    /// Extract token information from a JWT.
436    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), // 1 hour
475            "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        // Wait a bit to ensure expiry
495        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), // 1 hour
507            "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)) // 1 hour
528        ).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}