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            crv: None,
89            x: None,
90        };
91
92        // The decoding key must come from the PUBLIC half. `DecodingKey::from_rsa_pem`
93        // expects a public-key PEM; handed a private one it still constructs, but every
94        // later `validate_token` fails with `InvalidSignature`. Deriving it from the JWK
95        // we just built keeps both halves provably in sync with what `/jwks` publishes.
96        let decoding_key = jwk.to_decoding_key()?;
97
98        Ok(Self {
99            encoding_key,
100            decoding_key,
101            issuer,
102            kid: Some(kid_val),
103            alg: Algorithm::RS256,
104            public_jwk: Some(jwk),
105        })
106    }
107
108    /// Creates a TokenManager for asymmetric signing with Ed25519 (EdDSA).
109    /// `private_key_pem` must be a valid Ed25519 private key in PKCS#8 PEM
110    /// format (`-----BEGIN PRIVATE KEY-----`), e.g. as produced by
111    /// `openssl genpkey -algorithm ed25519`.
112    ///
113    /// Mirrors `new_asymmetric` (RS256): OP/external verification should use
114    /// this path when downstream resource servers require EdDSA-signed
115    /// tokens; internal resource servers can continue to use `new` (HS256).
116    /// The published JWK (`public_jwk`) is the OKP shape from RFC 8037, so
117    /// pair this with #188 (`Jwk`'s OKP support) to publish a verifiable
118    /// `/jwks.json` for the resulting deployment.
119    pub fn new_ed25519(
120        private_key_pem: &[u8],
121        issuer: Option<String>,
122        kid: Option<String>,
123    ) -> Result<Self, AuthError> {
124        let encoding_key = EncodingKey::from_ed_pem(private_key_pem)
125            .map_err(|e| AuthError::Token(e.to_string()))?;
126
127        let pem_str = std::str::from_utf8(private_key_pem)
128            .map_err(|_| AuthError::Token("Invalid PEM UTF-8".into()))?;
129
130        use ed25519_dalek::pkcs8::DecodePrivateKey;
131        let signing_key = ed25519_dalek::SigningKey::from_pkcs8_pem(pem_str)
132            .map_err(|e| AuthError::Token(format!("Failed to parse Ed25519 key: {}", e)))?;
133
134        use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
135        let x = URL_SAFE_NO_PAD.encode(signing_key.verifying_key().to_bytes());
136
137        let kid_val = kid.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
138
139        let jwk = crate::token::jwk::Jwk {
140            kid: Some(kid_val.clone()),
141            kty: "OKP".to_string(),
142            alg: Some("EdDSA".to_string()),
143            n: None,
144            e: None,
145            crv: Some("Ed25519".to_string()),
146            x: Some(x),
147        };
148
149        // Same rationale as `new_asymmetric`: derive the decoding key from
150        // the JWK we just built (the public half) rather than from the
151        // private PEM, so both provably agree with what `/jwks` publishes.
152        // See the regression note on that constructor and the test below
153        // named after it.
154        let decoding_key = jwk.to_decoding_key()?;
155
156        Ok(Self {
157            encoding_key,
158            decoding_key,
159            issuer,
160            kid: Some(kid_val),
161            alg: Algorithm::EdDSA,
162            public_jwk: Some(jwk),
163        })
164    }
165
166    pub fn public_jwk(&self) -> Option<crate::token::jwk::Jwk> {
167        self.public_jwk.clone()
168    }
169
170    pub fn with_issuer(mut self, issuer: String) -> Self {
171        self.issuer = Some(issuer);
172        self
173    }
174
175    /// Issues a token for a user identity.
176    pub fn issue_user_token(
177        &self,
178        identity: Identity,
179        expires_in_secs: u64,
180        scope: Option<String>,
181        aud: Option<String>,
182    ) -> Result<String, AuthError> {
183        self.issue_user_token_with_extra(identity, expires_in_secs, scope, aud, HashMap::new())
184    }
185
186    /// Issues a token for a user identity, stamping the given `extra` claims
187    /// onto the token in addition to the standard/core claims.
188    ///
189    /// This lets a host application (e.g. a resource server built on top of
190    /// this engine) attach domain-specific claims — such as `api_key_id`,
191    /// `project_id`, or `roles` — so downstream consumers (an API gateway or
192    /// authorization proxy) can read them directly off the token without a
193    /// database round-trip. Keys in `extra` take precedence over any
194    /// same-named field set elsewhere in `extra` by this method; they cannot
195    /// override the top-level standard claims (`sub`, `aud`, `exp`, etc.)
196    /// since those are not part of the flattened map.
197    pub fn issue_user_token_with_extra(
198        &self,
199        identity: Identity,
200        expires_in_secs: u64,
201        scope: Option<String>,
202        aud: Option<String>,
203        extra: HashMap<String, serde_json::Value>,
204    ) -> Result<String, AuthError> {
205        let now = chrono::Utc::now().timestamp() as usize;
206        let expiration = now + expires_in_secs as usize;
207
208        let claims = Claims {
209            iss: self.issuer.clone(),
210            sub: identity.external_id.clone(),
211            aud,
212            exp: expiration,
213            iat: now,
214            nbf: Some(now),
215            jti: Some(uuid::Uuid::new_v4().to_string()),
216            scope,
217            identity: Some(identity),
218            extra,
219        };
220
221        let mut header = Header::new(self.alg);
222        if let Some(ref kid) = self.kid {
223            header.kid = Some(kid.clone());
224        }
225
226        encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
227    }
228
229    /// Issues an OIDC-conformant ID token.
230    pub fn issue_id_token(
231        &self,
232        identity: Identity,
233        client_id: &str,
234        nonce: Option<String>,
235        expires_in_secs: u64,
236    ) -> Result<String, AuthError> {
237        self.issue_id_token_with_extra(identity, client_id, nonce, expires_in_secs, HashMap::new())
238    }
239
240    /// Issues an OIDC-conformant ID token, stamping the given `extra` claims
241    /// onto the token in addition to the standard/core claims.
242    ///
243    /// `nonce` is a reserved claim key: `extra` is merged into the token
244    /// first, then the explicit `nonce` parameter is applied on top. So if
245    /// `nonce` is `Some(_)`, it always wins over any `"nonce"` entry passed
246    /// in `extra`. If `nonce` is `None`, an `extra["nonce"]` value (if any)
247    /// is left as-is. This preserves OIDC `nonce` semantics — it reflects
248    /// what the client sent in the authorization request — and keeps it from
249    /// being accidentally clobbered by unrelated custom claims.
250    pub fn issue_id_token_with_extra(
251        &self,
252        identity: Identity,
253        client_id: &str,
254        nonce: Option<String>,
255        expires_in_secs: u64,
256        extra: HashMap<String, serde_json::Value>,
257    ) -> Result<String, AuthError> {
258        let now = chrono::Utc::now().timestamp() as usize;
259        let expiration = now + expires_in_secs as usize;
260
261        let mut claims = Claims {
262            iss: self.issuer.clone(),
263            sub: identity.external_id.clone(),
264            aud: Some(client_id.to_string()),
265            exp: expiration,
266            iat: now,
267            nbf: Some(now),
268            jti: Some(uuid::Uuid::new_v4().to_string()),
269            scope: None,
270            identity: Some(identity),
271            extra,
272        };
273
274        if let Some(n) = nonce {
275            claims
276                .extra
277                .insert("nonce".to_string(), serde_json::Value::String(n));
278        }
279
280        let mut header = Header::new(self.alg);
281        if let Some(ref kid) = self.kid {
282            header.kid = Some(kid.clone());
283        }
284
285        encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
286    }
287
288    /// Issues a machine-to-machine (M2M) token for a client.
289    pub fn issue_client_token(
290        &self,
291        client_id: &str,
292        expires_in_secs: u64,
293        scope: Option<String>,
294        aud: Option<String>,
295    ) -> Result<String, AuthError> {
296        self.issue_client_token_with_extra(client_id, expires_in_secs, scope, aud, HashMap::new())
297    }
298
299    /// Issues a machine-to-machine (M2M) token for a client, stamping the
300    /// given `extra` claims onto the token in addition to the standard/core
301    /// claims. See [`Self::issue_user_token_with_extra`] for the rationale.
302    pub fn issue_client_token_with_extra(
303        &self,
304        client_id: &str,
305        expires_in_secs: u64,
306        scope: Option<String>,
307        aud: Option<String>,
308        extra: HashMap<String, serde_json::Value>,
309    ) -> Result<String, AuthError> {
310        let now = chrono::Utc::now().timestamp() as usize;
311        let expiration = now + expires_in_secs as usize;
312
313        let claims = Claims {
314            iss: self.issuer.clone(),
315            sub: client_id.to_string(),
316            aud,
317            exp: expiration,
318            iat: now,
319            nbf: Some(now),
320            jti: Some(uuid::Uuid::new_v4().to_string()),
321            scope,
322            identity: None,
323            extra,
324        };
325
326        let mut header = Header::new(self.alg);
327        if let Some(ref kid) = self.kid {
328            header.kid = Some(kid.clone());
329        }
330
331        encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
332    }
333
334    /// Issues a token with an explicit `typ` header and no `aud`, for
335    /// callers minting something that is not a standard OIDC ID/access/user
336    /// token and needs its own wire-format `typ` so verifiers can tell it
337    /// apart from those (e.g. `authkestra-op`'s device/service attestations,
338    /// whose contract requires `typ: "webank-attest+jws"` rather than the
339    /// default `"JWT"`). Additive alongside the `issue_*_token*` family
340    /// above; those are unchanged.
341    pub fn issue_custom_token(
342        &self,
343        sub: String,
344        expires_in_secs: u64,
345        typ: &str,
346        extra: HashMap<String, serde_json::Value>,
347    ) -> Result<String, AuthError> {
348        let now = chrono::Utc::now().timestamp() as usize;
349        let claims = Claims {
350            iss: self.issuer.clone(),
351            sub,
352            aud: None,
353            exp: now + expires_in_secs as usize,
354            iat: now,
355            nbf: Some(now),
356            jti: Some(uuid::Uuid::new_v4().to_string()),
357            scope: None,
358            identity: None,
359            extra,
360        };
361
362        let mut header = Header::new(self.alg);
363        header.typ = Some(typ.to_string());
364        if let Some(ref kid) = self.kid {
365            header.kid = Some(kid.clone());
366        }
367
368        encode(&header, &claims, &self.encoding_key).map_err(|e| AuthError::Token(e.to_string()))
369    }
370
371    pub fn validate_token(
372        &self,
373        token: &str,
374        expected_aud: Option<&str>,
375    ) -> Result<Claims, AuthError> {
376        let mut validation = Validation::new(self.alg);
377        if let Some(aud) = expected_aud {
378            validation.set_audience(&[aud]);
379        } else {
380            validation.validate_aud = false;
381        }
382        if let Some(ref iss) = self.issuer {
383            validation.set_issuer(&[iss]);
384        }
385
386        let token_data = decode::<Claims>(token, &self.decoding_key, &validation)
387            .map_err(|e| AuthError::Token(e.to_string()))?;
388
389        Ok(token_data.claims)
390    }
391}
392
393#[cfg(test)]
394mod tests {
395
396    use super::*;
397    use crate::auth::state::Identity;
398    use std::collections::HashMap;
399
400    /// Throwaway RSA-2048 private key, test-only.
401    const TEST_RSA_PRIVATE_KEY_PEM: &[u8] = b"-----BEGIN PRIVATE KEY-----
402MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDA5hJIcQ+2rxMz
403VM8ZH5WAmguCr0xmNDAdy0IzzsUeFLG7BebB7izOkU36J4t8t5tUaQwrBMnx2Fvt
404VqJjbdE242UDpvWF/8m9zJ2HR5298cbwT5cGMKLB0HWzDMahugs+Bbh2lCgwyLZk
405Tr3Diwxp5SwFew/Wb+Ke9cNG9Hu5IFH3BCuJ839d9hfqisIeYrBPfb52xxckM37R
4067zSGu/eDP/HZAeLkQuptZJW4A3u7xni14u4qyqXDqsHsYFNgJaxMSAwWgBRY6HNu
407TnvBArTXCiVfL+F73B2L6mdYr64g+QS9nK9v97MlJu/E3mSduz54pren4mpCHc9m
408/S2+VjCZAgMBAAECggEAASC9qQbGnL7XuExRDOIn/m4bWx92ehjo0lCTibhpY3LW
409umbSbpfbhmmuSj3CjW9VZsaM3hBTgSjoTX72lbY/eIUXD7c0memUK5pV4XcEIrQw
410AZlPIye6ckx4I7ZGnKasO8FoAel9dd7DXw36AuBK3LBzJwtzkEFsBc0e3/wixqmG
411UJBbbt/+5ya7CxyjuePaQhKtkLD5R6DpvN2XnCYq5nHJNJdvSVg1pOzsTHYIf+Ee
4122Rz42fGsfFKqeEQCcBFRZaGb/ELeP4c6UZdktZAvmHb1p1fursVZc6X9JXmiJ2OJ
413Kv2H2tMKuysP8L0fXFOMgkH2SVt6rcdHkO6xhlhWsQKBgQDqR8rAJeEE5BFoXA8T
414VVW6CLMlW51x4ey7PEGOaYh39dTG2Q+GZQBZ9G+SZk3f5Y85UCACSyc//4qaz/c3
4150nWsegZ+JPyymmuc79wzIAFFvXB7pL6wyn0Ed1P620kOZTtA8iBcXrsuxL+KP7iu
416MXfWmU1QiZpbndILtyDnY+70uwKBgQDSyCljWkydQCaPU+fiAXLxP8CvcJTSSNQD
417mVUlwJ+OpHnU+Alsi1rBavMgUtLlYbFqzH7NmYrLC8Yadq3ZOwLt0VEK0r8qstAL
4187QCDUD2WNuQjpZupRnXuMUl3iXB96i2gb+VQKGuUAJvVWjdIbYa4+Gu+sBMfcDcX
419dBihDLuEuwKBgAgX4tEwfc2Fc3R/eaXZVNTQaB/qQk4k1+C//CPHUYeTXn5gEUE7
420S//PiesszZPmgkQgmHp7zidP1KH0fT3Yb2g97ut8q54f54fMYXcCrAiUusYKsuu4
421kwkMdkI8QRHWPW3I74VBYIYFFfjYqrCZ1OH8+cbGeiagFRmCggh8U0zxAoGAVW3u
4226Ge22Z0gg8LcHsu7jG/sZq7Ygool8/d3fT+e669Z+ak2GJo6hF4WgClRdMqtn72W
423PzpV+ImjFyK2v26dd0n48MwN0v56N/ss1Av3iiRhPtlmR6tZLNspDZvUzhPVvkrb
424xCs9vtSoVEamVWKe0eVNthGjDoDqs0TInq2MavUCgYB6REavSJs/CLkSS7iimjxZ
425G7g5YQi9/p1lXLOEUDiwEmvRr0XTwzzxUsIc535IXhh/ZUYpthenW+qBBzn85pEC
426TowIqciHu5redqlQ8rITA8/AOY98vaDIhppDg1rfpnHHaZHFbXD/keYAEbhBtbvf
427a0QMqKUcs8+YTy5R5K6qtw==
428-----END PRIVATE KEY-----";
429
430    /// Throwaway Ed25519 private key (PKCS#8 PEM), test-only. Generated with
431    /// `openssl genpkey -algorithm ed25519`.
432    const TEST_ED25519_PRIVATE_KEY_PEM: &[u8] = b"-----BEGIN PRIVATE KEY-----
433MC4CAQAwBQYDK2VwBCIEIKIPR2jojpdobYr1M/pjIRuMONpZGYQ+y5yxSqKX9T9/
434-----END PRIVATE KEY-----";
435
436    /// A second, distinct throwaway Ed25519 private key, test-only — used to
437    /// prove a token signed by one key is rejected by another key's manager.
438    const TEST_ED25519_PRIVATE_KEY_PEM_B: &[u8] = b"-----BEGIN PRIVATE KEY-----
439MC4CAQAwBQYDK2VwBCIEIPlsnSfvh53rJ+Tlbo8e7cgq2mIkWQ1NCM5paVeinUh8
440-----END PRIVATE KEY-----";
441
442    #[test]
443    fn test_claims_serialization() {
444        let mut extra = HashMap::new();
445        extra.insert(
446            "custom".to_string(),
447            serde_json::Value::String("value".to_string()),
448        );
449
450        let claims = Claims {
451            iss: Some("issuer".to_string()),
452            sub: "user123".to_string(),
453            aud: Some("audience".to_string()),
454            exp: 1000,
455            iat: 500,
456            nbf: Some(500),
457            jti: Some("jti".to_string()),
458            scope: Some("openid profile".to_string()),
459            identity: Some(Identity {
460                provider_id: "google".to_string(),
461                external_id: "user123".to_string(),
462                email: Some("user@example.com".to_string()),
463                username: Some("user".to_string()),
464                attributes: HashMap::new(),
465            }),
466            extra,
467        };
468
469        let serialized = serde_json::to_string(&claims).unwrap();
470        let deserialized: Claims = serde_json::from_str(&serialized).unwrap();
471
472        assert_eq!(deserialized.iss, claims.iss);
473        assert_eq!(deserialized.sub, claims.sub);
474        assert_eq!(deserialized.extra.get("custom").unwrap(), "value");
475    }
476
477    #[test]
478    fn test_token_manager_issuance() {
479        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
480        let identity = Identity {
481            provider_id: "mock".to_string(),
482            external_id: "user123".to_string(),
483            email: None,
484            username: None,
485            attributes: HashMap::new(),
486        };
487
488        let token = manager
489            .issue_user_token(identity, 3600, None, None)
490            .unwrap();
491        let claims = manager.validate_token(&token, None).unwrap();
492
493        assert_eq!(claims.iss, Some("issuer".to_string()));
494        assert_eq!(claims.sub, "user123");
495        assert!(claims.jti.is_some());
496        assert!(claims.nbf.is_some());
497    }
498
499    #[test]
500    fn test_token_manager_asymmetric_issuance() {
501        let manager = TokenManager::new_asymmetric(
502            TEST_RSA_PRIVATE_KEY_PEM,
503            Some("issuer".to_string()),
504            Some("my-kid-123".to_string()),
505        )
506        .unwrap();
507
508        let identity = Identity {
509            provider_id: "mock".to_string(),
510            external_id: "user123".to_string(),
511            email: None,
512            username: None,
513            attributes: HashMap::new(),
514        };
515
516        let token = manager
517            .issue_user_token(identity, 3600, None, None)
518            .unwrap();
519
520        // Decode directly via jsonwebtoken to prove independent verification
521        let jwk = manager.public_jwk().unwrap();
522        assert_eq!(jwk.kid.as_deref(), Some("my-kid-123"));
523
524        let decoding_key = jwk.to_decoding_key().unwrap();
525        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
526        validation.set_issuer(&["issuer"]);
527
528        let token_data =
529            jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap();
530        assert_eq!(token_data.claims.sub, "user123");
531        assert_eq!(token_data.header.kid.as_deref(), Some("my-kid-123"));
532    }
533
534    /// Regression test for the asymmetric decoding key being derived from the
535    /// private half instead of the public one.
536    ///
537    /// `test_token_manager_asymmetric_issuance` above verifies via the
538    /// published JWK, which exercises `Jwk::to_decoding_key` rather than the
539    /// manager's own `decoding_key` — so it stays green either way. This one
540    /// goes through `validate_token`, which is the path `/reissue` and
541    /// `/userinfo` actually take.
542    #[test]
543    fn test_asymmetric_manager_validates_its_own_tokens() {
544        let manager = TokenManager::new_asymmetric(
545            TEST_RSA_PRIVATE_KEY_PEM,
546            Some("issuer".to_string()),
547            Some("my-kid-123".to_string()),
548        )
549        .unwrap();
550
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 token = manager
560            .issue_user_token(identity, 3600, None, None)
561            .unwrap();
562
563        let claims = manager.validate_token(&token, None).unwrap();
564        assert_eq!(claims.sub, "user123");
565        assert_eq!(claims.iss, Some("issuer".to_string()));
566
567        // Same round trip for the attestation shape `/reissue` presents.
568        let attestation = manager
569            .issue_custom_token(
570                "device-1".to_string(),
571                60,
572                "webank-attest+jws",
573                HashMap::new(),
574            )
575            .unwrap();
576
577        let attest_claims = manager.validate_token(&attestation, None).unwrap();
578        assert_eq!(attest_claims.sub, "device-1");
579    }
580
581    #[test]
582    fn test_issue_id_token() {
583        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
584        let identity = Identity {
585            provider_id: "mock".to_string(),
586            external_id: "user123".to_string(),
587            email: None,
588            username: None,
589            attributes: HashMap::new(),
590        };
591
592        let token = manager
593            .issue_id_token(identity, "client-1", Some("nonce123".to_string()), 3600)
594            .unwrap();
595
596        let claims = manager.validate_token(&token, None).unwrap();
597
598        assert_eq!(claims.iss, Some("issuer".to_string()));
599        assert_eq!(claims.sub, "user123");
600        assert_eq!(claims.aud, Some("client-1".to_string()));
601        assert_eq!(claims.extra.get("nonce").unwrap(), "nonce123");
602    }
603    #[test]
604    fn test_token_manager_audience_validation() {
605        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
606        let identity = Identity {
607            provider_id: "mock".to_string(),
608            external_id: "user123".to_string(),
609            email: None,
610            username: None,
611            attributes: HashMap::new(),
612        };
613
614        // Issue token for "client-1"
615        let token = manager
616            .issue_id_token(identity, "client-1", None, 3600)
617            .unwrap();
618
619        // Validate with correct audience
620        let claims = manager.validate_token(&token, Some("client-1")).unwrap();
621        assert_eq!(claims.aud, Some("client-1".to_string()));
622
623        // Validate with incorrect audience (should fail)
624        let err = manager
625            .validate_token(&token, Some("client-2"))
626            .unwrap_err();
627        assert!(err.to_string().contains("InvalidAudience"));
628    }
629
630    #[test]
631    fn test_issue_custom_token_sets_typ_and_no_aud() {
632        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
633
634        let mut extra = HashMap::new();
635        extra.insert("cnf".to_string(), serde_json::json!({"jkt": "abc"}));
636
637        let token = manager
638            .issue_custom_token("device-1".to_string(), 60, "webank-attest+jws", extra)
639            .unwrap();
640
641        let header = jsonwebtoken::decode_header(&token).unwrap();
642        assert_eq!(header.typ.as_deref(), Some("webank-attest+jws"));
643
644        let claims = manager.validate_token(&token, None).unwrap();
645        assert_eq!(claims.sub, "device-1");
646        assert_eq!(claims.aud, None);
647        assert_eq!(claims.extra.get("cnf").unwrap()["jkt"], "abc");
648    }
649
650    #[test]
651    fn test_issue_user_token_with_extra_round_trips_custom_claims() {
652        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
653        let identity = Identity {
654            provider_id: "mock".to_string(),
655            external_id: "user123".to_string(),
656            email: None,
657            username: None,
658            attributes: HashMap::new(),
659        };
660
661        let mut extra = HashMap::new();
662        extra.insert("api_key_id".to_string(), serde_json::json!("key-abc"));
663        extra.insert("project_id".to_string(), serde_json::json!("proj-42"));
664
665        let token = manager
666            .issue_user_token_with_extra(identity, 3600, None, None, extra)
667            .unwrap();
668        let claims = manager.validate_token(&token, None).unwrap();
669
670        assert_eq!(
671            claims.extra.get("api_key_id"),
672            Some(&serde_json::json!("key-abc"))
673        );
674        assert_eq!(
675            claims.extra.get("project_id"),
676            Some(&serde_json::json!("proj-42"))
677        );
678    }
679
680    #[test]
681    fn test_issue_user_token_wrapper_still_has_empty_extra() {
682        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
683        let identity = Identity {
684            provider_id: "mock".to_string(),
685            external_id: "user123".to_string(),
686            email: None,
687            username: None,
688            attributes: HashMap::new(),
689        };
690
691        let token = manager
692            .issue_user_token(identity, 3600, None, None)
693            .unwrap();
694        let claims = manager.validate_token(&token, None).unwrap();
695
696        assert!(claims.extra.is_empty());
697    }
698
699    #[test]
700    fn test_issue_client_token_with_extra_round_trips_custom_claims() {
701        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
702
703        let mut extra = HashMap::new();
704        extra.insert("roles".to_string(), serde_json::json!(["admin", "billing"]));
705
706        let token = manager
707            .issue_client_token_with_extra("client-1", 3600, None, None, extra)
708            .unwrap();
709        let claims = manager.validate_token(&token, None).unwrap();
710
711        assert_eq!(
712            claims.extra.get("roles"),
713            Some(&serde_json::json!(["admin", "billing"]))
714        );
715    }
716
717    #[test]
718    fn test_issue_client_token_wrapper_still_has_empty_extra() {
719        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
720
721        let token = manager
722            .issue_client_token("client-1", 3600, None, None)
723            .unwrap();
724        let claims = manager.validate_token(&token, None).unwrap();
725
726        assert!(claims.extra.is_empty());
727    }
728
729    #[test]
730    fn test_issue_id_token_with_extra_round_trips_custom_claims() {
731        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
732        let identity = Identity {
733            provider_id: "mock".to_string(),
734            external_id: "user123".to_string(),
735            email: None,
736            username: None,
737            attributes: HashMap::new(),
738        };
739
740        let mut extra = HashMap::new();
741        extra.insert("org_id".to_string(), serde_json::json!("org-7"));
742
743        let token = manager
744            .issue_id_token_with_extra(
745                identity,
746                "client-1",
747                Some("nonce123".to_string()),
748                3600,
749                extra,
750            )
751            .unwrap();
752        let claims = manager.validate_token(&token, None).unwrap();
753
754        assert_eq!(
755            claims.extra.get("org_id"),
756            Some(&serde_json::json!("org-7"))
757        );
758        assert_eq!(
759            claims.extra.get("nonce"),
760            Some(&serde_json::json!("nonce123"))
761        );
762    }
763
764    #[test]
765    fn test_issue_id_token_with_extra_explicit_nonce_wins_over_extra_nonce() {
766        // Documents the precedence chosen in `issue_id_token_with_extra`:
767        // the explicit `nonce` parameter always overrides a `"nonce"` entry
768        // supplied via `extra`.
769        let manager = TokenManager::new(b"secret", Some("issuer".to_string()));
770        let identity = Identity {
771            provider_id: "mock".to_string(),
772            external_id: "user123".to_string(),
773            email: None,
774            username: None,
775            attributes: HashMap::new(),
776        };
777
778        let mut extra = HashMap::new();
779        extra.insert(
780            "nonce".to_string(),
781            serde_json::json!("attacker-supplied-nonce"),
782        );
783
784        let token = manager
785            .issue_id_token_with_extra(
786                identity,
787                "client-1",
788                Some("real-nonce".to_string()),
789                3600,
790                extra,
791            )
792            .unwrap();
793        let claims = manager.validate_token(&token, None).unwrap();
794
795        assert_eq!(
796            claims.extra.get("nonce"),
797            Some(&serde_json::json!("real-nonce"))
798        );
799    }
800
801    #[test]
802    fn test_token_manager_ed25519_issuance() {
803        let manager = TokenManager::new_ed25519(
804            TEST_ED25519_PRIVATE_KEY_PEM,
805            Some("issuer".to_string()),
806            Some("my-ed25519-kid".to_string()),
807        )
808        .unwrap();
809
810        let identity = Identity {
811            provider_id: "mock".to_string(),
812            external_id: "user123".to_string(),
813            email: None,
814            username: None,
815            attributes: HashMap::new(),
816        };
817
818        let token = manager
819            .issue_user_token(identity, 3600, None, None)
820            .unwrap();
821
822        let header = jsonwebtoken::decode_header(&token).unwrap();
823        assert_eq!(header.alg, Algorithm::EdDSA);
824        assert_eq!(header.kid.as_deref(), Some("my-ed25519-kid"));
825    }
826
827    /// Proves #187 and #188 actually compose end-to-end: mint a token with
828    /// the Ed25519 constructor, then verify it using ONLY the decoding key
829    /// derived from the published JWK (the OKP shape `/jwks.json` would
830    /// serve) — not the manager's own internal decoding key. This is the
831    /// round trip a real resource server performs against a real JWKS
832    /// endpoint.
833    #[test]
834    fn test_ed25519_token_round_trips_via_published_jwk() {
835        let manager = TokenManager::new_ed25519(
836            TEST_ED25519_PRIVATE_KEY_PEM,
837            Some("issuer".to_string()),
838            Some("my-ed25519-kid".to_string()),
839        )
840        .unwrap();
841
842        let identity = Identity {
843            provider_id: "mock".to_string(),
844            external_id: "user123".to_string(),
845            email: None,
846            username: None,
847            attributes: HashMap::new(),
848        };
849
850        let token = manager
851            .issue_user_token(identity, 3600, None, None)
852            .unwrap();
853
854        let jwk = manager.public_jwk().unwrap();
855        assert_eq!(jwk.kid.as_deref(), Some("my-ed25519-kid"));
856
857        let decoding_key = jwk.to_decoding_key().unwrap();
858        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::EdDSA);
859        validation.set_issuer(&["issuer"]);
860
861        let token_data =
862            jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap();
863        assert_eq!(token_data.claims.sub, "user123");
864        assert_eq!(token_data.header.kid.as_deref(), Some("my-ed25519-kid"));
865    }
866
867    /// Same shape as `test_asymmetric_manager_validates_its_own_tokens`
868    /// (RS256): goes through `validate_token`, the path `/reissue` and
869    /// `/userinfo` actually take, rather than the published-JWK path above.
870    #[test]
871    fn test_ed25519_manager_validates_its_own_tokens() {
872        let manager = TokenManager::new_ed25519(
873            TEST_ED25519_PRIVATE_KEY_PEM,
874            Some("issuer".to_string()),
875            Some("my-ed25519-kid".to_string()),
876        )
877        .unwrap();
878
879        let identity = Identity {
880            provider_id: "mock".to_string(),
881            external_id: "user123".to_string(),
882            email: None,
883            username: None,
884            attributes: HashMap::new(),
885        };
886
887        let token = manager
888            .issue_user_token(identity, 3600, None, None)
889            .unwrap();
890
891        let claims = manager.validate_token(&token, None).unwrap();
892        assert_eq!(claims.sub, "user123");
893        assert_eq!(claims.iss, Some("issuer".to_string()));
894    }
895
896    /// `public_jwk()` for an Ed25519 manager must emit spec-correct OKP JSON
897    /// (RFC 8037 §2): `kty: "OKP"`, `crv: "Ed25519"`, `x` present and
898    /// base64url-encoded to exactly 32 bytes, and no stray RSA fields (`n`,
899    /// `e`) on the wire.
900    #[test]
901    fn test_ed25519_public_jwk_is_spec_correct_okp_json() {
902        let manager = TokenManager::new_ed25519(
903            TEST_ED25519_PRIVATE_KEY_PEM,
904            Some("issuer".to_string()),
905            Some("my-ed25519-kid".to_string()),
906        )
907        .unwrap();
908
909        let jwk = manager.public_jwk().unwrap();
910        assert_eq!(jwk.kty, "OKP");
911        assert_eq!(jwk.alg.as_deref(), Some("EdDSA"));
912        assert_eq!(jwk.crv.as_deref(), Some("Ed25519"));
913        assert!(jwk.n.is_none());
914        assert!(jwk.e.is_none());
915
916        let x = jwk.x.as_ref().expect("OKP JWK must have 'x'");
917        use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
918        let decoded = URL_SAFE_NO_PAD
919            .decode(x)
920            .expect("'x' must be valid base64url (unpadded)");
921        assert_eq!(decoded.len(), 32, "Ed25519 public key must be 32 bytes");
922
923        let value = serde_json::to_value(&jwk).unwrap();
924        let obj = value.as_object().unwrap();
925        assert_eq!(obj.get("kty").unwrap(), "OKP");
926        assert_eq!(obj.get("crv").unwrap(), "Ed25519");
927        assert_eq!(obj.get("alg").unwrap(), "EdDSA");
928        assert!(obj.contains_key("x"));
929        assert!(
930            !obj.contains_key("n"),
931            "OKP JWK must not serialize the RSA 'n' field, got: {}",
932            value
933        );
934        assert!(
935            !obj.contains_key("e"),
936            "OKP JWK must not serialize the RSA 'e' field, got: {}",
937            value
938        );
939    }
940
941    /// Same spec-correctness check for the RSA shape, unchanged by the OKP
942    /// addition: no stray `crv`/`x` fields on the wire.
943    #[test]
944    fn test_rsa_public_jwk_still_omits_okp_fields() {
945        let manager = TokenManager::new_asymmetric(
946            TEST_RSA_PRIVATE_KEY_PEM,
947            Some("issuer".to_string()),
948            Some("my-kid-123".to_string()),
949        )
950        .unwrap();
951
952        let jwk = manager.public_jwk().unwrap();
953        let value = serde_json::to_value(&jwk).unwrap();
954        let obj = value.as_object().unwrap();
955        assert_eq!(obj.get("kty").unwrap(), "RSA");
956        assert!(obj.contains_key("n"));
957        assert!(obj.contains_key("e"));
958        assert!(
959            !obj.contains_key("crv"),
960            "RSA JWK must not serialize the OKP 'crv' field, got: {}",
961            value
962        );
963        assert!(
964            !obj.contains_key("x"),
965            "RSA JWK must not serialize the OKP 'x' field, got: {}",
966            value
967        );
968    }
969
970    /// Wrong-key rejection: a token signed by one Ed25519 manager must be
971    /// rejected when verified against a *different* Ed25519 manager's
972    /// published JWK — proving the JWKS round trip actually checks the
973    /// signature rather than trivially accepting any well-formed EdDSA JWT.
974    #[test]
975    fn test_ed25519_token_rejected_by_wrong_key_jwk() {
976        let signer = TokenManager::new_ed25519(
977            TEST_ED25519_PRIVATE_KEY_PEM,
978            Some("issuer".to_string()),
979            Some("kid-a".to_string()),
980        )
981        .unwrap();
982        let other = TokenManager::new_ed25519(
983            TEST_ED25519_PRIVATE_KEY_PEM_B,
984            Some("issuer".to_string()),
985            Some("kid-b".to_string()),
986        )
987        .unwrap();
988
989        let identity = Identity {
990            provider_id: "mock".to_string(),
991            external_id: "user123".to_string(),
992            email: None,
993            username: None,
994            attributes: HashMap::new(),
995        };
996
997        let token = signer.issue_user_token(identity, 3600, None, None).unwrap();
998
999        // Verifying with the signer's own manager still works.
1000        assert!(signer.validate_token(&token, None).is_ok());
1001
1002        // Verifying with a different key's manager must fail.
1003        let err = other.validate_token(&token, None).unwrap_err();
1004        assert!(err.to_string().contains("InvalidSignature"));
1005
1006        // Same result going through the published-JWK path a real resource
1007        // server would use.
1008        let wrong_jwk = other.public_jwk().unwrap();
1009        let decoding_key = wrong_jwk.to_decoding_key().unwrap();
1010        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::EdDSA);
1011        validation.set_issuer(&["issuer"]);
1012        let err = jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation).unwrap_err();
1013        assert_eq!(
1014            err.kind(),
1015            &jsonwebtoken::errors::ErrorKind::InvalidSignature
1016        );
1017    }
1018
1019    /// A tampered payload (claims byte flipped after signing, signature
1020    /// left as-is) must be rejected, whichever algorithm signed it —
1021    /// guards against a validator that only checks structural shape.
1022    #[test]
1023    fn test_ed25519_tampered_token_rejected() {
1024        let manager = TokenManager::new_ed25519(
1025            TEST_ED25519_PRIVATE_KEY_PEM,
1026            Some("issuer".to_string()),
1027            Some("my-ed25519-kid".to_string()),
1028        )
1029        .unwrap();
1030
1031        let identity = Identity {
1032            provider_id: "mock".to_string(),
1033            external_id: "user123".to_string(),
1034            email: None,
1035            username: None,
1036            attributes: HashMap::new(),
1037        };
1038
1039        let token = manager
1040            .issue_user_token(identity, 3600, None, None)
1041            .unwrap();
1042
1043        let mut parts: Vec<&str> = token.split('.').collect();
1044        assert_eq!(parts.len(), 3);
1045        // Corrupt a single character in the base64url payload segment.
1046        let mut payload = parts[1].to_string();
1047        let last = payload.pop().unwrap();
1048        let replacement = if last == 'A' { 'B' } else { 'A' };
1049        payload.push(replacement);
1050        parts[1] = &payload;
1051        let tampered = parts.join(".");
1052
1053        let err = manager.validate_token(&tampered, None).unwrap_err();
1054        assert!(
1055            err.to_string().contains("InvalidSignature") || err.to_string().contains("Json"),
1056            "unexpected error for tampered token: {err}"
1057        );
1058    }
1059}
1060pub mod jwk;
authkestra_engine/token/mod.rs

authkestra_engine/token/
mod.rs
