authkestra_engine/token/
mod.rs

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    // Standard OIDC claims
10    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    // Engine-specific core fields
19    pub scope: Option<String>,
20    /// Optional identity data for user-centric tokens.
21    /// If None, this is likely a machine-to-machine token.
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub identity: Option<Identity>,
24
25    // Isolated custom claims
26    #[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    /// Creates a TokenManager for symmetric signing (HS256).
42    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    /// Creates a TokenManager for asymmetric signing (RS256).
54    /// `private_key_pem` must be a valid RSA private key in PEM format.
55    /// OP/external verification should use this path; internal resource servers
56    /// can continue to use `new` (HS256).
57    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
65        let pem_str = std::str::from_utf8(private_key_pem)
66            .map_err(|_| AuthError::Token("Invalid PEM UTF-8".into()))?;
67
68        use rsa::pkcs1::DecodeRsaPrivateKey;
69        use rsa::pkcs8::DecodePrivateKey;
70        let rsa_key = rsa::RsaPrivateKey::from_pkcs8_pem(pem_str)
71            .or_else(|_| rsa::RsaPrivateKey::from_pkcs1_pem(pem_str))
72            .map_err(|e| AuthError::Token(format!("Failed to parse RSA key: {}", e)))?;
73
74        use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
75        use rsa::traits::PublicKeyParts;
76
77        let n = URL_SAFE_NO_PAD.encode(rsa_key.n().to_bytes_be());
78        let e = URL_SAFE_NO_PAD.encode(rsa_key.e().to_bytes_be());
79
80        let kid_val = kid.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
81
82        let jwk = crate::token::jwk::Jwk {
83            kid: Some(kid_val.clone()),
84            kty: "RSA".to_string(),
85            alg: Some("RS256".to_string()),
86            n: Some(n),
87            e: Some(e),
88        };
89
90        // The decoding key must come from the PUBLIC half. `DecodingKey::from_rsa_pem`
91        // expects a public-key PEM; handed a private one it still constructs, but every
92        // later `validate_token` fails with `InvalidSignature`. Deriving it from the JWK
93        // we just built keeps both halves provably in sync with what `/jwks` publishes.
94        let decoding_key = jwk.to_decoding_key()?;
95
96        Ok(Self {
97            encoding_key,
98            decoding_key,
99            issuer,
100            kid: Some(kid_val),
101            alg: Algorithm::RS256,
102            public_jwk: Some(jwk),
103        })
104    }
105
106    pub fn public_jwk(&self) -> Option<crate::token::jwk::Jwk> {
107        self.public_jwk.clone()
108    }
109
110    pub fn with_issuer(mut self, issuer: String) -> Self {
111        self.issuer = Some(issuer);
112        self
113    }
114
115    /// Issues a token for a user identity.
116    pub fn issue_user_token(
117        &self,
118        identity: Identity,
119        expires_in_secs: u64,
120        scope: Option<String>,
121        aud: Option<String>,
122    ) -> Result<String, AuthError> {
123        self.issue_user_token_with_extra(identity, expires_in_secs, scope, aud, HashMap::new())
124    }
125
126    /// Issues a token for a user identity, stamping the given `extra` claims
127    /// onto the token in addition to the standard/core claims.
128    ///
129    /// This lets a host application (e.g. a resource server built on top of
130    /// this engine) attach domain-specific claims — such as `api_key_id`,
131    /// `project_id`, or `roles` — so downstream consumers (an API gateway or
132    /// authorization proxy) can read them directly off the token without a
133    /// database round-trip. Keys in `extra` take precedence over any
134    /// same-named field set elsewhere in `extra` by this method; they cannot
135    /// override the top-level standard claims (`sub`, `aud`, `exp`, etc.)
136    /// since those are not part of the flattened map.
137    pub fn issue_user_token_with_extra(
138        &self,
139        identity: Identity,
140        expires_in_secs: u64,
141        scope: Option<String>,
142        aud: Option<String>,
143        extra: HashMap<String, serde_json::Value>,
144    ) -> Result<String, AuthError> {
145        let now = chrono::Utc::now().timestamp() as usize;
146        let expiration = now + expires_in_secs as usize;
147
148        let claims = Claims {
149            iss: self.issuer.clone(),
150            sub: identity.external_id.clone(),
151            aud,
152            exp: expiration,
153            iat: now,
154            nbf: Some(now),
155            jti: Some(uuid::Uuid::new_v4().to_string()),
156            scope,
157            identity: Some(identity),
158            extra,
159        };
160
161        let mut header = Header::new(self.alg);
162        if let Some(ref kid) = self.kid {
163            header.kid = Some(kid.clone());
164        }
165
166        encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
167    }
168
169    /// Issues an OIDC-conformant ID token.
170    pub fn issue_id_token(
171        &self,
172        identity: Identity,
173        client_id: &str,
174        nonce: Option<String>,
175        expires_in_secs: u64,
176    ) -> Result<String, AuthError> {
177        self.issue_id_token_with_extra(identity, client_id, nonce, expires_in_secs, HashMap::new())
178    }
179
180    /// Issues an OIDC-conformant ID token, stamping the given `extra` claims
181    /// onto the token in addition to the standard/core claims.
182    ///
183    /// `nonce` is a reserved claim key: `extra` is merged into the token
184    /// first, then the explicit `nonce` parameter is applied on top. So if
185    /// `nonce` is `Some(_)`, it always wins over any `"nonce"` entry passed
186    /// in `extra`. If `nonce` is `None`, an `extra["nonce"]` value (if any)
187    /// is left as-is. This preserves OIDC `nonce` semantics — it reflects
188    /// what the client sent in the authorization request — and keeps it from
189    /// being accidentally clobbered by unrelated custom claims.
190    pub fn issue_id_token_with_extra(
191        &self,
192        identity: Identity,
193        client_id: &str,
194        nonce: Option<String>,
195        expires_in_secs: u64,
196        extra: HashMap<String, serde_json::Value>,
197    ) -> Result<String, AuthError> {
198        let now = chrono::Utc::now().timestamp() as usize;
199        let expiration = now + expires_in_secs as usize;
200
201        let mut claims = Claims {
202            iss: self.issuer.clone(),
203            sub: identity.external_id.clone(),
204            aud: Some(client_id.to_string()),
205            exp: expiration,
206            iat: now,
207            nbf: Some(now),
208            jti: Some(uuid::Uuid::new_v4().to_string()),
209            scope: None,
210            identity: Some(identity),
211            extra,
212        };
213
214        if let Some(n) = nonce {
215            claims
216                .extra
217                .insert("nonce".to_string(), serde_json::Value::String(n));
218        }
219
220        let mut header = Header::new(self.alg);
221        if let Some(ref kid) = self.kid {
222            header.kid = Some(kid.clone());
223        }
224
225        encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
226    }
227
228    /// Issues a machine-to-machine (M2M) token for a client.
229    pub fn issue_client_token(
230        &self,
231        client_id: &str,
232        expires_in_secs: u64,
233        scope: Option<String>,
234        aud: Option<String>,
235    ) -> Result<String, AuthError> {
236        self.issue_client_token_with_extra(client_id, expires_in_secs, scope, aud, HashMap::new())
237    }
238
239    /// Issues a machine-to-machine (M2M) token for a client, stamping the
240    /// given `extra` claims onto the token in addition to the standard/core
241    /// claims. See [`Self::issue_user_token_with_extra`] for the rationale.
242    pub fn issue_client_token_with_extra(
243        &self,
244        client_id: &str,
245        expires_in_secs: u64,
246        scope: Option<String>,
247        aud: Option<String>,
248        extra: HashMap<String, serde_json::Value>,
249    ) -> Result<String, AuthError> {
250        let now = chrono::Utc::now().timestamp() as usize;
251        let expiration = now + expires_in_secs as usize;
252
253        let claims = Claims {
254            iss: self.issuer.clone(),
255            sub: client_id.to_string(),
256            aud,
257            exp: expiration,
258            iat: now,
259            nbf: Some(now),
260            jti: Some(uuid::Uuid::new_v4().to_string()),
261            scope,
262            identity: None,
263            extra,
264        };
265
266        let mut header = Header::new(self.alg);
267        if let Some(ref kid) = self.kid {
268            header.kid = Some(kid.clone());
269        }
270
271        encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
272    }
273
274    /// Issues a token with an explicit `typ` header and no `aud`, for
275    /// callers minting something that is not a standard OIDC ID/access/user
276    /// token and needs its own wire-format `typ` so verifiers can tell it
277    /// apart from those (e.g. `authkestra-op`'s device/service attestations,
278    /// whose contract requires `typ: "webank-attest+jws"` rather than the
279    /// default `"JWT"`). Additive alongside the `issue_*_token*` family
280    /// above; those are unchanged.
281    pub fn issue_custom_token(
282        &self,
283        sub: String,
284        expires_in_secs: u64,
285        typ: &str,
286        extra: HashMap<String, serde_json::Value>,
287    ) -> Result<String, AuthError> {
288        let now = chrono::Utc::now().timestamp() as usize;
289        let claims = Claims {
290            iss: self.issuer.clone(),
291            sub,
292            aud: None,
293            exp: now + expires_in_secs as usize,
294            iat: now,
295            nbf: Some(now),
296            jti: Some(uuid::Uuid::new_v4().to_string()),
297            scope: None,
298            identity: None,
299            extra,
300        };
301
302        let mut header = Header::new(self.alg);
303        header.typ = Some(typ.to_string());
304        if let Some(ref kid) = self.kid {
305            header.kid = Some(kid.clone());
306        }
307
308        encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
309    }
310
311    pub fn validate_token(
312        &self,
313        token: &str,
314        expected_aud: Option<&str>,
315    ) -> Result<Claims, AuthError> {
316        let mut validation = Validation::new(self.alg);
317        if let Some(aud) = expected_aud {
318            validation.set_audience(&[aud]);
319        } else {
320            validation.validate_aud = false;
321        }
322        if let Some(ref iss) = self.issuer {
323            validation.set_issuer(&[iss]);
324        }
325
326        let token_data = decode::<Claims>(token, &self.decoding_key, &validation)
327            .map_err(|e| AuthError::Token(e.to_string()))?;
328
329        Ok(token_data.claims)
330    }
331}
332
333#[cfg(test)]
334mod tests {
335
336    use super::*;
337    use crate::auth::state::Identity;
338    use std::collections::HashMap;
339
340    /// Throwaway RSA-2048 private key, test-only.
341    const TEST_RSA_PRIVATE_KEY_PEM: &[u8] = b"-----BEGIN PRIVATE KEY-----
342MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDA5hJIcQ+2rxMz
343VM8ZH5WAmguCr0xmNDAdy0IzzsUeFLG7BebB7izOkU36J4t8t5tUaQwrBMnx2Fvt
344VqJjbdE242UDpvWF/8m9zJ2HR5298cbwT5cGMKLB0HWzDMahugs+Bbh2lCgwyLZk
345Tr3Diwxp5SwFew/Wb+Ke9cNG9Hu5IFH3BCuJ839d9hfqisIeYrBPfb52xxckM37R
3467zSGu/eDP/HZAeLkQuptZJW4A3u7xni14u4qyqXDqsHsYFNgJaxMSAwWgBRY6HNu
347TnvBArTXCiVfL+F73B2L6mdYr64g+QS9nK9v97MlJu/E3mSduz54pren4mpCHc9m
348/S2+VjCZAgMBAAECggEAASC9qQbGnL7XuExRDOIn/m4bWx92ehjo0lCTibhpY3LW
349umbSbpfbhmmuSj3CjW9VZsaM3hBTgSjoTX72lbY/eIUXD7c0memUK5pV4XcEIrQw
350AZlPIye6ckx4I7ZGnKasO8FoAel9dd7DXw36AuBK3LBzJwtzkEFsBc0e3/wixqmG
351UJBbbt/+5ya7CxyjuePaQhKtkLD5R6DpvN2XnCYq5nHJNJdvSVg1pOzsTHYIf+Ee
3522Rz42fGsfFKqeEQCcBFRZaGb/ELeP4c6UZdktZAvmHb1p1fursVZc6X9JXmiJ2OJ
353Kv2H2tMKuysP8L0fXFOMgkH2SVt6rcdHkO6xhlhWsQKBgQDqR8rAJeEE5BFoXA8T
354VVW6CLMlW51x4ey7PEGOaYh39dTG2Q+GZQBZ9G+SZk3f5Y85UCACSyc//4qaz/c3
3550nWsegZ+JPyymmuc79wzIAFFvXB7pL6wyn0Ed1P620kOZTtA8iBcXrsuxL+KP7iu
356MXfWmU1QiZpbndILtyDnY+70uwKBgQDSyCljWkydQCaPU+fiAXLxP8CvcJTSSNQD
357mVUlwJ+OpHnU+Alsi1rBavMgUtLlYbFqzH7NmYrLC8Yadq3ZOwLt0VEK0r8qstAL
3587QCDUD2WNuQjpZupRnXuMUl3iXB96i2gb+VQKGuUAJvVWjdIbYa4+Gu+sBMfcDcX
359dBihDLuEuwKBgAgX4tEwfc2Fc3R/eaXZVNTQaB/qQk4k1+C//CPHUYeTXn5gEUE7
360S//PiesszZPmgkQgmHp7zidP1KH0fT3Yb2g97ut8q54f54fMYXcCrAiUusYKsuu4
361kwkMdkI8QRHWPW3I74VBYIYFFfjYqrCZ1OH8+cbGeiagFRmCggh8U0zxAoGAVW3u
3626Ge22Z0gg8LcHsu7jG/sZq7Ygool8/d3fT+e669Z+ak2GJo6hF4WgClRdMqtn72W
363PzpV+ImjFyK2v26dd0n48MwN0v56N/ss1Av3iiRhPtlmR6tZLNspDZvUzhPVvkrb
364xCs9vtSoVEamVWKe0eVNthGjDoDqs0TInq2MavUCgYB6REavSJs/CLkSS7iimjxZ
365G7g5YQi9/p1lXLOEUDiwEmvRr0XTwzzxUsIc535IXhh/ZUYpthenW+qBBzn85pEC
366TowIqciHu5redqlQ8rITA8/AOY98vaDIhppDg1rfpnHHaZHFbXD/keYAEbhBtbvf
367a0QMqKUcs8+YTy5R5K6qtw==
368-----END PRIVATE KEY-----";
369
370    #[test]
371    fn test_claims_serialization() {
372        let mut extra = HashMap::new();
373        extra.insert(
374            "custom".to_string(),
375            serde_json::Value::String("value".to_string()),
376        );
377
378        let claims = Claims {
379            iss: Some("issuer".to_string()),
380            sub: "user123".to_string(),
381            aud: Some("audience".to_string()),
382            exp: 1000,
383            iat: 500,
384            nbf: Some(500),
385            jti: Some("jti".to_string()),
386            scope: Some("openid profile".to_string()),
387            identity: Some(Identity {
388                provider_id: "google".to_string(),
389                external_id: "user123".to_string(),
390                email: Some("user@example.com".to_string()),
391                username: Some("user".to_string()),
392                attributes: HashMap::new(),
393            }),
394            extra,
395        };
396
397        let serialized = serde_json::to_string(&claims).unwrap();
398        let deserialized: Claims = serde_json::from_str(&serialized).unwrap();
399
400        assert_eq!(deserialized.iss, claims.iss);
401        assert_eq!(deserialized.sub, claims.sub);
402        assert_eq!(deserialized.extra.get("custom").unwrap(), "value");
403    }
404
405    #[test]
406    fn test_token_manager_issuance() {
407        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
408        let identity = Identity {
409            provider_id: "mock".to_string(),
410            external_id: "user123".to_string(),
411            email: None,
412            username: None,
413            attributes: HashMap::new(),
414        };
415
416        let token = manager
417            .issue_user_token(identity, 3600, None, None)
418            .unwrap();
419        let claims = manager.validate_token(&token, None).unwrap();
420
421        assert_eq!(claims.iss, Some("issuer".to_string()));
422        assert_eq!(claims.sub, "user123");
423        assert!(claims.jti.is_some());
424        assert!(claims.nbf.is_some());
425    }
426
427    #[test]
428    fn test_token_manager_asymmetric_issuance() {
429        let manager = TokenManager::new_asymmetric(
430            TEST_RSA_PRIVATE_KEY_PEM,
431            Some("issuer".to_string()),
432            Some("my-kid-123".to_string()),
433        )
434        .unwrap();
435
436        let identity = Identity {
437            provider_id: "mock".to_string(),
438            external_id: "user123".to_string(),
439            email: None,
440            username: None,
441            attributes: HashMap::new(),
442        };
443
444        let token = manager
445            .issue_user_token(identity, 3600, None, None)
446            .unwrap();
447
448        // Decode directly via jsonwebtoken to prove independent verification
449        let jwk = manager.public_jwk().unwrap();
450        assert_eq!(jwk.kid.as_deref(), Some("my-kid-123"));
451
452        let decoding_key = jwk.to_decoding_key().unwrap();
453        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
454        validation.set_issuer(&["issuer"]);
455
456        let token_data =
457            jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap();
458        assert_eq!(token_data.claims.sub, "user123");
459        assert_eq!(token_data.header.kid.as_deref(), Some("my-kid-123"));
460    }
461
462    /// Regression test for the asymmetric decoding key being derived from the
463    /// private half instead of the public one.
464    ///
465    /// `test_token_manager_asymmetric_issuance` above verifies via the
466    /// published JWK, which exercises `Jwk::to_decoding_key` rather than the
467    /// manager's own `decoding_key` — so it stays green either way. This one
468    /// goes through `validate_token`, which is the path `/reissue` and
469    /// `/userinfo` actually take.
470    #[test]
471    fn test_asymmetric_manager_validates_its_own_tokens() {
472        let manager = TokenManager::new_asymmetric(
473            TEST_RSA_PRIVATE_KEY_PEM,
474            Some("issuer".to_string()),
475            Some("my-kid-123".to_string()),
476        )
477        .unwrap();
478
479        let identity = Identity {
480            provider_id: "mock".to_string(),
481            external_id: "user123".to_string(),
482            email: None,
483            username: None,
484            attributes: HashMap::new(),
485        };
486
487        let token = manager
488            .issue_user_token(identity, 3600, None, None)
489            .unwrap();
490
491        let claims = manager.validate_token(&token, None).unwrap();
492        assert_eq!(claims.sub, "user123");
493        assert_eq!(claims.iss, Some("issuer".to_string()));
494
495        // Same round trip for the attestation shape `/reissue` presents.
496        let attestation = manager
497            .issue_custom_token(
498                "device-1".to_string(),
499                60,
500                "webank-attest+jws",
501                HashMap::new(),
502            )
503            .unwrap();
504
505        let attest_claims = manager.validate_token(&attestation, None).unwrap();
506        assert_eq!(attest_claims.sub, "device-1");
507    }
508
509    #[test]
510    fn test_issue_id_token() {
511        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
512        let identity = Identity {
513            provider_id: "mock".to_string(),
514            external_id: "user123".to_string(),
515            email: None,
516            username: None,
517            attributes: HashMap::new(),
518        };
519
520        let token = manager
521            .issue_id_token(identity, "client-1", Some("nonce123".to_string()), 3600)
522            .unwrap();
523
524        let claims = manager.validate_token(&token, None).unwrap();
525
526        assert_eq!(claims.iss, Some("issuer".to_string()));
527        assert_eq!(claims.sub, "user123");
528        assert_eq!(claims.aud, Some("client-1".to_string()));
529        assert_eq!(claims.extra.get("nonce").unwrap(), "nonce123");
530    }
531    #[test]
532    fn test_token_manager_audience_validation() {
533        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
534        let identity = Identity {
535            provider_id: "mock".to_string(),
536            external_id: "user123".to_string(),
537            email: None,
538            username: None,
539            attributes: HashMap::new(),
540        };
541
542        // Issue token for "client-1"
543        let token = manager
544            .issue_id_token(identity, "client-1", None, 3600)
545            .unwrap();
546
547        // Validate with correct audience
548        let claims = manager.validate_token(&token, Some("client-1")).unwrap();
549        assert_eq!(claims.aud, Some("client-1".to_string()));
550
551        // Validate with incorrect audience (should fail)
552        let err = manager
553            .validate_token(&token, Some("client-2"))
554            .unwrap_err();
555        assert!(err.to_string().contains("InvalidAudience"));
556    }
557
558    #[test]
559    fn test_issue_custom_token_sets_typ_and_no_aud() {
560        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
561
562        let mut extra = HashMap::new();
563        extra.insert("cnf".to_string(), serde_json::json!({"jkt": "abc"}));
564
565        let token = manager
566            .issue_custom_token("device-1".to_string(), 60, "webank-attest+jws", extra)
567            .unwrap();
568
569        let header = jsonwebtoken::decode_header(&token).unwrap();
570        assert_eq!(header.typ.as_deref(), Some("webank-attest+jws"));
571
572        let claims = manager.validate_token(&token, None).unwrap();
573        assert_eq!(claims.sub, "device-1");
574        assert_eq!(claims.aud, None);
575        assert_eq!(claims.extra.get("cnf").unwrap()["jkt"], "abc");
576    }
577
578    #[test]
579    fn test_issue_user_token_with_extra_round_trips_custom_claims() {
580        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
581        let identity = Identity {
582            provider_id: "mock".to_string(),
583            external_id: "user123".to_string(),
584            email: None,
585            username: None,
586            attributes: HashMap::new(),
587        };
588
589        let mut extra = HashMap::new();
590        extra.insert("api_key_id".to_string(), serde_json::json!("key-abc"));
591        extra.insert("project_id".to_string(), serde_json::json!("proj-42"));
592
593        let token = manager
594            .issue_user_token_with_extra(identity, 3600, None, None, extra)
595            .unwrap();
596        let claims = manager.validate_token(&token, None).unwrap();
597
598        assert_eq!(
599            claims.extra.get("api_key_id"),
600            Some(&serde_json::json!("key-abc"))
601        );
602        assert_eq!(
603            claims.extra.get("project_id"),
604            Some(&serde_json::json!("proj-42"))
605        );
606    }
607
608    #[test]
609    fn test_issue_user_token_wrapper_still_has_empty_extra() {
610        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
611        let identity = Identity {
612            provider_id: "mock".to_string(),
613            external_id: "user123".to_string(),
614            email: None,
615            username: None,
616            attributes: HashMap::new(),
617        };
618
619        let token = manager
620            .issue_user_token(identity, 3600, None, None)
621            .unwrap();
622        let claims = manager.validate_token(&token, None).unwrap();
623
624        assert!(claims.extra.is_empty());
625    }
626
627    #[test]
628    fn test_issue_client_token_with_extra_round_trips_custom_claims() {
629        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
630
631        let mut extra = HashMap::new();
632        extra.insert("roles".to_string(), serde_json::json!(["admin", "billing"]));
633
634        let token = manager
635            .issue_client_token_with_extra("client-1", 3600, None, None, extra)
636            .unwrap();
637        let claims = manager.validate_token(&token, None).unwrap();
638
639        assert_eq!(
640            claims.extra.get("roles"),
641            Some(&serde_json::json!(["admin", "billing"]))
642        );
643    }
644
645    #[test]
646    fn test_issue_client_token_wrapper_still_has_empty_extra() {
647        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
648
649        let token = manager
650            .issue_client_token("client-1", 3600, None, None)
651            .unwrap();
652        let claims = manager.validate_token(&token, None).unwrap();
653
654        assert!(claims.extra.is_empty());
655    }
656
657    #[test]
658    fn test_issue_id_token_with_extra_round_trips_custom_claims() {
659        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
660        let identity = Identity {
661            provider_id: "mock".to_string(),
662            external_id: "user123".to_string(),
663            email: None,
664            username: None,
665            attributes: HashMap::new(),
666        };
667
668        let mut extra = HashMap::new();
669        extra.insert("org_id".to_string(), serde_json::json!("org-7"));
670
671        let token = manager
672            .issue_id_token_with_extra(
673                identity,
674                "client-1",
675                Some("nonce123".to_string()),
676                3600,
677                extra,
678            )
679            .unwrap();
680        let claims = manager.validate_token(&token, None).unwrap();
681
682        assert_eq!(
683            claims.extra.get("org_id"),
684            Some(&serde_json::json!("org-7"))
685        );
686        assert_eq!(
687            claims.extra.get("nonce"),
688            Some(&serde_json::json!("nonce123"))
689        );
690    }
691
692    #[test]
693    fn test_issue_id_token_with_extra_explicit_nonce_wins_over_extra_nonce() {
694        // Documents the precedence chosen in `issue_id_token_with_extra`:
695        // the explicit `nonce` parameter always overrides a `"nonce"` entry
696        // supplied via `extra`.
697        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
698        let identity = Identity {
699            provider_id: "mock".to_string(),
700            external_id: "user123".to_string(),
701            email: None,
702            username: None,
703            attributes: HashMap::new(),
704        };
705
706        let mut extra = HashMap::new();
707        extra.insert(
708            "nonce".to_string(),
709            serde_json::json!("attacker-supplied-nonce"),
710        );
711
712        let token = manager
713            .issue_id_token_with_extra(
714                identity,
715                "client-1",
716                Some("real-nonce".to_string()),
717                3600,
718                extra,
719            )
720            .unwrap();
721        let claims = manager.validate_token(&token, None).unwrap();
722
723        assert_eq!(
724            claims.extra.get("nonce"),
725            Some(&serde_json::json!("real-nonce"))
726        );
727    }
728}
729pub mod jwk;
authkestra_engine/token/mod.rs

authkestra_engine/token/
mod.rs
