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    // Authkestra-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        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    /// Issues a token for a user identity.
112    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    /// Issues a token for a user identity, stamping the given `extra` claims
123    /// onto the token in addition to the standard/core claims.
124    ///
125    /// This lets a host application (e.g. a resource server built on top of
126    /// this engine) attach domain-specific claims — such as `api_key_id`,
127    /// `project_id`, or `roles` — so downstream consumers (an API gateway or
128    /// authorization proxy) can read them directly off the token without a
129    /// database round-trip. Keys in `extra` take precedence over any
130    /// same-named field set elsewhere in `extra` by this method; they cannot
131    /// override the top-level standard claims (`sub`, `aud`, `exp`, etc.)
132    /// since those are not part of the flattened map.
133    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    /// Issues an OIDC-conformant ID token.
166    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    /// Issues an OIDC-conformant ID token, stamping the given `extra` claims
177    /// onto the token in addition to the standard/core claims.
178    ///
179    /// `nonce` is a reserved claim key: `extra` is merged into the token
180    /// first, then the explicit `nonce` parameter is applied on top. So if
181    /// `nonce` is `Some(_)`, it always wins over any `"nonce"` entry passed
182    /// in `extra`. If `nonce` is `None`, an `extra["nonce"]` value (if any)
183    /// is left as-is. This preserves OIDC `nonce` semantics — it reflects
184    /// what the client sent in the authorization request — and keeps it from
185    /// being accidentally clobbered by unrelated custom claims.
186    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    /// Issues a machine-to-machine (M2M) token for a client.
225    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    /// Issues a machine-to-machine (M2M) token for a client, stamping the
236    /// given `extra` claims onto the token in addition to the standard/core
237    /// claims. See [`Self::issue_user_token_with_extra`] for the rationale.
238    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 validate_token(
271        &self,
272        token: &str,
273        expected_aud: Option<&str>,
274    ) -> Result<Claims, AuthError> {
275        let mut validation = Validation::new(self.alg);
276        if let Some(aud) = expected_aud {
277            validation.set_audience(&[aud]);
278        } else {
279            validation.validate_aud = false;
280        }
281        if let Some(ref iss) = self.issuer {
282            validation.set_issuer(&[iss]);
283        }
284
285        let token_data = decode::<Claims>(token, &self.decoding_key, &validation)
286            .map_err(|e| AuthError::Token(e.to_string()))?;
287
288        Ok(token_data.claims)
289    }
290}
291
292#[cfg(test)]
293mod tests {
294
295    use super::*;
296    use crate::auth::state::Identity;
297    use std::collections::HashMap;
298
299    #[test]
300    fn test_claims_serialization() {
301        let mut extra = HashMap::new();
302        extra.insert(
303            "custom".to_string(),
304            serde_json::Value::String("value".to_string()),
305        );
306
307        let claims = Claims {
308            iss: Some("issuer".to_string()),
309            sub: "user123".to_string(),
310            aud: Some("audience".to_string()),
311            exp: 1000,
312            iat: 500,
313            nbf: Some(500),
314            jti: Some("jti".to_string()),
315            scope: Some("openid profile".to_string()),
316            identity: Some(Identity {
317                provider_id: "google".to_string(),
318                external_id: "user123".to_string(),
319                email: Some("user@example.com".to_string()),
320                username: Some("user".to_string()),
321                attributes: HashMap::new(),
322            }),
323            extra,
324        };
325
326        let serialized = serde_json::to_string(&claims).unwrap();
327        let deserialized: Claims = serde_json::from_str(&serialized).unwrap();
328
329        assert_eq!(deserialized.iss, claims.iss);
330        assert_eq!(deserialized.sub, claims.sub);
331        assert_eq!(deserialized.extra.get("custom").unwrap(), "value");
332    }
333
334    #[test]
335    fn test_token_manager_issuance() {
336        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
337        let identity = Identity {
338            provider_id: "mock".to_string(),
339            external_id: "user123".to_string(),
340            email: None,
341            username: None,
342            attributes: HashMap::new(),
343        };
344
345        let token = manager
346            .issue_user_token(identity, 3600, None, None)
347            .unwrap();
348        let claims = manager.validate_token(&token, None).unwrap();
349
350        assert_eq!(claims.iss, Some("issuer".to_string()));
351        assert_eq!(claims.sub, "user123");
352        assert!(claims.jti.is_some());
353        assert!(claims.nbf.is_some());
354    }
355
356    #[test]
357    fn test_token_manager_asymmetric_issuance() {
358        let pem = b"-----BEGIN PRIVATE KEY-----
359MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDA5hJIcQ+2rxMz
360VM8ZH5WAmguCr0xmNDAdy0IzzsUeFLG7BebB7izOkU36J4t8t5tUaQwrBMnx2Fvt
361VqJjbdE242UDpvWF/8m9zJ2HR5298cbwT5cGMKLB0HWzDMahugs+Bbh2lCgwyLZk
362Tr3Diwxp5SwFew/Wb+Ke9cNG9Hu5IFH3BCuJ839d9hfqisIeYrBPfb52xxckM37R
3637zSGu/eDP/HZAeLkQuptZJW4A3u7xni14u4qyqXDqsHsYFNgJaxMSAwWgBRY6HNu
364TnvBArTXCiVfL+F73B2L6mdYr64g+QS9nK9v97MlJu/E3mSduz54pren4mpCHc9m
365/S2+VjCZAgMBAAECggEAASC9qQbGnL7XuExRDOIn/m4bWx92ehjo0lCTibhpY3LW
366umbSbpfbhmmuSj3CjW9VZsaM3hBTgSjoTX72lbY/eIUXD7c0memUK5pV4XcEIrQw
367AZlPIye6ckx4I7ZGnKasO8FoAel9dd7DXw36AuBK3LBzJwtzkEFsBc0e3/wixqmG
368UJBbbt/+5ya7CxyjuePaQhKtkLD5R6DpvN2XnCYq5nHJNJdvSVg1pOzsTHYIf+Ee
3692Rz42fGsfFKqeEQCcBFRZaGb/ELeP4c6UZdktZAvmHb1p1fursVZc6X9JXmiJ2OJ
370Kv2H2tMKuysP8L0fXFOMgkH2SVt6rcdHkO6xhlhWsQKBgQDqR8rAJeEE5BFoXA8T
371VVW6CLMlW51x4ey7PEGOaYh39dTG2Q+GZQBZ9G+SZk3f5Y85UCACSyc//4qaz/c3
3720nWsegZ+JPyymmuc79wzIAFFvXB7pL6wyn0Ed1P620kOZTtA8iBcXrsuxL+KP7iu
373MXfWmU1QiZpbndILtyDnY+70uwKBgQDSyCljWkydQCaPU+fiAXLxP8CvcJTSSNQD
374mVUlwJ+OpHnU+Alsi1rBavMgUtLlYbFqzH7NmYrLC8Yadq3ZOwLt0VEK0r8qstAL
3757QCDUD2WNuQjpZupRnXuMUl3iXB96i2gb+VQKGuUAJvVWjdIbYa4+Gu+sBMfcDcX
376dBihDLuEuwKBgAgX4tEwfc2Fc3R/eaXZVNTQaB/qQk4k1+C//CPHUYeTXn5gEUE7
377S//PiesszZPmgkQgmHp7zidP1KH0fT3Yb2g97ut8q54f54fMYXcCrAiUusYKsuu4
378kwkMdkI8QRHWPW3I74VBYIYFFfjYqrCZ1OH8+cbGeiagFRmCggh8U0zxAoGAVW3u
3796Ge22Z0gg8LcHsu7jG/sZq7Ygool8/d3fT+e669Z+ak2GJo6hF4WgClRdMqtn72W
380PzpV+ImjFyK2v26dd0n48MwN0v56N/ss1Av3iiRhPtlmR6tZLNspDZvUzhPVvkrb
381xCs9vtSoVEamVWKe0eVNthGjDoDqs0TInq2MavUCgYB6REavSJs/CLkSS7iimjxZ
382G7g5YQi9/p1lXLOEUDiwEmvRr0XTwzzxUsIc535IXhh/ZUYpthenW+qBBzn85pEC
383TowIqciHu5redqlQ8rITA8/AOY98vaDIhppDg1rfpnHHaZHFbXD/keYAEbhBtbvf
384a0QMqKUcs8+YTy5R5K6qtw==
385-----END PRIVATE KEY-----";
386
387        let manager = TokenManager::new_asymmetric(
388            pem,
389            Some("issuer".to_string()),
390            Some("my-kid-123".to_string()),
391        )
392        .unwrap();
393
394        let identity = Identity {
395            provider_id: "mock".to_string(),
396            external_id: "user123".to_string(),
397            email: None,
398            username: None,
399            attributes: HashMap::new(),
400        };
401
402        let token = manager
403            .issue_user_token(identity, 3600, None, None)
404            .unwrap();
405
406        // Decode directly via jsonwebtoken to prove independent verification
407        let jwk = manager.public_jwk().unwrap();
408        assert_eq!(jwk.kid.as_deref(), Some("my-kid-123"));
409
410        let decoding_key = jwk.to_decoding_key().unwrap();
411        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
412        validation.set_issuer(&["issuer"]);
413
414        let token_data =
415            jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap();
416        assert_eq!(token_data.claims.sub, "user123");
417        assert_eq!(token_data.header.kid.as_deref(), Some("my-kid-123"));
418    }
419
420    #[test]
421    fn test_issue_id_token() {
422        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
423        let identity = Identity {
424            provider_id: "mock".to_string(),
425            external_id: "user123".to_string(),
426            email: None,
427            username: None,
428            attributes: HashMap::new(),
429        };
430
431        let token = manager
432            .issue_id_token(identity, "client-1", Some("nonce123".to_string()), 3600)
433            .unwrap();
434
435        let claims = manager.validate_token(&token, None).unwrap();
436
437        assert_eq!(claims.iss, Some("issuer".to_string()));
438        assert_eq!(claims.sub, "user123");
439        assert_eq!(claims.aud, Some("client-1".to_string()));
440        assert_eq!(claims.extra.get("nonce").unwrap(), "nonce123");
441    }
442    #[test]
443    fn test_token_manager_audience_validation() {
444        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
445        let identity = Identity {
446            provider_id: "mock".to_string(),
447            external_id: "user123".to_string(),
448            email: None,
449            username: None,
450            attributes: HashMap::new(),
451        };
452
453        // Issue token for "client-1"
454        let token = manager
455            .issue_id_token(identity, "client-1", None, 3600)
456            .unwrap();
457
458        // Validate with correct audience
459        let claims = manager.validate_token(&token, Some("client-1")).unwrap();
460        assert_eq!(claims.aud, Some("client-1".to_string()));
461
462        // Validate with incorrect audience (should fail)
463        let err = manager
464            .validate_token(&token, Some("client-2"))
465            .unwrap_err();
466        assert!(err.to_string().contains("InvalidAudience"));
467    }
468
469    #[test]
470    fn test_issue_user_token_with_extra_round_trips_custom_claims() {
471        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
472        let identity = Identity {
473            provider_id: "mock".to_string(),
474            external_id: "user123".to_string(),
475            email: None,
476            username: None,
477            attributes: HashMap::new(),
478        };
479
480        let mut extra = HashMap::new();
481        extra.insert("api_key_id".to_string(), serde_json::json!("key-abc"));
482        extra.insert("project_id".to_string(), serde_json::json!("proj-42"));
483
484        let token = manager
485            .issue_user_token_with_extra(identity, 3600, None, None, extra)
486            .unwrap();
487        let claims = manager.validate_token(&token, None).unwrap();
488
489        assert_eq!(
490            claims.extra.get("api_key_id"),
491            Some(&serde_json::json!("key-abc"))
492        );
493        assert_eq!(
494            claims.extra.get("project_id"),
495            Some(&serde_json::json!("proj-42"))
496        );
497    }
498
499    #[test]
500    fn test_issue_user_token_wrapper_still_has_empty_extra() {
501        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
502        let identity = Identity {
503            provider_id: "mock".to_string(),
504            external_id: "user123".to_string(),
505            email: None,
506            username: None,
507            attributes: HashMap::new(),
508        };
509
510        let token = manager
511            .issue_user_token(identity, 3600, None, None)
512            .unwrap();
513        let claims = manager.validate_token(&token, None).unwrap();
514
515        assert!(claims.extra.is_empty());
516    }
517
518    #[test]
519    fn test_issue_client_token_with_extra_round_trips_custom_claims() {
520        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
521
522        let mut extra = HashMap::new();
523        extra.insert("roles".to_string(), serde_json::json!(["admin", "billing"]));
524
525        let token = manager
526            .issue_client_token_with_extra("client-1", 3600, None, None, extra)
527            .unwrap();
528        let claims = manager.validate_token(&token, None).unwrap();
529
530        assert_eq!(
531            claims.extra.get("roles"),
532            Some(&serde_json::json!(["admin", "billing"]))
533        );
534    }
535
536    #[test]
537    fn test_issue_client_token_wrapper_still_has_empty_extra() {
538        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
539
540        let token = manager
541            .issue_client_token("client-1", 3600, None, None)
542            .unwrap();
543        let claims = manager.validate_token(&token, None).unwrap();
544
545        assert!(claims.extra.is_empty());
546    }
547
548    #[test]
549    fn test_issue_id_token_with_extra_round_trips_custom_claims() {
550        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
551        let identity = Identity {
552            provider_id: "mock".to_string(),
553            external_id: "user123".to_string(),
554            email: None,
555            username: None,
556            attributes: HashMap::new(),
557        };
558
559        let mut extra = HashMap::new();
560        extra.insert("org_id".to_string(), serde_json::json!("org-7"));
561
562        let token = manager
563            .issue_id_token_with_extra(
564                identity,
565                "client-1",
566                Some("nonce123".to_string()),
567                3600,
568                extra,
569            )
570            .unwrap();
571        let claims = manager.validate_token(&token, None).unwrap();
572
573        assert_eq!(
574            claims.extra.get("org_id"),
575            Some(&serde_json::json!("org-7"))
576        );
577        assert_eq!(
578            claims.extra.get("nonce"),
579            Some(&serde_json::json!("nonce123"))
580        );
581    }
582
583    #[test]
584    fn test_issue_id_token_with_extra_explicit_nonce_wins_over_extra_nonce() {
585        // Documents the precedence chosen in `issue_id_token_with_extra`:
586        // the explicit `nonce` parameter always overrides a `"nonce"` entry
587        // supplied via `extra`.
588        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
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 mut extra = HashMap::new();
598        extra.insert(
599            "nonce".to_string(),
600            serde_json::json!("attacker-supplied-nonce"),
601        );
602
603        let token = manager
604            .issue_id_token_with_extra(
605                identity,
606                "client-1",
607                Some("real-nonce".to_string()),
608                3600,
609                extra,
610            )
611            .unwrap();
612        let claims = manager.validate_token(&token, None).unwrap();
613
614        assert_eq!(
615            claims.extra.get("nonce"),
616            Some(&serde_json::json!("real-nonce"))
617        );
618    }
619}
620pub mod jwk;
authkestra_engine/token/mod.rs

authkestra_engine/token/
mod.rs
