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, PartialEq, Eq)]
18#[serde(untagged)]
19pub enum Audience {
20 Single(String),
21 Multiple(Vec<String>),
22}
23
24impl Audience {
25 pub fn contains(&self, value: &str) -> bool {
29 match self {
30 Audience::Single(s) => s == value,
31 Audience::Multiple(values) => values.iter().any(|s| s == value),
32 }
33 }
34}
35
36impl From<String> for Audience {
37 fn from(value: String) -> Self {
38 Audience::Single(value)
39 }
40}
41
42impl From<&str> for Audience {
43 fn from(value: &str) -> Self {
44 Audience::Single(value.to_string())
45 }
46}
47
48fn take_jti(extra: &mut HashMap<String, serde_json::Value>) -> String {
65 match extra.remove("jti") {
66 Some(serde_json::Value::String(jti)) => jti,
67 _ => uuid::Uuid::new_v4().to_string(),
68 }
69}
70
71#[derive(Debug, Serialize, Deserialize, Clone)]
72pub struct Claims {
73 pub iss: Option<String>,
75 pub sub: String,
76 pub aud: Option<Audience>,
77 pub exp: usize,
78 pub iat: usize,
79 pub nbf: Option<usize>,
80 pub jti: Option<String>,
81
82 pub scope: Option<String>,
84 #[serde(skip_serializing_if = "Option::is_none")]
87 pub identity: Option<Identity>,
88
89 #[serde(flatten)]
91 pub extra: HashMap<String, serde_json::Value>,
92}
93
94#[derive(Clone)]
95pub struct TokenManager {
96 encoding_key: EncodingKey,
97 decoding_key: DecodingKey,
98 issuer: Option<String>,
99 kid: Option<String>,
100 alg: Algorithm,
101 public_jwk: Option<crate::token::jwk::Jwk>,
102}
103
104impl TokenManager {
105 pub fn new(secret: &[u8], issuer: Option<String>) -> Self {
107 Self {
108 encoding_key: EncodingKey::from_secret(secret),
109 decoding_key: DecodingKey::from_secret(secret),
110 issuer,
111 kid: None,
112 alg: Algorithm::HS256,
113 public_jwk: None,
114 }
115 }
116
117 pub fn new_asymmetric(
122 private_key_pem: &[u8],
123 issuer: Option<String>,
124 kid: Option<String>,
125 ) -> Result<Self, AuthError> {
126 let encoding_key = EncodingKey::from_rsa_pem(private_key_pem)
127 .map_err(|e| AuthError::Token(e.to_string()))?;
128
129 let pem_str = std::str::from_utf8(private_key_pem)
130 .map_err(|_| AuthError::Token("Invalid PEM UTF-8".into()))?;
131
132 use rsa::pkcs1::DecodeRsaPrivateKey;
133 use rsa::pkcs8::DecodePrivateKey;
134 let rsa_key = rsa::RsaPrivateKey::from_pkcs8_pem(pem_str)
135 .or_else(|_| rsa::RsaPrivateKey::from_pkcs1_pem(pem_str))
136 .map_err(|e| AuthError::Token(format!("Failed to parse RSA key: {}", e)))?;
137
138 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
139 use rsa::traits::PublicKeyParts;
140
141 let n = URL_SAFE_NO_PAD.encode(rsa_key.n().to_bytes_be());
142 let e = URL_SAFE_NO_PAD.encode(rsa_key.e().to_bytes_be());
143
144 let kid_val = kid.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
145
146 let jwk = crate::token::jwk::Jwk {
147 kid: Some(kid_val.clone()),
148 kty: "RSA".to_string(),
149 alg: Some("RS256".to_string()),
150 n: Some(n),
151 e: Some(e),
152 crv: None,
153 x: None,
154 };
155
156 let decoding_key = jwk.to_decoding_key()?;
161
162 Ok(Self {
163 encoding_key,
164 decoding_key,
165 issuer,
166 kid: Some(kid_val),
167 alg: Algorithm::RS256,
168 public_jwk: Some(jwk),
169 })
170 }
171
172 pub fn new_ed25519(
184 private_key_pem: &[u8],
185 issuer: Option<String>,
186 kid: Option<String>,
187 ) -> Result<Self, AuthError> {
188 let encoding_key = EncodingKey::from_ed_pem(private_key_pem)
189 .map_err(|e| AuthError::Token(e.to_string()))?;
190
191 let pem_str = std::str::from_utf8(private_key_pem)
192 .map_err(|_| AuthError::Token("Invalid PEM UTF-8".into()))?;
193
194 use ed25519_dalek::pkcs8::DecodePrivateKey;
195 let signing_key = ed25519_dalek::SigningKey::from_pkcs8_pem(pem_str)
196 .map_err(|e| AuthError::Token(format!("Failed to parse Ed25519 key: {}", e)))?;
197
198 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
199 let x = URL_SAFE_NO_PAD.encode(signing_key.verifying_key().to_bytes());
200
201 let kid_val = kid.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
202
203 let jwk = crate::token::jwk::Jwk {
204 kid: Some(kid_val.clone()),
205 kty: "OKP".to_string(),
206 alg: Some("EdDSA".to_string()),
207 n: None,
208 e: None,
209 crv: Some("Ed25519".to_string()),
210 x: Some(x),
211 };
212
213 let decoding_key = jwk.to_decoding_key()?;
219
220 Ok(Self {
221 encoding_key,
222 decoding_key,
223 issuer,
224 kid: Some(kid_val),
225 alg: Algorithm::EdDSA,
226 public_jwk: Some(jwk),
227 })
228 }
229
230 pub fn public_jwk(&self) -> Option<crate::token::jwk::Jwk> {
231 self.public_jwk.clone()
232 }
233
234 pub fn with_issuer(mut self, issuer: String) -> Self {
235 self.issuer = Some(issuer);
236 self
237 }
238
239 pub fn issue_user_token(
241 &self,
242 identity: Identity,
243 expires_in_secs: u64,
244 scope: Option<String>,
245 aud: Option<String>,
246 ) -> Result<String, AuthError> {
247 self.issue_user_token_with_extra(identity, expires_in_secs, scope, aud, HashMap::new())
248 }
249
250 pub fn issue_user_token_with_extra(
268 &self,
269 identity: Identity,
270 expires_in_secs: u64,
271 scope: Option<String>,
272 aud: Option<String>,
273 mut extra: HashMap<String, serde_json::Value>,
274 ) -> Result<String, AuthError> {
275 let now = chrono::Utc::now().timestamp() as usize;
276 let expiration = now + expires_in_secs as usize;
277 let jti = take_jti(&mut extra);
278
279 let claims = Claims {
280 iss: self.issuer.clone(),
281 sub: identity.external_id.clone(),
282 aud: aud.map(Audience::from),
283 exp: expiration,
284 iat: now,
285 nbf: Some(now),
286 jti: Some(jti),
287 scope,
288 identity: Some(identity),
289 extra,
290 };
291
292 let mut header = Header::new(self.alg);
293 if let Some(ref kid) = self.kid {
294 header.kid = Some(kid.clone());
295 }
296
297 encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
298 }
299
300 pub fn issue_id_token(
302 &self,
303 identity: Identity,
304 client_id: &str,
305 nonce: Option<String>,
306 expires_in_secs: u64,
307 ) -> Result<String, AuthError> {
308 self.issue_id_token_with_extra(identity, client_id, nonce, expires_in_secs, HashMap::new())
309 }
310
311 pub fn issue_id_token_with_extra(
327 &self,
328 identity: Identity,
329 client_id: &str,
330 nonce: Option<String>,
331 expires_in_secs: u64,
332 mut extra: HashMap<String, serde_json::Value>,
333 ) -> Result<String, AuthError> {
334 let now = chrono::Utc::now().timestamp() as usize;
335 let expiration = now + expires_in_secs as usize;
336 let jti = take_jti(&mut extra);
337
338 let mut claims = Claims {
339 iss: self.issuer.clone(),
340 sub: identity.external_id.clone(),
341 aud: Some(Audience::from(client_id)),
342 exp: expiration,
343 iat: now,
344 nbf: Some(now),
345 jti: Some(jti),
346 scope: None,
347 identity: Some(identity),
348 extra,
349 };
350
351 if let Some(n) = nonce {
352 claims
353 .extra
354 .insert("nonce".to_string(), serde_json::Value::String(n));
355 }
356
357 let mut header = Header::new(self.alg);
358 if let Some(ref kid) = self.kid {
359 header.kid = Some(kid.clone());
360 }
361
362 encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
363 }
364
365 pub fn issue_client_token(
367 &self,
368 client_id: &str,
369 expires_in_secs: u64,
370 scope: Option<String>,
371 aud: Option<String>,
372 ) -> Result<String, AuthError> {
373 self.issue_client_token_with_extra(client_id, expires_in_secs, scope, aud, HashMap::new())
374 }
375
376 pub fn issue_client_token_with_extra(
381 &self,
382 client_id: &str,
383 expires_in_secs: u64,
384 scope: Option<String>,
385 aud: Option<String>,
386 mut extra: HashMap<String, serde_json::Value>,
387 ) -> Result<String, AuthError> {
388 let now = chrono::Utc::now().timestamp() as usize;
389 let expiration = now + expires_in_secs as usize;
390 let jti = take_jti(&mut extra);
391
392 let claims = Claims {
393 iss: self.issuer.clone(),
394 sub: client_id.to_string(),
395 aud: aud.map(Audience::from),
396 exp: expiration,
397 iat: now,
398 nbf: Some(now),
399 jti: Some(jti),
400 scope,
401 identity: None,
402 extra,
403 };
404
405 let mut header = Header::new(self.alg);
406 if let Some(ref kid) = self.kid {
407 header.kid = Some(kid.clone());
408 }
409
410 encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
411 }
412
413 pub fn issue_custom_token(
422 &self,
423 sub: String,
424 expires_in_secs: u64,
425 typ: &str,
426 mut extra: HashMap<String, serde_json::Value>,
427 ) -> Result<String, AuthError> {
428 let now = chrono::Utc::now().timestamp() as usize;
429 let jti = take_jti(&mut extra);
430 let claims = Claims {
431 iss: self.issuer.clone(),
432 sub,
433 aud: None,
434 exp: now + expires_in_secs as usize,
435 iat: now,
436 nbf: Some(now),
437 jti: Some(jti),
438 scope: None,
439 identity: None,
440 extra,
441 };
442
443 let mut header = Header::new(self.alg);
444 header.typ = Some(typ.to_string());
445 if let Some(ref kid) = self.kid {
446 header.kid = Some(kid.clone());
447 }
448
449 encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
450 }
451
452 pub fn validate_token(
453 &self,
454 token: &str,
455 expected_aud: Option<&str>,
456 ) -> Result<Claims, AuthError> {
457 let mut validation = Validation::new(self.alg);
458 if let Some(aud) = expected_aud {
459 validation.set_audience(&[aud]);
460 } else {
461 validation.validate_aud = false;
462 }
463 if let Some(ref iss) = self.issuer {
464 validation.set_issuer(&[iss]);
465 }
466
467 let token_data = decode::<Claims>(token, &self.decoding_key, &validation)
468 .map_err(|e| AuthError::Token(e.to_string()))?;
469
470 Ok(token_data.claims)
471 }
472}
473
474#[cfg(test)]
475mod tests {
476
477 use super::*;
478 use crate::auth::state::Identity;
479 use std::collections::HashMap;
480
481 const TEST_RSA_PRIVATE_KEY_PEM: &[u8] = b"-----BEGIN PRIVATE KEY-----
483MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDA5hJIcQ+2rxMz
484VM8ZH5WAmguCr0xmNDAdy0IzzsUeFLG7BebB7izOkU36J4t8t5tUaQwrBMnx2Fvt
485VqJjbdE242UDpvWF/8m9zJ2HR5298cbwT5cGMKLB0HWzDMahugs+Bbh2lCgwyLZk
486Tr3Diwxp5SwFew/Wb+Ke9cNG9Hu5IFH3BCuJ839d9hfqisIeYrBPfb52xxckM37R
4877zSGu/eDP/HZAeLkQuptZJW4A3u7xni14u4qyqXDqsHsYFNgJaxMSAwWgBRY6HNu
488TnvBArTXCiVfL+F73B2L6mdYr64g+QS9nK9v97MlJu/E3mSduz54pren4mpCHc9m
489/S2+VjCZAgMBAAECggEAASC9qQbGnL7XuExRDOIn/m4bWx92ehjo0lCTibhpY3LW
490umbSbpfbhmmuSj3CjW9VZsaM3hBTgSjoTX72lbY/eIUXD7c0memUK5pV4XcEIrQw
491AZlPIye6ckx4I7ZGnKasO8FoAel9dd7DXw36AuBK3LBzJwtzkEFsBc0e3/wixqmG
492UJBbbt/+5ya7CxyjuePaQhKtkLD5R6DpvN2XnCYq5nHJNJdvSVg1pOzsTHYIf+Ee
4932Rz42fGsfFKqeEQCcBFRZaGb/ELeP4c6UZdktZAvmHb1p1fursVZc6X9JXmiJ2OJ
494Kv2H2tMKuysP8L0fXFOMgkH2SVt6rcdHkO6xhlhWsQKBgQDqR8rAJeEE5BFoXA8T
495VVW6CLMlW51x4ey7PEGOaYh39dTG2Q+GZQBZ9G+SZk3f5Y85UCACSyc//4qaz/c3
4960nWsegZ+JPyymmuc79wzIAFFvXB7pL6wyn0Ed1P620kOZTtA8iBcXrsuxL+KP7iu
497MXfWmU1QiZpbndILtyDnY+70uwKBgQDSyCljWkydQCaPU+fiAXLxP8CvcJTSSNQD
498mVUlwJ+OpHnU+Alsi1rBavMgUtLlYbFqzH7NmYrLC8Yadq3ZOwLt0VEK0r8qstAL
4997QCDUD2WNuQjpZupRnXuMUl3iXB96i2gb+VQKGuUAJvVWjdIbYa4+Gu+sBMfcDcX
500dBihDLuEuwKBgAgX4tEwfc2Fc3R/eaXZVNTQaB/qQk4k1+C//CPHUYeTXn5gEUE7
501S//PiesszZPmgkQgmHp7zidP1KH0fT3Yb2g97ut8q54f54fMYXcCrAiUusYKsuu4
502kwkMdkI8QRHWPW3I74VBYIYFFfjYqrCZ1OH8+cbGeiagFRmCggh8U0zxAoGAVW3u
5036Ge22Z0gg8LcHsu7jG/sZq7Ygool8/d3fT+e669Z+ak2GJo6hF4WgClRdMqtn72W
504PzpV+ImjFyK2v26dd0n48MwN0v56N/ss1Av3iiRhPtlmR6tZLNspDZvUzhPVvkrb
505xCs9vtSoVEamVWKe0eVNthGjDoDqs0TInq2MavUCgYB6REavSJs/CLkSS7iimjxZ
506G7g5YQi9/p1lXLOEUDiwEmvRr0XTwzzxUsIc535IXhh/ZUYpthenW+qBBzn85pEC
507TowIqciHu5redqlQ8rITA8/AOY98vaDIhppDg1rfpnHHaZHFbXD/keYAEbhBtbvf
508a0QMqKUcs8+YTy5R5K6qtw==
509-----END PRIVATE KEY-----";
510
511 const TEST_ED25519_PRIVATE_KEY_PEM: &[u8] = b"-----BEGIN PRIVATE KEY-----
514MC4CAQAwBQYDK2VwBCIEIKIPR2jojpdobYr1M/pjIRuMONpZGYQ+y5yxSqKX9T9/
515-----END PRIVATE KEY-----";
516
517 const TEST_ED25519_PRIVATE_KEY_PEM_B: &[u8] = b"-----BEGIN PRIVATE KEY-----
520MC4CAQAwBQYDK2VwBCIEIPlsnSfvh53rJ+Tlbo8e7cgq2mIkWQ1NCM5paVeinUh8
521-----END PRIVATE KEY-----";
522
523 #[test]
524 fn test_claims_serialization() {
525 let mut extra = HashMap::new();
526 extra.insert(
527 "custom".to_string(),
528 serde_json::Value::String("value".to_string()),
529 );
530
531 let claims = Claims {
532 iss: Some("issuer".to_string()),
533 sub: "user123".to_string(),
534 aud: Some(Audience::from("audience")),
535 exp: 1000,
536 iat: 500,
537 nbf: Some(500),
538 jti: Some("jti".to_string()),
539 scope: Some("openid profile".to_string()),
540 identity: Some(Identity {
541 provider_id: "google".to_string(),
542 external_id: "user123".to_string(),
543 email: Some("user@example.com".to_string()),
544 username: Some("user".to_string()),
545 attributes: HashMap::new(),
546 }),
547 extra,
548 };
549
550 let serialized = serde_json::to_string(&claims).unwrap();
551 let deserialized: Claims = serde_json::from_str(&serialized).unwrap();
552
553 assert_eq!(deserialized.iss, claims.iss);
554 assert_eq!(deserialized.sub, claims.sub);
555 assert_eq!(deserialized.extra.get("custom").unwrap(), "value");
556 }
557
558 #[test]
559 fn test_token_manager_issuance() {
560 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
561 let identity = Identity {
562 provider_id: "mock".to_string(),
563 external_id: "user123".to_string(),
564 email: None,
565 username: None,
566 attributes: HashMap::new(),
567 };
568
569 let token = manager
570 .issue_user_token(identity, 3600, None, None)
571 .unwrap();
572 let claims = manager.validate_token(&token, None).unwrap();
573
574 assert_eq!(claims.iss, Some("issuer".to_string()));
575 assert_eq!(claims.sub, "user123");
576 assert!(claims.jti.is_some());
577 assert!(claims.nbf.is_some());
578 }
579
580 #[test]
581 fn test_token_manager_asymmetric_issuance() {
582 let manager = TokenManager::new_asymmetric(
583 TEST_RSA_PRIVATE_KEY_PEM,
584 Some("issuer".to_string()),
585 Some("my-kid-123".to_string()),
586 )
587 .unwrap();
588
589 let identity = Identity {
590 provider_id: "mock".to_string(),
591 external_id: "user123".to_string(),
592 email: None,
593 username: None,
594 attributes: HashMap::new(),
595 };
596
597 let token = manager
598 .issue_user_token(identity, 3600, None, None)
599 .unwrap();
600
601 let jwk = manager.public_jwk().unwrap();
603 assert_eq!(jwk.kid.as_deref(), Some("my-kid-123"));
604
605 let decoding_key = jwk.to_decoding_key().unwrap();
606 let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
607 validation.set_issuer(&["issuer"]);
608
609 let token_data =
610 jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap();
611 assert_eq!(token_data.claims.sub, "user123");
612 assert_eq!(token_data.header.kid.as_deref(), Some("my-kid-123"));
613 }
614
615 #[test]
624 fn test_asymmetric_manager_validates_its_own_tokens() {
625 let manager = TokenManager::new_asymmetric(
626 TEST_RSA_PRIVATE_KEY_PEM,
627 Some("issuer".to_string()),
628 Some("my-kid-123".to_string()),
629 )
630 .unwrap();
631
632 let identity = Identity {
633 provider_id: "mock".to_string(),
634 external_id: "user123".to_string(),
635 email: None,
636 username: None,
637 attributes: HashMap::new(),
638 };
639
640 let token = manager
641 .issue_user_token(identity, 3600, None, None)
642 .unwrap();
643
644 let claims = manager.validate_token(&token, None).unwrap();
645 assert_eq!(claims.sub, "user123");
646 assert_eq!(claims.iss, Some("issuer".to_string()));
647
648 let attestation = manager
650 .issue_custom_token(
651 "device-1".to_string(),
652 60,
653 "webank-attest+jws",
654 HashMap::new(),
655 )
656 .unwrap();
657
658 let attest_claims = manager.validate_token(&attestation, None).unwrap();
659 assert_eq!(attest_claims.sub, "device-1");
660 }
661
662 #[test]
663 fn test_issue_id_token() {
664 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
665 let identity = Identity {
666 provider_id: "mock".to_string(),
667 external_id: "user123".to_string(),
668 email: None,
669 username: None,
670 attributes: HashMap::new(),
671 };
672
673 let token = manager
674 .issue_id_token(identity, "client-1", Some("nonce123".to_string()), 3600)
675 .unwrap();
676
677 let claims = manager.validate_token(&token, None).unwrap();
678
679 assert_eq!(claims.iss, Some("issuer".to_string()));
680 assert_eq!(claims.sub, "user123");
681 assert_eq!(claims.aud, Some(Audience::from("client-1")));
682 assert_eq!(claims.extra.get("nonce").unwrap(), "nonce123");
683 }
684 #[test]
685 fn test_token_manager_audience_validation() {
686 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
687 let identity = Identity {
688 provider_id: "mock".to_string(),
689 external_id: "user123".to_string(),
690 email: None,
691 username: None,
692 attributes: HashMap::new(),
693 };
694
695 let token = manager
697 .issue_id_token(identity, "client-1", None, 3600)
698 .unwrap();
699
700 let claims = manager.validate_token(&token, Some("client-1")).unwrap();
702 assert_eq!(claims.aud, Some(Audience::from("client-1")));
703
704 let err = manager
706 .validate_token(&token, Some("client-2"))
707 .unwrap_err();
708 assert!(err.to_string().contains("InvalidAudience"));
709 }
710
711 #[test]
712 fn test_issue_custom_token_sets_typ_and_no_aud() {
713 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
714
715 let mut extra = HashMap::new();
716 extra.insert("cnf".to_string(), serde_json::json!({"jkt": "abc"}));
717
718 let token = manager
719 .issue_custom_token("device-1".to_string(), 60, "webank-attest+jws", extra)
720 .unwrap();
721
722 let header = jsonwebtoken::decode_header(&token).unwrap();
723 assert_eq!(header.typ.as_deref(), Some("webank-attest+jws"));
724
725 let claims = manager.validate_token(&token, None).unwrap();
726 assert_eq!(claims.sub, "device-1");
727 assert_eq!(claims.aud, None);
728 assert_eq!(claims.extra.get("cnf").unwrap()["jkt"], "abc");
729 }
730
731 #[test]
735 fn test_issue_custom_token_jti_override_round_trips() {
736 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
737
738 let mut extra = HashMap::new();
739 extra.insert(
740 "jti".to_string(),
741 serde_json::json!("caller-supplied-cuid2"),
742 );
743
744 let token = manager
745 .issue_custom_token("device-1".to_string(), 60, "webank-attest+jws", extra)
746 .unwrap();
747 let claims = manager
748 .validate_token(&token, None)
749 .expect("TokenManager must be able to decode the token it just issued");
750
751 assert_eq!(claims.jti, Some("caller-supplied-cuid2".to_string()));
752 assert!(!claims.extra.contains_key("jti"));
753 }
754
755 #[test]
756 fn test_issue_user_token_with_extra_round_trips_custom_claims() {
757 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
758 let identity = Identity {
759 provider_id: "mock".to_string(),
760 external_id: "user123".to_string(),
761 email: None,
762 username: None,
763 attributes: HashMap::new(),
764 };
765
766 let mut extra = HashMap::new();
767 extra.insert("api_key_id".to_string(), serde_json::json!("key-abc"));
768 extra.insert("project_id".to_string(), serde_json::json!("proj-42"));
769
770 let token = manager
771 .issue_user_token_with_extra(identity, 3600, None, None, extra)
772 .unwrap();
773 let claims = manager.validate_token(&token, None).unwrap();
774
775 assert_eq!(
776 claims.extra.get("api_key_id"),
777 Some(&serde_json::json!("key-abc"))
778 );
779 assert_eq!(
780 claims.extra.get("project_id"),
781 Some(&serde_json::json!("proj-42"))
782 );
783 }
784
785 #[test]
786 fn test_issue_user_token_wrapper_still_has_empty_extra() {
787 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
788 let identity = Identity {
789 provider_id: "mock".to_string(),
790 external_id: "user123".to_string(),
791 email: None,
792 username: None,
793 attributes: HashMap::new(),
794 };
795
796 let token = manager
797 .issue_user_token(identity, 3600, None, None)
798 .unwrap();
799 let claims = manager.validate_token(&token, None).unwrap();
800
801 assert!(claims.extra.is_empty());
802 }
803
804 #[test]
812 fn test_issue_user_token_with_extra_jti_override_no_duplicate_key() {
813 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
814 let identity = Identity {
815 provider_id: "mock".to_string(),
816 external_id: "user123".to_string(),
817 email: None,
818 username: None,
819 attributes: HashMap::new(),
820 };
821
822 let mut extra = HashMap::new();
823 extra.insert(
824 "jti".to_string(),
825 serde_json::json!("caller-supplied-cuid2"),
826 );
827
828 let token = manager
829 .issue_user_token_with_extra(identity, 3600, None, None, extra)
830 .unwrap();
831
832 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
833 let payload_b64 = token.split('.').nth(1).unwrap();
834 let payload_json = String::from_utf8(URL_SAFE_NO_PAD.decode(payload_b64).unwrap()).unwrap();
835 assert_eq!(
836 payload_json.matches("\"jti\"").count(),
837 1,
838 "wire payload must contain exactly one jti key, got: {payload_json}"
839 );
840
841 let claims = manager
842 .validate_token(&token, None)
843 .expect("TokenManager must be able to decode the token it just issued");
844 assert_eq!(claims.jti, Some("caller-supplied-cuid2".to_string()));
845 }
846
847 #[test]
850 fn test_issue_user_token_with_extra_default_jti_is_still_generated_uuid() {
851 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
852 let identity = Identity {
853 provider_id: "mock".to_string(),
854 external_id: "user123".to_string(),
855 email: None,
856 username: None,
857 attributes: HashMap::new(),
858 };
859
860 let token = manager
861 .issue_user_token_with_extra(identity, 3600, None, None, HashMap::new())
862 .unwrap();
863 let claims = manager.validate_token(&token, None).unwrap();
864
865 let jti = claims.jti.expect("jti must still be generated by default");
866 assert!(
867 uuid::Uuid::parse_str(&jti).is_ok(),
868 "default-path jti must still be a UUID, got: {jti}"
869 );
870 }
871
872 #[test]
873 fn test_issue_client_token_with_extra_round_trips_custom_claims() {
874 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
875
876 let mut extra = HashMap::new();
877 extra.insert("roles".to_string(), serde_json::json!(["admin", "billing"]));
878
879 let token = manager
880 .issue_client_token_with_extra("client-1", 3600, None, None, extra)
881 .unwrap();
882 let claims = manager.validate_token(&token, None).unwrap();
883
884 assert_eq!(
885 claims.extra.get("roles"),
886 Some(&serde_json::json!(["admin", "billing"]))
887 );
888 }
889
890 #[test]
891 fn test_issue_client_token_wrapper_still_has_empty_extra() {
892 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
893
894 let token = manager
895 .issue_client_token("client-1", 3600, None, None)
896 .unwrap();
897 let claims = manager.validate_token(&token, None).unwrap();
898
899 assert!(claims.extra.is_empty());
900 }
901
902 #[test]
906 fn test_issue_client_token_with_extra_jti_override_round_trips() {
907 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
908
909 let mut extra = HashMap::new();
910 extra.insert(
911 "jti".to_string(),
912 serde_json::json!("caller-supplied-cuid2"),
913 );
914
915 let token = manager
916 .issue_client_token_with_extra("client-1", 3600, None, None, extra)
917 .unwrap();
918 let claims = manager
919 .validate_token(&token, None)
920 .expect("TokenManager must be able to decode the token it just issued");
921
922 assert_eq!(claims.jti, Some("caller-supplied-cuid2".to_string()));
923 assert!(!claims.extra.contains_key("jti"));
924 }
925
926 #[test]
927 fn test_issue_id_token_with_extra_round_trips_custom_claims() {
928 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
929 let identity = Identity {
930 provider_id: "mock".to_string(),
931 external_id: "user123".to_string(),
932 email: None,
933 username: None,
934 attributes: HashMap::new(),
935 };
936
937 let mut extra = HashMap::new();
938 extra.insert("org_id".to_string(), serde_json::json!("org-7"));
939
940 let token = manager
941 .issue_id_token_with_extra(
942 identity,
943 "client-1",
944 Some("nonce123".to_string()),
945 3600,
946 extra,
947 )
948 .unwrap();
949 let claims = manager.validate_token(&token, None).unwrap();
950
951 assert_eq!(
952 claims.extra.get("org_id"),
953 Some(&serde_json::json!("org-7"))
954 );
955 assert_eq!(
956 claims.extra.get("nonce"),
957 Some(&serde_json::json!("nonce123"))
958 );
959 }
960
961 #[test]
962 fn test_issue_id_token_with_extra_explicit_nonce_wins_over_extra_nonce() {
963 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
967 let identity = Identity {
968 provider_id: "mock".to_string(),
969 external_id: "user123".to_string(),
970 email: None,
971 username: None,
972 attributes: HashMap::new(),
973 };
974
975 let mut extra = HashMap::new();
976 extra.insert(
977 "nonce".to_string(),
978 serde_json::json!("attacker-supplied-nonce"),
979 );
980
981 let token = manager
982 .issue_id_token_with_extra(
983 identity,
984 "client-1",
985 Some("real-nonce".to_string()),
986 3600,
987 extra,
988 )
989 .unwrap();
990 let claims = manager.validate_token(&token, None).unwrap();
991
992 assert_eq!(
993 claims.extra.get("nonce"),
994 Some(&serde_json::json!("real-nonce"))
995 );
996 }
997
998 #[test]
1002 fn test_issue_id_token_with_extra_jti_override_round_trips() {
1003 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
1004 let identity = Identity {
1005 provider_id: "mock".to_string(),
1006 external_id: "user123".to_string(),
1007 email: None,
1008 username: None,
1009 attributes: HashMap::new(),
1010 };
1011
1012 let mut extra = HashMap::new();
1013 extra.insert(
1014 "jti".to_string(),
1015 serde_json::json!("caller-supplied-cuid2"),
1016 );
1017
1018 let token = manager
1019 .issue_id_token_with_extra(identity, "client-1", None, 3600, extra)
1020 .unwrap();
1021 let claims = manager
1022 .validate_token(&token, None)
1023 .expect("TokenManager must be able to decode the token it just issued");
1024
1025 assert_eq!(claims.jti, Some("caller-supplied-cuid2".to_string()));
1026 assert!(!claims.extra.contains_key("jti"));
1027 }
1028
1029 #[test]
1030 fn test_token_manager_ed25519_issuance() {
1031 let manager = TokenManager::new_ed25519(
1032 TEST_ED25519_PRIVATE_KEY_PEM,
1033 Some("issuer".to_string()),
1034 Some("my-ed25519-kid".to_string()),
1035 )
1036 .unwrap();
1037
1038 let identity = Identity {
1039 provider_id: "mock".to_string(),
1040 external_id: "user123".to_string(),
1041 email: None,
1042 username: None,
1043 attributes: HashMap::new(),
1044 };
1045
1046 let token = manager
1047 .issue_user_token(identity, 3600, None, None)
1048 .unwrap();
1049
1050 let header = jsonwebtoken::decode_header(&token).unwrap();
1051 assert_eq!(header.alg, Algorithm::EdDSA);
1052 assert_eq!(header.kid.as_deref(), Some("my-ed25519-kid"));
1053 }
1054
1055 #[test]
1062 fn test_ed25519_token_round_trips_via_published_jwk() {
1063 let manager = TokenManager::new_ed25519(
1064 TEST_ED25519_PRIVATE_KEY_PEM,
1065 Some("issuer".to_string()),
1066 Some("my-ed25519-kid".to_string()),
1067 )
1068 .unwrap();
1069
1070 let identity = Identity {
1071 provider_id: "mock".to_string(),
1072 external_id: "user123".to_string(),
1073 email: None,
1074 username: None,
1075 attributes: HashMap::new(),
1076 };
1077
1078 let token = manager
1079 .issue_user_token(identity, 3600, None, None)
1080 .unwrap();
1081
1082 let jwk = manager.public_jwk().unwrap();
1083 assert_eq!(jwk.kid.as_deref(), Some("my-ed25519-kid"));
1084
1085 let decoding_key = jwk.to_decoding_key().unwrap();
1086 let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::EdDSA);
1087 validation.set_issuer(&["issuer"]);
1088
1089 let token_data =
1090 jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap();
1091 assert_eq!(token_data.claims.sub, "user123");
1092 assert_eq!(token_data.header.kid.as_deref(), Some("my-ed25519-kid"));
1093 }
1094
1095 #[test]
1099 fn test_ed25519_manager_validates_its_own_tokens() {
1100 let manager = TokenManager::new_ed25519(
1101 TEST_ED25519_PRIVATE_KEY_PEM,
1102 Some("issuer".to_string()),
1103 Some("my-ed25519-kid".to_string()),
1104 )
1105 .unwrap();
1106
1107 let identity = Identity {
1108 provider_id: "mock".to_string(),
1109 external_id: "user123".to_string(),
1110 email: None,
1111 username: None,
1112 attributes: HashMap::new(),
1113 };
1114
1115 let token = manager
1116 .issue_user_token(identity, 3600, None, None)
1117 .unwrap();
1118
1119 let claims = manager.validate_token(&token, None).unwrap();
1120 assert_eq!(claims.sub, "user123");
1121 assert_eq!(claims.iss, Some("issuer".to_string()));
1122 }
1123
1124 #[test]
1129 fn test_ed25519_public_jwk_is_spec_correct_okp_json() {
1130 let manager = TokenManager::new_ed25519(
1131 TEST_ED25519_PRIVATE_KEY_PEM,
1132 Some("issuer".to_string()),
1133 Some("my-ed25519-kid".to_string()),
1134 )
1135 .unwrap();
1136
1137 let jwk = manager.public_jwk().unwrap();
1138 assert_eq!(jwk.kty, "OKP");
1139 assert_eq!(jwk.alg.as_deref(), Some("EdDSA"));
1140 assert_eq!(jwk.crv.as_deref(), Some("Ed25519"));
1141 assert!(jwk.n.is_none());
1142 assert!(jwk.e.is_none());
1143
1144 let x = jwk.x.as_ref().expect("OKP JWK must have 'x'");
1145 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
1146 let decoded = URL_SAFE_NO_PAD
1147 .decode(x)
1148 .expect("'x' must be valid base64url (unpadded)");
1149 assert_eq!(decoded.len(), 32, "Ed25519 public key must be 32 bytes");
1150
1151 let value = serde_json::to_value(&jwk).unwrap();
1152 let obj = value.as_object().unwrap();
1153 assert_eq!(obj.get("kty").unwrap(), "OKP");
1154 assert_eq!(obj.get("crv").unwrap(), "Ed25519");
1155 assert_eq!(obj.get("alg").unwrap(), "EdDSA");
1156 assert!(obj.contains_key("x"));
1157 assert!(
1158 !obj.contains_key("n"),
1159 "OKP JWK must not serialize the RSA 'n' field, got: {}",
1160 value
1161 );
1162 assert!(
1163 !obj.contains_key("e"),
1164 "OKP JWK must not serialize the RSA 'e' field, got: {}",
1165 value
1166 );
1167 }
1168
1169 #[test]
1172 fn test_rsa_public_jwk_still_omits_okp_fields() {
1173 let manager = TokenManager::new_asymmetric(
1174 TEST_RSA_PRIVATE_KEY_PEM,
1175 Some("issuer".to_string()),
1176 Some("my-kid-123".to_string()),
1177 )
1178 .unwrap();
1179
1180 let jwk = manager.public_jwk().unwrap();
1181 let value = serde_json::to_value(&jwk).unwrap();
1182 let obj = value.as_object().unwrap();
1183 assert_eq!(obj.get("kty").unwrap(), "RSA");
1184 assert!(obj.contains_key("n"));
1185 assert!(obj.contains_key("e"));
1186 assert!(
1187 !obj.contains_key("crv"),
1188 "RSA JWK must not serialize the OKP 'crv' field, got: {}",
1189 value
1190 );
1191 assert!(
1192 !obj.contains_key("x"),
1193 "RSA JWK must not serialize the OKP 'x' field, got: {}",
1194 value
1195 );
1196 }
1197
1198 #[test]
1203 fn test_ed25519_token_rejected_by_wrong_key_jwk() {
1204 let signer = TokenManager::new_ed25519(
1205 TEST_ED25519_PRIVATE_KEY_PEM,
1206 Some("issuer".to_string()),
1207 Some("kid-a".to_string()),
1208 )
1209 .unwrap();
1210 let other = TokenManager::new_ed25519(
1211 TEST_ED25519_PRIVATE_KEY_PEM_B,
1212 Some("issuer".to_string()),
1213 Some("kid-b".to_string()),
1214 )
1215 .unwrap();
1216
1217 let identity = Identity {
1218 provider_id: "mock".to_string(),
1219 external_id: "user123".to_string(),
1220 email: None,
1221 username: None,
1222 attributes: HashMap::new(),
1223 };
1224
1225 let token = signer.issue_user_token(identity, 3600, None, None).unwrap();
1226
1227 assert!(signer.validate_token(&token, None).is_ok());
1229
1230 let err = other.validate_token(&token, None).unwrap_err();
1232 assert!(err.to_string().contains("InvalidSignature"));
1233
1234 let wrong_jwk = other.public_jwk().unwrap();
1237 let decoding_key = wrong_jwk.to_decoding_key().unwrap();
1238 let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::EdDSA);
1239 validation.set_issuer(&["issuer"]);
1240 let err = jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap_err();
1241 assert_eq!(
1242 err.kind(),
1243 &jsonwebtoken::errors::ErrorKind::InvalidSignature
1244 );
1245 }
1246
1247 #[test]
1251 fn test_ed25519_tampered_token_rejected() {
1252 let manager = TokenManager::new_ed25519(
1253 TEST_ED25519_PRIVATE_KEY_PEM,
1254 Some("issuer".to_string()),
1255 Some("my-ed25519-kid".to_string()),
1256 )
1257 .unwrap();
1258
1259 let identity = Identity {
1260 provider_id: "mock".to_string(),
1261 external_id: "user123".to_string(),
1262 email: None,
1263 username: None,
1264 attributes: HashMap::new(),
1265 };
1266
1267 let token = manager
1268 .issue_user_token(identity, 3600, None, None)
1269 .unwrap();
1270
1271 let mut parts: Vec<&str> = token.split('.').collect();
1272 assert_eq!(parts.len(), 3);
1273 let mut payload = parts[1].to_string();
1275 let last = payload.pop().unwrap();
1276 let replacement = if last == 'A' { 'B' } else { 'A' };
1277 payload.push(replacement);
1278 parts[1] = &payload;
1279 let tampered = parts.join(".");
1280
1281 let err = manager.validate_token(&tampered, None).unwrap_err();
1282 assert!(
1283 err.to_string().contains("InvalidSignature") || err.to_string().contains("Json"),
1284 "unexpected error for tampered token: {err}"
1285 );
1286 }
1287
1288 #[test]
1296 fn test_validate_token_accepts_array_aud() {
1297 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
1298
1299 let raw_claims = serde_json::json!({
1300 "iss": "issuer",
1301 "sub": "user123",
1302 "aud": ["client-1", "client-2"],
1303 "exp": (chrono::Utc::now().timestamp() as usize) + 3600,
1304 "iat": chrono::Utc::now().timestamp() as usize,
1305 });
1306 let token = jsonwebtoken::encode(
1307 &jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256),
1308 &raw_claims,
1309 &jsonwebtoken::EncodingKey::from_secret(b"secret"),
1310 )
1311 .unwrap();
1312
1313 let claims = manager
1314 .validate_token(&token, None)
1315 .expect("multi-audience subject token must deserialize");
1316 let aud = claims.aud.expect("aud claim must be present");
1317 assert!(aud.contains("client-1"));
1318 assert!(aud.contains("client-2"));
1319 assert!(!aud.contains("client-3"));
1320 }
1321
1322 #[test]
1327 fn test_validate_token_still_accepts_string_aud() {
1328 let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
1329 let identity = Identity {
1330 provider_id: "mock".to_string(),
1331 external_id: "user123".to_string(),
1332 email: None,
1333 username: None,
1334 attributes: HashMap::new(),
1335 };
1336
1337 let token = manager
1338 .issue_id_token(identity, "client-1", None, 3600)
1339 .unwrap();
1340
1341 let claims = manager.validate_token(&token, None).unwrap();
1342 let aud = claims.aud.expect("aud claim must be present");
1343 assert!(aud.contains("client-1"));
1344 assert!(!aud.contains("client-2"));
1345 }
1346
1347 #[test]
1352 fn test_audience_single_round_trips_as_bare_string() {
1353 let single = Audience::Single("client-1".to_string());
1354 let serialized = serde_json::to_string(&single).unwrap();
1355 assert_eq!(serialized, "\"client-1\"");
1356
1357 let deserialized: Audience = serde_json::from_str(&serialized).unwrap();
1358 assert_eq!(deserialized, single);
1359
1360 let multiple = Audience::Multiple(vec!["client-1".to_string(), "client-2".to_string()]);
1361 let serialized_multi = serde_json::to_string(&multiple).unwrap();
1362 assert_eq!(serialized_multi, "[\"client-1\",\"client-2\"]");
1363
1364 let deserialized_multi: Audience = serde_json::from_str(&serialized_multi).unwrap();
1365 assert_eq!(deserialized_multi, multiple);
1366 }
1367}
1368pub mod jwk;