1use crate::auth::{error::AuthError, state::Identity};
2
3use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7#[derive(Debug, Serialize, Deserialize, Clone)]
8pub struct Claims {
9 pub iss: Option<String>,
11 pub sub: String,
12 pub aud: Option<String>,
13 pub exp: usize,
14 pub iat: usize,
15 pub nbf: Option<usize>,
16 pub jti: Option<String>,
17
18 pub scope: Option<String>,
20 #[serde(skip_serializing_if = "Option::is_none")]
23 pub identity: Option<Identity>,
24
25 #[serde(flatten)]
27 pub extra: HashMap<String, serde_json::Value>,
28}
29
30#[derive(Clone)]
31pub struct TokenManager {
32 encoding_key: EncodingKey,
33 decoding_key: DecodingKey,
34 issuer: Option<String>,
35 kid: Option<String>,
36 alg: Algorithm,
37 public_jwk: Option<crate::token::jwk::Jwk>,
38}
39
40impl TokenManager {
41 pub fn new(secret: &[u8], issuer: Option<String>) -> Self {
43 Self {
44 encoding_key: EncodingKey::from_secret(secret),
45 decoding_key: DecodingKey::from_secret(secret),
46 issuer,
47 kid: None,
48 alg: Algorithm::HS256,
49 public_jwk: None,
50 }
51 }
52
53 pub fn new_asymmetric(
58 private_key_pem: &[u8],
59 issuer: Option<String>,
60 kid: Option<String>,
61 ) -> Result<Self, AuthError> {
62 let encoding_key = EncodingKey::from_rsa_pem(private_key_pem)
63 .map_err(|e| AuthError::Token(e.to_string()))?;
64 let decoding_key = DecodingKey::from_rsa_pem(private_key_pem)
65 .map_err(|e| AuthError::Token(e.to_string()))?;
66
67 let pem_str = std::str::from_utf8(private_key_pem)
68 .map_err(|_| AuthError::Token("Invalid PEM UTF-8".into()))?;
69
70 use rsa::pkcs1::DecodeRsaPrivateKey;
71 use rsa::pkcs8::DecodePrivateKey;
72 let rsa_key = rsa::RsaPrivateKey::from_pkcs8_pem(pem_str)
73 .or_else(|_| rsa::RsaPrivateKey::from_pkcs1_pem(pem_str))
74 .map_err(|e| AuthError::Token(format!("Failed to parse RSA key: {}", e)))?;
75
76 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
77 use rsa::traits::PublicKeyParts;
78
79 let n = URL_SAFE_NO_PAD.encode(rsa_key.n().to_bytes_be());
80 let e = URL_SAFE_NO_PAD.encode(rsa_key.e().to_bytes_be());
81
82 let kid_val = kid.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
83
84 let jwk = crate::token::jwk::Jwk {
85 kid: Some(kid_val.clone()),
86 kty: "RSA".to_string(),
87 alg: Some("RS256".to_string()),
88 n: Some(n),
89 e: Some(e),
90 };
91
92 Ok(Self {
93 encoding_key,
94 decoding_key,
95 issuer,
96 kid: Some(kid_val),
97 alg: Algorithm::RS256,
98 public_jwk: Some(jwk),
99 })
100 }
101
102 pub fn public_jwk(&self) -> Option<crate::token::jwk::Jwk> {
103 self.public_jwk.clone()
104 }
105
106 pub fn with_issuer(mut self, issuer: String) -> Self {
107 self.issuer = Some(issuer);
108 self
109 }
110
111 pub fn issue_user_token(
113 &self,
114 identity: Identity,
115 expires_in_secs: u64,
116 scope: Option<String>,
117 aud: Option<String>,
118 ) -> Result<String, AuthError> {
119 self.issue_user_token_with_extra(identity, expires_in_secs, scope, aud, HashMap::new())
120 }
121
122 pub fn issue_user_token_with_extra(
134 &self,
135 identity: Identity,
136 expires_in_secs: u64,
137 scope: Option<String>,
138 aud: Option<String>,
139 extra: HashMap<String, serde_json::Value>,
140 ) -> Result<String, AuthError> {
141 let now = chrono::Utc::now().timestamp() as usize;
142 let expiration = now + expires_in_secs as usize;
143
144 let claims = Claims {
145 iss: self.issuer.clone(),
146 sub: identity.external_id.clone(),
147 aud,
148 exp: expiration,
149 iat: now,
150 nbf: Some(now),
151 jti: Some(uuid::Uuid::new_v4().to_string()),
152 scope,
153 identity: Some(identity),
154 extra,
155 };
156
157 let mut header = Header::new(self.alg);
158 if let Some(ref kid) = self.kid {
159 header.kid = Some(kid.clone());
160 }
161
162 encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
163 }
164
165 pub fn issue_id_token(
167 &self,
168 identity: Identity,
169 client_id: &str,
170 nonce: Option<String>,
171 expires_in_secs: u64,
172 ) -> Result<String, AuthError> {
173 self.issue_id_token_with_extra(identity, client_id, nonce, expires_in_secs, HashMap::new())
174 }
175
176 pub fn issue_id_token_with_extra(
187 &self,
188 identity: Identity,
189 client_id: &str,
190 nonce: Option<String>,
191 expires_in_secs: u64,
192 extra: HashMap<String, serde_json::Value>,
193 ) -> Result<String, AuthError> {
194 let now = chrono::Utc::now().timestamp() as usize;
195 let expiration = now + expires_in_secs as usize;
196
197 let mut claims = Claims {
198 iss: self.issuer.clone(),
199 sub: identity.external_id.clone(),
200 aud: Some(client_id.to_string()),
201 exp: expiration,
202 iat: now,
203 nbf: Some(now),
204 jti: Some(uuid::Uuid::new_v4().to_string()),
205 scope: None,
206 identity: Some(identity),
207 extra,
208 };
209
210 if let Some(n) = nonce {
211 claims
212 .extra
213 .insert("nonce".to_string(), serde_json::Value::String(n));
214 }
215
216 let mut header = Header::new(self.alg);
217 if let Some(ref kid) = self.kid {
218 header.kid = Some(kid.clone());
219 }
220
221 encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
222 }
223
224 pub fn issue_client_token(
226 &self,
227 client_id: &str,
228 expires_in_secs: u64,
229 scope: Option<String>,
230 aud: Option<String>,
231 ) -> Result<String, AuthError> {
232 self.issue_client_token_with_extra(client_id, expires_in_secs, scope, aud, HashMap::new())
233 }
234
235 pub fn issue_client_token_with_extra(
239 &self,
240 client_id: &str,
241 expires_in_secs: u64,
242 scope: Option<String>,
243 aud: Option<String>,
244 extra: HashMap<String, serde_json::Value>,
245 ) -> Result<String, AuthError> {
246 let now = chrono::Utc::now().timestamp() as usize;
247 let expiration = now + expires_in_secs as usize;
248
249 let claims = Claims {
250 iss: self.issuer.clone(),
251 sub: client_id.to_string(),
252 aud,
253 exp: expiration,
254 iat: now,
255 nbf: Some(now),
256 jti: Some(uuid::Uuid::new_v4().to_string()),
257 scope,
258 identity: None,
259 extra,
260 };
261
262 let mut header = Header::new(self.alg);
263 if let Some(ref kid) = self.kid {
264 header.kid = Some(kid.clone());
265 }
266
267 encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
268 }
269
270 pub fn issue_custom_token(
278 &self,
279 sub: String,
280 expires_in_secs: u64,
281 typ: &str,
282 extra: HashMap<String, serde_json::Value>,
283 ) -> Result<String, AuthError> {
284 let now = chrono::Utc::now().timestamp() as usize;
285 let claims = Claims {
286 iss: self.issuer.clone(),
287 sub,
288 aud: None,
289 exp: now + expires_in_secs as usize,
290 iat: now,
291 nbf: Some(now),
292 jti: Some(uuid::Uuid::new_v4().to_string()),
293 scope: None,
294 identity: None,
295 extra,
296 };
297
298 let mut header = Header::new(self.alg);
299 header.typ = Some(typ.to_string());
300 if let Some(ref kid) = self.kid {
301 header.kid = Some(kid.clone());
302 }
303
304 encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
305 }
306
307 pub fn validate_token(
308 &self,
309 token: &str,
310 expected_aud: Option<&str>,
311 ) -> Result<Claims, AuthError> {
312 let mut validation = Validation::new(self.alg);
313 if let Some(aud) = expected_aud {
314 validation.set_audience(&[aud]);
315 } else {
316 validation.validate_aud = false;
317 }
318 if let Some(ref iss) = self.issuer {
319 validation.set_issuer(&[iss]);
320 }
321
322 let token_data = decode::<Claims>(token, &self.decoding_key, &validation)
323 .map_err(|e| AuthError::Token(e.to_string()))?;
324
325 Ok(token_data.claims)
326 }
327}
328
329#[cfg(test)]
330mod tests {
331
332 use super::*;
333 use crate::auth::state::Identity;
334 use std::collections::HashMap;
335
336 #[test]
337 fn test_claims_serialization() {
338 let mut extra = HashMap::new();
339 extra.insert(
340 "custom".to_string(),
341 serde_json::Value::String("value".to_string()),
342 );
343
344 let claims = Claims {
345 iss: Some("issuer".to_string()),
346 sub: "user123".to_string(),
347 aud: Some("audience".to_string()),
348 exp: 1000,
349 iat: 500,
350 nbf: Some(500),
351 jti: Some("jti".to_string()),
352 scope: Some("openid profile".to_string()),
353 identity: Some(Identity {
354 provider_id: "google".to_string(),
355 external_id: "user123".to_string(),
356 email: Some("user@example.com".to_string()),
357 username: Some("user".to_string()),
358 attributes: HashMap::new(),
359 }),
360 extra,
361 };
362
363 let serialized = serde_json::to_string(&claims).unwrap();
364 let deserialized: Claims = serde_json::from_str(&serialized).unwrap();
365
366 assert_eq!(deserialized.iss, claims.iss);
367 assert_eq!(deserialized.sub, claims.sub);
368 assert_eq!(deserialized.extra.get("custom").unwrap(), "value");
369 }
370
371 #[test]
372 fn test_token_manager_issuance() {
373 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
374 let identity = Identity {
375 provider_id: "mock".to_string(),
376 external_id: "user123".to_string(),
377 email: None,
378 username: None,
379 attributes: HashMap::new(),
380 };
381
382 let token = manager
383 .issue_user_token(identity, 3600, None, None)
384 .unwrap();
385 let claims = manager.validate_token(&token, None).unwrap();
386
387 assert_eq!(claims.iss, Some("issuer".to_string()));
388 assert_eq!(claims.sub, "user123");
389 assert!(claims.jti.is_some());
390 assert!(claims.nbf.is_some());
391 }
392
393 #[test]
394 fn test_token_manager_asymmetric_issuance() {
395 let pem = b"-----BEGIN PRIVATE KEY-----
396MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDA5hJIcQ+2rxMz
397VM8ZH5WAmguCr0xmNDAdy0IzzsUeFLG7BebB7izOkU36J4t8t5tUaQwrBMnx2Fvt
398VqJjbdE242UDpvWF/8m9zJ2HR5298cbwT5cGMKLB0HWzDMahugs+Bbh2lCgwyLZk
399Tr3Diwxp5SwFew/Wb+Ke9cNG9Hu5IFH3BCuJ839d9hfqisIeYrBPfb52xxckM37R
4007zSGu/eDP/HZAeLkQuptZJW4A3u7xni14u4qyqXDqsHsYFNgJaxMSAwWgBRY6HNu
401TnvBArTXCiVfL+F73B2L6mdYr64g+QS9nK9v97MlJu/E3mSduz54pren4mpCHc9m
402/S2+VjCZAgMBAAECggEAASC9qQbGnL7XuExRDOIn/m4bWx92ehjo0lCTibhpY3LW
403umbSbpfbhmmuSj3CjW9VZsaM3hBTgSjoTX72lbY/eIUXD7c0memUK5pV4XcEIrQw
404AZlPIye6ckx4I7ZGnKasO8FoAel9dd7DXw36AuBK3LBzJwtzkEFsBc0e3/wixqmG
405UJBbbt/+5ya7CxyjuePaQhKtkLD5R6DpvN2XnCYq5nHJNJdvSVg1pOzsTHYIf+Ee
4062Rz42fGsfFKqeEQCcBFRZaGb/ELeP4c6UZdktZAvmHb1p1fursVZc6X9JXmiJ2OJ
407Kv2H2tMKuysP8L0fXFOMgkH2SVt6rcdHkO6xhlhWsQKBgQDqR8rAJeEE5BFoXA8T
408VVW6CLMlW51x4ey7PEGOaYh39dTG2Q+GZQBZ9G+SZk3f5Y85UCACSyc//4qaz/c3
4090nWsegZ+JPyymmuc79wzIAFFvXB7pL6wyn0Ed1P620kOZTtA8iBcXrsuxL+KP7iu
410MXfWmU1QiZpbndILtyDnY+70uwKBgQDSyCljWkydQCaPU+fiAXLxP8CvcJTSSNQD
411mVUlwJ+OpHnU+Alsi1rBavMgUtLlYbFqzH7NmYrLC8Yadq3ZOwLt0VEK0r8qstAL
4127QCDUD2WNuQjpZupRnXuMUl3iXB96i2gb+VQKGuUAJvVWjdIbYa4+Gu+sBMfcDcX
413dBihDLuEuwKBgAgX4tEwfc2Fc3R/eaXZVNTQaB/qQk4k1+C//CPHUYeTXn5gEUE7
414S//PiesszZPmgkQgmHp7zidP1KH0fT3Yb2g97ut8q54f54fMYXcCrAiUusYKsuu4
415kwkMdkI8QRHWPW3I74VBYIYFFfjYqrCZ1OH8+cbGeiagFRmCggh8U0zxAoGAVW3u
4166Ge22Z0gg8LcHsu7jG/sZq7Ygool8/d3fT+e669Z+ak2GJo6hF4WgClRdMqtn72W
417PzpV+ImjFyK2v26dd0n48MwN0v56N/ss1Av3iiRhPtlmR6tZLNspDZvUzhPVvkrb
418xCs9vtSoVEamVWKe0eVNthGjDoDqs0TInq2MavUCgYB6REavSJs/CLkSS7iimjxZ
419G7g5YQi9/p1lXLOEUDiwEmvRr0XTwzzxUsIc535IXhh/ZUYpthenW+qBBzn85pEC
420TowIqciHu5redqlQ8rITA8/AOY98vaDIhppDg1rfpnHHaZHFbXD/keYAEbhBtbvf
421a0QMqKUcs8+YTy5R5K6qtw==
422-----END PRIVATE KEY-----";
423
424 let manager = TokenManager::new_asymmetric(
425 pem,
426 Some("issuer".to_string()),
427 Some("my-kid-123".to_string()),
428 )
429 .unwrap();
430
431 let identity = Identity {
432 provider_id: "mock".to_string(),
433 external_id: "user123".to_string(),
434 email: None,
435 username: None,
436 attributes: HashMap::new(),
437 };
438
439 let token = manager
440 .issue_user_token(identity, 3600, None, None)
441 .unwrap();
442
443 let jwk = manager.public_jwk().unwrap();
445 assert_eq!(jwk.kid.as_deref(), Some("my-kid-123"));
446
447 let decoding_key = jwk.to_decoding_key().unwrap();
448 let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
449 validation.set_issuer(&["issuer"]);
450
451 let token_data =
452 jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap();
453 assert_eq!(token_data.claims.sub, "user123");
454 assert_eq!(token_data.header.kid.as_deref(), Some("my-kid-123"));
455 }
456
457 #[test]
458 fn test_issue_id_token() {
459 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
460 let identity = Identity {
461 provider_id: "mock".to_string(),
462 external_id: "user123".to_string(),
463 email: None,
464 username: None,
465 attributes: HashMap::new(),
466 };
467
468 let token = manager
469 .issue_id_token(identity, "client-1", Some("nonce123".to_string()), 3600)
470 .unwrap();
471
472 let claims = manager.validate_token(&token, None).unwrap();
473
474 assert_eq!(claims.iss, Some("issuer".to_string()));
475 assert_eq!(claims.sub, "user123");
476 assert_eq!(claims.aud, Some("client-1".to_string()));
477 assert_eq!(claims.extra.get("nonce").unwrap(), "nonce123");
478 }
479 #[test]
480 fn test_token_manager_audience_validation() {
481 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
482 let identity = Identity {
483 provider_id: "mock".to_string(),
484 external_id: "user123".to_string(),
485 email: None,
486 username: None,
487 attributes: HashMap::new(),
488 };
489
490 let token = manager
492 .issue_id_token(identity, "client-1", None, 3600)
493 .unwrap();
494
495 let claims = manager.validate_token(&token, Some("client-1")).unwrap();
497 assert_eq!(claims.aud, Some("client-1".to_string()));
498
499 let err = manager
501 .validate_token(&token, Some("client-2"))
502 .unwrap_err();
503 assert!(err.to_string().contains("InvalidAudience"));
504 }
505
506 #[test]
507 fn test_issue_custom_token_sets_typ_and_no_aud() {
508 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
509
510 let mut extra = HashMap::new();
511 extra.insert("cnf".to_string(), serde_json::json!({"jkt": "abc"}));
512
513 let token = manager
514 .issue_custom_token("device-1".to_string(), 60, "webank-attest+jws", extra)
515 .unwrap();
516
517 let header = jsonwebtoken::decode_header(&token).unwrap();
518 assert_eq!(header.typ.as_deref(), Some("webank-attest+jws"));
519
520 let claims = manager.validate_token(&token, None).unwrap();
521 assert_eq!(claims.sub, "device-1");
522 assert_eq!(claims.aud, None);
523 assert_eq!(claims.extra.get("cnf").unwrap()["jkt"], "abc");
524 }
525
526 #[test]
527 fn test_issue_user_token_with_extra_round_trips_custom_claims() {
528 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
529 let identity = Identity {
530 provider_id: "mock".to_string(),
531 external_id: "user123".to_string(),
532 email: None,
533 username: None,
534 attributes: HashMap::new(),
535 };
536
537 let mut extra = HashMap::new();
538 extra.insert("api_key_id".to_string(), serde_json::json!("key-abc"));
539 extra.insert("project_id".to_string(), serde_json::json!("proj-42"));
540
541 let token = manager
542 .issue_user_token_with_extra(identity, 3600, None, None, extra)
543 .unwrap();
544 let claims = manager.validate_token(&token, None).unwrap();
545
546 assert_eq!(
547 claims.extra.get("api_key_id"),
548 Some(&serde_json::json!("key-abc"))
549 );
550 assert_eq!(
551 claims.extra.get("project_id"),
552 Some(&serde_json::json!("proj-42"))
553 );
554 }
555
556 #[test]
557 fn test_issue_user_token_wrapper_still_has_empty_extra() {
558 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
559 let identity = Identity {
560 provider_id: "mock".to_string(),
561 external_id: "user123".to_string(),
562 email: None,
563 username: None,
564 attributes: HashMap::new(),
565 };
566
567 let token = manager
568 .issue_user_token(identity, 3600, None, None)
569 .unwrap();
570 let claims = manager.validate_token(&token, None).unwrap();
571
572 assert!(claims.extra.is_empty());
573 }
574
575 #[test]
576 fn test_issue_client_token_with_extra_round_trips_custom_claims() {
577 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
578
579 let mut extra = HashMap::new();
580 extra.insert("roles".to_string(), serde_json::json!(["admin", "billing"]));
581
582 let token = manager
583 .issue_client_token_with_extra("client-1", 3600, None, None, extra)
584 .unwrap();
585 let claims = manager.validate_token(&token, None).unwrap();
586
587 assert_eq!(
588 claims.extra.get("roles"),
589 Some(&serde_json::json!(["admin", "billing"]))
590 );
591 }
592
593 #[test]
594 fn test_issue_client_token_wrapper_still_has_empty_extra() {
595 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
596
597 let token = manager
598 .issue_client_token("client-1", 3600, None, None)
599 .unwrap();
600 let claims = manager.validate_token(&token, None).unwrap();
601
602 assert!(claims.extra.is_empty());
603 }
604
605 #[test]
606 fn test_issue_id_token_with_extra_round_trips_custom_claims() {
607 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
608 let identity = Identity {
609 provider_id: "mock".to_string(),
610 external_id: "user123".to_string(),
611 email: None,
612 username: None,
613 attributes: HashMap::new(),
614 };
615
616 let mut extra = HashMap::new();
617 extra.insert("org_id".to_string(), serde_json::json!("org-7"));
618
619 let token = manager
620 .issue_id_token_with_extra(
621 identity,
622 "client-1",
623 Some("nonce123".to_string()),
624 3600,
625 extra,
626 )
627 .unwrap();
628 let claims = manager.validate_token(&token, None).unwrap();
629
630 assert_eq!(
631 claims.extra.get("org_id"),
632 Some(&serde_json::json!("org-7"))
633 );
634 assert_eq!(
635 claims.extra.get("nonce"),
636 Some(&serde_json::json!("nonce123"))
637 );
638 }
639
640 #[test]
641 fn test_issue_id_token_with_extra_explicit_nonce_wins_over_extra_nonce() {
642 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
646 let identity = Identity {
647 provider_id: "mock".to_string(),
648 external_id: "user123".to_string(),
649 email: None,
650 username: None,
651 attributes: HashMap::new(),
652 };
653
654 let mut extra = HashMap::new();
655 extra.insert(
656 "nonce".to_string(),
657 serde_json::json!("attacker-supplied-nonce"),
658 );
659
660 let token = manager
661 .issue_id_token_with_extra(
662 identity,
663 "client-1",
664 Some("real-nonce".to_string()),
665 3600,
666 extra,
667 )
668 .unwrap();
669 let claims = manager.validate_token(&token, None).unwrap();
670
671 assert_eq!(
672 claims.extra.get("nonce"),
673 Some(&serde_json::json!("real-nonce"))
674 );
675 }
676}
677pub mod jwk;