1use std::collections::{HashMap, HashSet};
16use std::sync::Arc;
17use std::sync::Mutex;
18use std::time::Duration as StdDuration;
19
20use crate::identity::IdentityRegistry;
21use axum::{
22 body::Body,
23 extract::{Query, State},
24 http::{header, StatusCode},
25 response::{IntoResponse, Json, Redirect, Response},
26};
27use axum_extra::extract::cookie::{Cookie, SameSite};
28use pep::oidc::pkce_cookie::PkceCookieManager;
29use pep::oidc_client::OidcClient;
30use pep::oidc_resource_server::ResourceServerClient;
31use pep::session_manager::WebSessionManager;
32use pep::{DevConfig, JwtClaims, JwtValidationOptions, OidcClientConfig};
33use serde::Deserialize;
34use sha2::{Digest, Sha256};
35use time::Duration as TimeDuration;
36
37use cedar_policy::{Context, Entities, EntityUid, Request};
38use std::str::FromStr;
39
40#[derive(Debug, Clone)]
46pub struct AuthConfig {
47 pub issuer_url: String,
49 pub client_id: String,
51 pub client_secret: Option<String>,
53 pub redirect_uri: String,
55 pub scope: String,
57 pub cookie_name: String,
59 pub dev_config: DevConfig,
61 pub validation_options: JwtValidationOptions,
63 pub pkce_cookie_secret: String,
65}
66
67impl AuthConfig {
68 pub fn from_toml(config_toml: &str) -> Option<Self> {
73 let table: toml::Table = toml::from_str(config_toml).ok()?;
74
75 let dev_config = table
77 .get("dev")
78 .and_then(|d| d.as_table())
79 .map(|d| DevConfig {
80 local_dev_mode: d
81 .get("local_dev_mode")
82 .and_then(|v| v.as_bool())
83 .unwrap_or(false),
84 local_dev_email: d
85 .get("local_dev_email")
86 .and_then(|v| v.as_str())
87 .map(String::from),
88 local_dev_name: d
89 .get("local_dev_name")
90 .and_then(|v| v.as_str())
91 .map(String::from),
92 local_dev_username: d
93 .get("local_dev_username")
94 .and_then(|v| v.as_str())
95 .map(String::from),
96 });
97
98 if let Some(ref dc) = dev_config {
100 if dc.local_dev_mode {
101 let oidc = Self::parse_oidc_section(&table);
103 return Some(Self {
104 issuer_url: oidc
105 .as_ref()
106 .map(|o| o.0.clone())
107 .unwrap_or_else(|| "https://auth.example.com".into()),
108 client_id: oidc
109 .as_ref()
110 .map(|o| o.1.clone())
111 .unwrap_or_else(|| "trustee".into()),
112 client_secret: oidc.as_ref().and_then(|o| o.2.clone()),
113 redirect_uri: oidc
114 .as_ref()
115 .map(|o| o.3.clone())
116 .unwrap_or_else(|| "http://localhost:3000/auth/callback".into()),
117 scope: oidc
118 .as_ref()
119 .map(|o| o.4.clone())
120 .unwrap_or_else(|| "openid profile email".into()),
121 cookie_name: "trustee_token".into(),
122 dev_config: dc.clone(),
123 validation_options: JwtValidationOptions::default(),
124 pkce_cookie_secret: oidc
125 .as_ref()
126 .map(|o| o.6.clone())
127 .unwrap_or_else(|| "trustee-default-pkce-secret-change-me".into()),
128 });
129 }
130 }
131
132 let (
134 issuer_url,
135 client_id,
136 client_secret,
137 redirect_uri,
138 scope,
139 validation_options,
140 pkce_secret,
141 ) = Self::parse_oidc_section(&table)?;
142
143 Some(Self {
144 issuer_url,
145 client_id,
146 client_secret,
147 redirect_uri,
148 scope,
149 cookie_name: "trustee_token".into(),
150 dev_config: dev_config.unwrap_or_default(),
151 validation_options,
152 pkce_cookie_secret: pkce_secret,
153 })
154 }
155
156 fn parse_oidc_section(
159 table: &toml::Table,
160 ) -> Option<(
161 String,
162 String,
163 Option<String>,
164 String,
165 String,
166 JwtValidationOptions,
167 String,
168 )> {
169 let oidc = table.get("oidc")?.as_table()?;
170
171 let issuer_url = oidc.get("issuer_url")?.as_str()?.to_string();
172 let client_id = oidc.get("client_id")?.as_str()?.to_string();
173 let client_secret = oidc
174 .get("client_secret")
175 .and_then(|v| v.as_str())
176 .map(String::from);
177 let redirect_uri = oidc
178 .get("redirect_uri")
179 .or_else(|| oidc.get("redirect_url")) .and_then(|v| v.as_str())
181 .unwrap_or("http://localhost:3000/auth/callback")
182 .to_string();
183 let scope = oidc
184 .get("scope")
185 .and_then(|v| v.as_str())
186 .unwrap_or("openid profile email")
187 .to_string();
188
189 let mut validation_options = JwtValidationOptions::default();
190 if let Some(skip) = oidc.get("skip_issuer_validation").and_then(|v| v.as_bool()) {
191 validation_options.skip_issuer_validation = skip;
192 }
193 if let Some(skip) = oidc
194 .get("skip_audience_validation")
195 .and_then(|v| v.as_bool())
196 {
197 validation_options.skip_audience_validation = skip;
198 }
199 validation_options.expected_audience = oidc
200 .get("expected_audience")
201 .and_then(|v| v.as_str())
202 .map(String::from);
203
204 let pkce_secret = oidc
205 .get("pkce_cookie_secret")
206 .and_then(|v| v.as_str())
207 .unwrap_or("trustee-default-pkce-secret-change-me")
208 .to_string();
209
210 Some((
211 issuer_url,
212 client_id,
213 client_secret,
214 redirect_uri,
215 scope,
216 validation_options,
217 pkce_secret,
218 ))
219 }
220
221 pub fn oidc_client_config(&self) -> OidcClientConfig {
223 OidcClientConfig {
224 issuer_url: self.issuer_url.clone(),
225 client_id: self.client_id.clone(),
226 client_secret: self.client_secret.clone(),
227 redirect_uri: self.redirect_uri.clone(),
228 scope: self.scope.clone(),
229 code_challenge_method: "S256".to_string(),
230 }
231 }
232}
233
234const CACHE_EXP_GRACE_SECS: i64 = 120;
242
243const CACHE_TTL_SECS: i64 = 300;
246
247const CACHE_MAX_ENTRIES: usize = 1024;
249
250fn validation_cache_key(token: &str) -> [u8; 16] {
253 let digest = Sha256::digest(token.as_bytes());
254 let mut key = [0u8; 16];
255 key.copy_from_slice(&digest[..16]);
256 key
257}
258
259fn unix_now() -> i64 {
260 std::time::SystemTime::now()
261 .duration_since(std::time::UNIX_EPOCH)
262 .map(|d| d.as_secs() as i64)
263 .unwrap_or(0)
264}
265
266struct ValidationCache {
276 inner: Mutex<ValidationCacheInner>,
277}
278
279#[derive(Default)]
280struct ValidationCacheInner {
281 entries: HashMap<[u8; 16], (JwtClaims, i64)>,
283 first_sight: HashSet<(String, String)>,
285}
286
287impl ValidationCache {
288 fn new() -> Self {
289 Self {
290 inner: Mutex::new(ValidationCacheInner::default()),
291 }
292 }
293
294 fn get(&self, key: &[u8; 16], now: i64) -> Option<JwtClaims> {
296 let inner = self.inner.lock().expect("validation cache poisoned");
297 inner
298 .entries
299 .get(key)
300 .filter(|(_, until)| *until > now)
301 .map(|(claims, _)| claims.clone())
302 }
303
304 fn put(&self, key: [u8; 16], claims: JwtClaims, now: i64) {
308 let valid_until = (claims.exp - CACHE_EXP_GRACE_SECS).min(now + CACHE_TTL_SECS);
309 let mut inner = self.inner.lock().expect("validation cache poisoned");
310 if inner.entries.len() >= CACHE_MAX_ENTRIES {
311 inner.entries.retain(|_, (_, until)| *until > now);
315 if inner.entries.len() >= CACHE_MAX_ENTRIES {
316 inner.entries.clear();
317 }
318 }
319 inner.entries.insert(key, (claims, valid_until));
320 }
321
322 fn mark_first_sight(&self, sub: &str, issuer: &str) -> bool {
324 let mut inner = self.inner.lock().expect("validation cache poisoned");
325 if inner.first_sight.len() >= CACHE_MAX_ENTRIES {
326 inner.first_sight.clear();
327 }
328 inner
329 .first_sight
330 .insert((sub.to_string(), issuer.to_string()))
331 }
332}
333
334#[derive(Clone)]
336pub struct AuthState {
337 pub oidc_client: OidcClient,
339 pub resource_server: ResourceServerClient,
341 pub client_config: OidcClientConfig,
343 pub config: AuthConfig,
345 pub pkce_manager: PkceCookieManager,
347 pub session_manager: Arc<WebSessionManager>,
349 pub issuer_fallbacks: Vec<String>,
355 pub cedar_authorizer: Option<Arc<pep::cedar::CedarAuthorizer>>,
357 validation_cache: Arc<ValidationCache>,
360 pub identity_registry: IdentityRegistry,
364}
365
366impl AuthState {
367 pub fn new(config: AuthConfig) -> Self {
369 Self::with_cedar(config, None)
370 }
371
372 pub fn with_cedar(
374 config: AuthConfig,
375 cedar_authorizer: Option<Arc<pep::cedar::CedarAuthorizer>>,
376 ) -> Self {
377 let pkce_manager = PkceCookieManager::new(
378 config.pkce_cookie_secret.as_bytes(),
379 "trustee_pkce_state",
380 StdDuration::from_secs(600),
381 );
382
383 let session_manager = Arc::new(WebSessionManager::new(
384 OidcClient::new(),
385 config.issuer_url.clone(),
386 config.client_id.clone(),
387 config.client_secret.clone(),
388 config.scope.clone(),
389 ));
390
391 Self {
392 oidc_client: OidcClient::new(),
393 resource_server: ResourceServerClient::new(),
394 client_config: config.oidc_client_config(),
395 pkce_manager,
396 session_manager,
397 config,
398 cedar_authorizer,
399 issuer_fallbacks: Vec::new(),
400 validation_cache: Arc::new(ValidationCache::new()),
401 identity_registry: IdentityRegistry::from_default_home(),
402 }
403 }
404
405 pub fn resolve_user_key(&self, claims: &JwtClaims) -> String {
411 if PrincipalKind::from_role(claim_role(claims).as_deref()) == PrincipalKind::Agent {
412 return claims.sub.clone();
413 }
414 self.identity_registry
415 .pin(claims.sub.as_str(), &jwt_user_key(claims))
416 }
417
418 pub fn with_issuer_fallbacks(mut self, issuers: Vec<String>) -> Self {
420 self.issuer_fallbacks = issuers;
421 self
422 }
423
424 pub fn is_dev_mode(&self) -> bool {
426 self.config.dev_config.local_dev_mode
427 }
428
429 pub async fn validate_token(&self, token: &str) -> anyhow::Result<JwtClaims> {
431 self.validate_token_on(&self.config.issuer_url, token).await
432 }
433
434 pub async fn validate_token_flexible(&self, token: &str) -> anyhow::Result<JwtClaims> {
447 let key = validation_cache_key(token);
448 let now = unix_now();
449 if let Some(cached) = self.validation_cache.get(&key, now) {
450 return Ok(cached);
451 }
452 let claims = self.validate_token_flexible_uncached(token).await?;
453 self.validation_cache.put(key, claims.clone(), now);
454 Ok(claims)
455 }
456
457 async fn validate_token_flexible_uncached(&self, token: &str) -> anyhow::Result<JwtClaims> {
458 match self.validate_token(token).await {
459 Ok(claims) => Ok(claims),
460 Err(primary) => {
461 let mut last = primary;
462 for si in &self.issuer_fallbacks {
463 if *si == self.config.issuer_url {
464 continue; }
466 match self.validate_token_on(si, token).await {
467 Ok(claims) => {
468 let first = self.validation_cache.mark_first_sight(&claims.sub, si);
472 let who = claims.preferred_username.as_deref().unwrap_or(&claims.sub);
473 if first {
474 tracing::info!(
475 "service principal {} authenticating via service-issuer fallback (iss={})",
476 who,
477 si
478 );
479 } else {
480 tracing::debug!(
481 "token validated via service-issuer fallback (iss={}) — principal {}",
482 si,
483 who
484 );
485 }
486 return Ok(claims);
487 }
488 Err(e) => last = e,
489 }
490 }
491 tracing::warn!(
494 "token validation failed — primary: {last}; all service-issuer fallbacks exhausted"
495 );
496 Err(anyhow::anyhow!(
497 "primary: {last}; all service-issuer fallbacks exhausted"
498 ))
499 }
500 }
501 }
502
503 pub async fn validate_token_on(
506 &self,
507 issuer_url: &str,
508 token: &str,
509 ) -> anyhow::Result<JwtClaims> {
510 let mut claims = self
511 .resource_server
512 .validate_jwt_with_options(
513 token,
514 issuer_url,
515 &self.config.client_id,
516 &self.config.validation_options,
517 )
518 .await
519 .map_err(|e| anyhow::anyhow!("Token validation failed: {}", e))?;
520
521 if let Err(e) = self
526 .resource_server
527 .enrich_claims_with_userinfo(&mut claims, token, issuer_url, None)
528 .await
529 {
530 tracing::error!(
531 "userinfo enrichment FAILED for sub {}: {} — principal carries NO role/groups; \
532 with Cedar enabled every request will be DENIED until enrichment succeeds",
533 claims.sub,
534 e
535 );
536 }
537
538 if claims.name.is_none() || claims.email.is_none() {
541 self.fill_userinfo_fields(&mut claims, token, issuer_url)
542 .await;
543 }
544
545 Ok(claims)
546 }
547
548 async fn fill_userinfo_fields(&self, claims: &mut JwtClaims, token: &str, issuer_url: &str) {
551 let userinfo_url = format!("{}/userinfo", issuer_url.trim_end_matches('/'));
556
557 let client = reqwest::Client::new();
558 let resp = client
559 .get(&userinfo_url)
560 .header("Authorization", format!("Bearer {}", token))
561 .header("Accept", "application/json")
562 .send()
563 .await;
564
565 let Ok(resp) = resp else {
566 tracing::debug!("Userinfo request failed for name/email enrichment");
567 return;
568 };
569
570 if !resp.status().is_success() {
571 tracing::debug!(
572 "Userinfo returned {} for name/email enrichment",
573 resp.status()
574 );
575 return;
576 }
577
578 let Ok(userinfo): Result<serde_json::Map<String, serde_json::Value>, _> = resp.json().await
579 else {
580 return;
581 };
582
583 tracing::debug!("Userinfo keys: {:?}", userinfo.keys().collect::<Vec<_>>());
584
585 if claims.name.is_none() {
586 if let Some(name) = userinfo.get("name").and_then(|v| v.as_str()) {
587 claims.name = Some(name.to_string());
588 }
589 }
590 if claims.email.is_none() {
591 if let Some(email) = userinfo.get("email").and_then(|v| v.as_str()) {
592 claims.email = Some(email.to_string());
593 }
594 }
595 if claims.preferred_username.is_none() {
596 if let Some(uname) = userinfo.get("preferred_username").and_then(|v| v.as_str()) {
597 claims.preferred_username = Some(uname.to_string());
598 }
599 }
600 }
601
602 fn check_cedar_authorized(&self, claims: &JwtClaims, action: &str) -> Result<(), ()> {
607 let Some(ref authorizer) = self.cedar_authorizer else {
608 return Ok(()); };
610
611 let principal_entity = match pep::cedar::build_principal_entity(claims) {
613 Ok(e) => e,
614 Err(e) => {
615 tracing::error!("Cedar: failed to build principal entity: {}", e);
616 return Err(());
617 }
618 };
619
620 let mut entities_vec = vec![principal_entity];
622
623 let app_uid = match EntityUid::from_str(r#"TrusteeApp::"default""#) {
625 Ok(uid) => uid,
626 Err(e) => {
627 tracing::error!("Cedar: failed to build TrusteeApp uid: {}", e);
628 return Err(());
629 }
630 };
631 let app_entity = match cedar_policy::Entity::new(
632 app_uid,
633 std::collections::HashMap::new(),
634 std::collections::HashSet::new(),
635 ) {
636 Ok(e) => e,
637 Err(e) => {
638 tracing::error!("Cedar: failed to build TrusteeApp entity: {}", e);
639 return Err(());
640 }
641 };
642 entities_vec.push(app_entity);
643
644 let entities = match Entities::from_entities(entities_vec, None) {
645 Ok(e) => e,
646 Err(e) => {
647 tracing::error!("Cedar: failed to build entities set: {}", e);
648 return Err(());
649 }
650 };
651
652 let principal_uid = match pep::cedar::build_principal_uid(claims) {
654 Ok(uid) => uid,
655 Err(e) => {
656 tracing::error!("Cedar: failed to build principal uid: {}", e);
657 return Err(());
658 }
659 };
660
661 let action_uid = match EntityUid::from_str(&format!("Action::\"{action}\"")) {
662 Ok(uid) => uid,
663 Err(e) => {
664 tracing::error!("Cedar: failed to build action uid: {}", e);
665 return Err(());
666 }
667 };
668
669 let resource_uid = match EntityUid::from_str(r#"TrusteeApp::"default""#) {
670 Ok(uid) => uid,
671 Err(e) => {
672 tracing::error!("Cedar: failed to build resource uid: {}", e);
673 return Err(());
674 }
675 };
676
677 let request = match Request::new(
678 principal_uid,
679 action_uid,
680 resource_uid,
681 Context::empty(),
682 None,
683 ) {
684 Ok(r) => r,
685 Err(e) => {
686 tracing::error!("Cedar: failed to build request: {}", e);
687 return Err(());
688 }
689 };
690
691 let response = authorizer.is_allowed_with_entities(&request, &entities);
692
693 if response.allowed() {
694 tracing::debug!(
695 "Cedar: authorized user {} (sub={})",
696 claims.email.as_deref().unwrap_or("unknown"),
697 claims.sub
698 );
699 Ok(())
700 } else {
701 tracing::warn!(
702 "Cedar: DENIED user {} (sub={}) — matched policies: {:?}, errors: {:?}",
703 claims.email.as_deref().unwrap_or("unknown"),
704 claims.sub,
705 response.matched_policies(),
706 response.errors()
707 );
708 Err(())
709 }
710 }
711}
712
713pub mod actions {
724 pub const LIST_MODELS: &str = "ListModels";
725 pub const LIST_SESSIONS: &str = "ListSessions";
726 pub const VIEW_SESSION: &str = "ViewSession";
727 pub const VIEW_HISTORY: &str = "ViewHistory";
728 pub const CREATE_SESSION: &str = "CreateSession";
729 pub const COMMAND_SESSION: &str = "CommandSession";
730 pub const CANCEL_SESSION: &str = "CancelSession";
731 pub const HANDOFF_SESSION: &str = "HandoffSession";
732 pub const RESUME_SESSION: &str = "ResumeSession";
733 pub const UPDATE_SESSION: &str = "UpdateSession";
734 pub const DELETE_SESSION: &str = "DeleteSession";
735 pub const VIEW_MCP_CREDENTIALS: &str = "ViewMcpCredentials";
736 pub const UPDATE_MCP_CREDENTIALS: &str = "UpdateMcpCredentials";
737}
738
739#[derive(Debug, Clone, Copy, PartialEq, Eq)]
743pub enum PrincipalKind {
744 Human,
745 Agent,
746}
747
748impl PrincipalKind {
749 pub fn from_role(role: Option<&str>) -> Self {
753 match role {
754 Some("agent") => Self::Agent,
755 _ => Self::Human,
756 }
757 }
758}
759
760fn claim_role(claims: &JwtClaims) -> Option<String> {
766 match claims.extra.get("role") {
767 Some(serde_json::Value::String(s)) => Some(s.clone()),
768 Some(serde_json::Value::Array(arr)) => arr
769 .iter()
770 .filter_map(|v| v.as_str())
771 .next()
772 .map(|s| s.to_string()),
773 _ => None,
774 }
775}
776
777#[derive(Debug, Clone)]
779pub struct AuthUser {
780 pub sub: String,
781 pub email: Option<String>,
782 pub name: Option<String>,
783 pub username: Option<String>,
784 pub is_dev: bool,
785 pub role: Option<String>,
787 pub kind: PrincipalKind,
789}
790
791impl From<JwtClaims> for AuthUser {
792 fn from(claims: JwtClaims) -> Self {
793 let role = claim_role(&claims);
794 let kind = PrincipalKind::from_role(role.as_deref());
795 Self {
796 sub: claims.sub,
797 email: claims.email,
798 name: claims.name,
799 username: claims.preferred_username,
800 is_dev: false,
801 role,
802 kind,
803 }
804 }
805}
806
807const SESSION_COOKIE_MAX_AGE: StdDuration = StdDuration::from_secs(3600);
809
810fn jwt_user_key(claims: &JwtClaims) -> String {
827 if PrincipalKind::from_role(claim_role(claims).as_deref()) == PrincipalKind::Agent {
828 return claims.sub.clone();
829 }
830 match claims.preferred_username.as_deref() {
831 Some(u) if !u.trim().is_empty() => u.trim().to_string(),
832 _ => {
833 tracing::debug!(
834 "user_key: preferred_username missing/empty for sub {} — falling back to sub",
835 claims.sub
836 );
837 claims.sub.clone()
838 }
839 }
840}
841
842fn dev_user_key(token: &str) -> Option<String> {
850 if let Some(name) = token.strip_prefix("dev:agent:") {
851 let name = name.trim();
852 if name.is_empty() || name.contains(':') {
853 return None;
854 }
855 return Some(format!("agent-{name}"));
856 }
857 let parts: Vec<&str> = token.splitn(4, ':').collect();
858 if parts.len() >= 4 {
859 Some(format!("dev:{}", parts[1]))
860 } else {
861 None
862 }
863}
864
865pub(crate) fn dispatch_allowed(kind: PrincipalKind, role: Option<&str>) -> bool {
877 kind == PrincipalKind::Human
878 && matches!(role, Some("admin") | Some("user") | Some("developer"))
879}
880
881pub async fn check_dispatch_admin(
896 auth: &Option<Arc<AuthState>>,
897 headers: &axum::http::HeaderMap,
898) -> Result<(), StatusCode> {
899 let Some(auth) = auth.as_ref() else {
900 return Ok(()); };
902 let Some(token) = headers
903 .get(header::AUTHORIZATION)
904 .and_then(|v| v.to_str().ok())
905 .and_then(|v| v.strip_prefix("Bearer "))
906 .map(|s| s.to_string())
907 else {
908 return Err(StatusCode::UNAUTHORIZED);
909 };
910 if token.starts_with("dev:") {
911 tracing::warn!("xagent dispatch rejected: dev tokens cannot dispatch agents");
912 return Err(StatusCode::FORBIDDEN);
913 }
914 let claims = auth.validate_token_flexible(&token).await.map_err(|e| {
915 tracing::warn!("xagent dispatch: caller token validation failed: {}", e);
916 StatusCode::UNAUTHORIZED
917 })?;
918 let user = AuthUser::from(claims);
919 if !dispatch_allowed(user.kind, user.role.as_deref()) {
920 tracing::warn!(
921 "xagent dispatch rejected: principal kind={:?} role={:?} — admin required",
922 user.kind,
923 user.role
924 );
925 return Err(StatusCode::FORBIDDEN);
926 }
927 Ok(())
928}
929
930impl AuthState {
931 pub async fn exchange_agent_token(
943 &self,
944 issuer_url: &str,
945 service_token: &str,
946 ) -> Result<(String, u64), StatusCode> {
947 let tr = self
948 .oidc_client
949 .exchange_token(
950 issuer_url,
951 &self.config.client_id,
952 None, service_token,
954 &self.config.client_id,
955 Some("openid groups"),
956 )
957 .await
958 .map_err(|e| {
959 tracing::error!(
960 "xagent dispatch: agent token exchange FAILED at {issuer_url}: {} — impersonation unavailable",
961 e
962 );
963 StatusCode::BAD_GATEWAY
964 })?;
965 Ok((tr.access_token, tr.expires_in.unwrap_or(900)))
966 }
967}
968
969pub async fn check_auth(
990 auth: &Option<Arc<AuthState>>,
991 headers: &axum::http::HeaderMap,
992 action: &str,
993) -> Result<(Option<String>, String), StatusCode> {
994 let Some(auth) = auth.as_ref() else {
995 return Ok((None, "default".to_string())); };
997
998 if let Some(token) = headers
1000 .get(header::AUTHORIZATION)
1001 .and_then(|v| v.to_str().ok())
1002 .and_then(|v| v.strip_prefix("Bearer "))
1003 .map(|s| s.to_string())
1004 {
1005 if token.starts_with("dev:") {
1007 if !auth.config.dev_config.local_dev_mode {
1008 tracing::warn!("Dev token presented but dev mode is disabled — rejecting");
1009 return Err(StatusCode::UNAUTHORIZED);
1010 }
1011 return match dev_user_key(&token) {
1012 Some(key) => Ok((None, key)),
1013 None => Err(StatusCode::UNAUTHORIZED),
1014 };
1015 }
1016
1017 return match auth.validate_token_flexible(&token).await {
1018 Ok(claims) => {
1019 if auth.check_cedar_authorized(&claims, action).is_err() {
1020 return Err(StatusCode::FORBIDDEN);
1021 }
1022 Ok((None, auth.resolve_user_key(&claims)))
1023 }
1024 Err(e) => {
1025 tracing::warn!("Bearer token validation failed: {}", e);
1026 Err(StatusCode::UNAUTHORIZED)
1027 }
1028 };
1029 }
1030
1031 let session_id = headers
1033 .get(header::COOKIE)
1034 .and_then(|v| v.to_str().ok())
1035 .and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
1036
1037 let Some(session_id) = session_id else {
1038 tracing::warn!("No auth token found in request");
1039 return Err(StatusCode::UNAUTHORIZED);
1040 };
1041
1042 if session_id.starts_with("dev:") {
1044 if !auth.config.dev_config.local_dev_mode {
1045 tracing::warn!("Dev cookie presented but dev mode is disabled — rejecting");
1046 return Err(StatusCode::UNAUTHORIZED);
1047 }
1048 return match dev_user_key(&session_id) {
1049 Some(key) => Ok((None, key)),
1050 None => Err(StatusCode::UNAUTHORIZED),
1051 };
1052 }
1053
1054 match auth.session_manager.get_token(&session_id).await {
1056 Ok(access_token) => match auth.validate_token(&access_token).await {
1057 Ok(claims) => {
1058 if auth.check_cedar_authorized(&claims, action).is_err() {
1060 return Err(StatusCode::FORBIDDEN);
1061 }
1062 let secure = auth.client_config.redirect_uri.starts_with("https");
1064 let cookie = create_auth_cookie(
1065 &auth.config.cookie_name,
1066 &session_id,
1067 SESSION_COOKIE_MAX_AGE,
1068 secure,
1069 );
1070 Ok((Some(cookie.to_string()), auth.resolve_user_key(&claims)))
1071 }
1072 Err(e) => {
1073 tracing::warn!(
1076 "Session token validation failed: {} — attempting force-refresh",
1077 e
1078 );
1079 match auth.session_manager.force_refresh(&session_id).await {
1080 Ok(new_token) => match auth.validate_token(&new_token).await {
1081 Ok(claims) => {
1082 if auth.check_cedar_authorized(&claims, action).is_err() {
1084 return Err(StatusCode::FORBIDDEN);
1085 }
1086 let secure = auth.client_config.redirect_uri.starts_with("https");
1087 let cookie = create_auth_cookie(
1088 &auth.config.cookie_name,
1089 &session_id,
1090 SESSION_COOKIE_MAX_AGE,
1091 secure,
1092 );
1093 Ok((Some(cookie.to_string()), auth.resolve_user_key(&claims)))
1094 }
1095 Err(e2) => {
1096 tracing::warn!(
1097 "Session token still invalid after force-refresh: {}",
1098 e2
1099 );
1100 Err(StatusCode::UNAUTHORIZED)
1101 }
1102 },
1103 Err(e2) => {
1104 tracing::warn!("Force-refresh failed: {}", e2);
1105 Err(StatusCode::UNAUTHORIZED)
1106 }
1107 }
1108 }
1109 },
1110 Err(e) => {
1111 tracing::warn!("Session lookup/refresh failed: {}", e);
1112 Err(StatusCode::UNAUTHORIZED)
1113 }
1114 }
1115}
1116
1117async fn resolve_access_token(
1123 auth: &AuthState,
1124 headers: &axum::http::HeaderMap,
1125) -> Result<String, StatusCode> {
1126 if let Some(token) = headers
1128 .get(header::AUTHORIZATION)
1129 .and_then(|v| v.to_str().ok())
1130 .and_then(|v| v.strip_prefix("Bearer "))
1131 .map(|s| s.to_string())
1132 {
1133 return Ok(token);
1134 }
1135
1136 let session_id = headers
1138 .get(header::COOKIE)
1139 .and_then(|v| v.to_str().ok())
1140 .and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
1141
1142 match session_id {
1143 Some(sid) if sid.starts_with("dev:") => {
1144 if !auth.config.dev_config.local_dev_mode {
1145 tracing::warn!(
1146 "Dev cookie in resolve_access_token but dev mode is disabled — rejecting"
1147 );
1148 Err(StatusCode::UNAUTHORIZED)
1149 } else {
1150 Ok(sid)
1151 }
1152 }
1153 Some(sid) => auth.session_manager.get_token(&sid).await.map_err(|e| {
1154 tracing::warn!("Failed to resolve session token: {}", e);
1155 StatusCode::UNAUTHORIZED
1156 }),
1157 None => Err(StatusCode::UNAUTHORIZED),
1158 }
1159}
1160
1161fn extract_token_from_cookies(cookie_header: &str, cookie_name: &str) -> Option<String> {
1163 for cookie in cookie_header.split(';') {
1164 let cookie = cookie.trim();
1165 if let Some(value) = cookie.strip_prefix(&format!("{}=", cookie_name)) {
1166 return Some(value.to_string());
1167 }
1168 }
1169 None
1170}
1171
1172pub fn auth_routes() -> axum::Router<crate::ServerState> {
1178 axum::Router::new()
1179 .route("/login", axum::routing::get(login_handler))
1180 .route("/callback", axum::routing::get(callback_handler))
1181 .route("/me", axum::routing::get(me_handler))
1182 .route("/logout", axum::routing::post(logout_handler))
1183 .route("/mcp/login", axum::routing::get(mcp_login_handler))
1184 .route("/mcp/callback", axum::routing::get(mcp_callback_handler))
1185 .route("/mcp/status", axum::routing::get(mcp_status_handler))
1186 .route("/mcp/logout", axum::routing::post(mcp_logout_handler))
1187}
1188
1189#[derive(Debug, Deserialize)]
1191pub struct CallbackQuery {
1192 pub code: Option<String>,
1193 pub state: Option<String>,
1194 pub error: Option<String>,
1195 pub error_description: Option<String>,
1196}
1197
1198async fn login_handler(State(state): State<crate::ServerState>) -> Result<Response, AuthError> {
1200 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1201
1202 if auth.is_dev_mode() {
1204 tracing::info!("Dev mode: creating dev session");
1205 let dev = &auth.config.dev_config;
1206 let dev_token = format!(
1207 "dev:{}:{}:{}",
1208 dev.local_dev_email.as_deref().unwrap_or("dev@localhost"),
1209 dev.local_dev_name.as_deref().unwrap_or("Dev User"),
1210 dev.local_dev_username.as_deref().unwrap_or("dev")
1211 );
1212 let cookie = create_auth_cookie(
1213 &auth.config.cookie_name,
1214 &dev_token,
1215 StdDuration::from_secs(86400),
1216 false,
1217 );
1218 return Ok(Response::builder()
1219 .status(StatusCode::FOUND)
1220 .header(header::LOCATION, "/")
1221 .header(header::SET_COOKIE, cookie.to_string())
1222 .body(Body::empty())
1223 .unwrap());
1224 }
1225
1226 let pkce_session = auth.pkce_manager.create();
1228 let challenge = OidcClient::generate_code_challenge(&pkce_session.verifier);
1229
1230 let auth_url = auth
1231 .oidc_client
1232 .build_authorization_url(&auth.client_config, &pkce_session.state, Some(&challenge))
1233 .await
1234 .map_err(|e| AuthError::OidcError(e.to_string()))?;
1235
1236 let secure = auth.client_config.redirect_uri.starts_with("https");
1240 let pkce_cookie = Cookie::build((
1241 auth.pkce_manager.cookie_name().to_string(),
1242 pkce_session.cookie_value,
1243 ))
1244 .path("/")
1245 .http_only(true)
1246 .same_site(SameSite::Lax)
1247 .secure(secure)
1248 .max_age(TimeDuration::seconds(
1249 auth.pkce_manager.ttl().as_secs() as i64
1250 ))
1251 .build();
1252
1253 Ok(Response::builder()
1254 .status(StatusCode::TEMPORARY_REDIRECT)
1255 .header(header::LOCATION, &auth_url)
1256 .header(header::SET_COOKIE, pkce_cookie.to_string())
1257 .body(Body::empty())
1258 .unwrap())
1259}
1260
1261async fn callback_handler(
1263 State(state): State<crate::ServerState>,
1264 Query(query): Query<CallbackQuery>,
1265 headers: axum::http::HeaderMap,
1266) -> Result<Response, AuthError> {
1267 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1268
1269 if let Some(error) = query.error {
1271 let desc = query.error_description.unwrap_or_default();
1272 tracing::error!("OIDC error: {} - {}", error, desc);
1273 return Ok(Redirect::temporary(&format!(
1274 "/?error={}&error_description={}",
1275 urlencoding::encode(&error),
1276 urlencoding::encode(&desc)
1277 ))
1278 .into_response());
1279 }
1280
1281 let code = query.code.ok_or(AuthError::MissingCode)?;
1282 let oauth_state = query.state.ok_or(AuthError::MissingState)?;
1283
1284 let cookie_header = headers
1286 .get(header::COOKIE)
1287 .and_then(|v| v.to_str().ok())
1288 .unwrap_or("");
1289 let pkce_value = extract_token_from_cookies(cookie_header, auth.pkce_manager.cookie_name())
1290 .ok_or(AuthError::InvalidState)?;
1291
1292 let verifier = auth
1294 .pkce_manager
1295 .verify(&pkce_value, &oauth_state)
1296 .ok_or(AuthError::InvalidState)?;
1297
1298 tracing::info!("Exchanging authorization code for tokens");
1300 let token_response = auth
1301 .oidc_client
1302 .exchange_code_for_tokens(&auth.client_config, &code, Some(&verifier))
1303 .await
1304 .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
1305
1306 let session_id = auth
1307 .session_manager
1308 .create_session(&token_response)
1309 .await
1310 .map_err(|e| AuthError::TokenExchangeFailed(format!("Session creation failed: {}", e)))?;
1311
1312 let max_age = SESSION_COOKIE_MAX_AGE;
1315
1316 let secure = auth.client_config.redirect_uri.starts_with("https");
1318 let cookie = create_auth_cookie(&auth.config.cookie_name, &session_id, max_age, secure);
1319
1320 let clear_pkce = Cookie::build((auth.pkce_manager.cookie_name().to_string(), ""))
1322 .path("/")
1323 .http_only(true)
1324 .same_site(SameSite::Lax)
1325 .max_age(TimeDuration::seconds(-1))
1326 .build();
1327
1328 tracing::info!("Authentication successful, redirecting to /");
1329
1330 Ok(Response::builder()
1331 .status(StatusCode::FOUND)
1332 .header(header::LOCATION, "/")
1333 .header(header::SET_COOKIE, cookie.to_string())
1334 .header(header::SET_COOKIE, clear_pkce.to_string())
1335 .body(Body::empty())
1336 .unwrap())
1337}
1338
1339async fn me_handler(
1341 State(state): State<crate::ServerState>,
1342 headers: axum::http::HeaderMap,
1343) -> Response {
1344 let Some(ref auth) = state.auth else {
1345 return axum::Json(serde_json::json!({
1347 "authenticated": true,
1348 "auth_enabled": false
1349 }))
1350 .into_response();
1351 };
1352
1353 let cookie_header = headers
1354 .get(header::COOKIE)
1355 .and_then(|v| v.to_str().ok())
1356 .unwrap_or("");
1357
1358 let bearer = headers
1360 .get(header::AUTHORIZATION)
1361 .and_then(|v| v.to_str().ok())
1362 .and_then(|v| v.strip_prefix("Bearer "))
1363 .map(String::from);
1364
1365 let token = bearer
1366 .clone()
1367 .or_else(|| extract_token_from_cookies(cookie_header, &auth.config.cookie_name));
1368
1369 let Some(cookie_value) = token else {
1370 return axum::Json(serde_json::json!({
1371 "authenticated": false,
1372 "auth_enabled": true
1373 }))
1374 .into_response();
1375 };
1376
1377 if cookie_value.starts_with("dev:") && auth.config.dev_config.local_dev_mode {
1380 let parts: Vec<&str> = cookie_value.splitn(4, ':').collect();
1381 if parts.len() >= 4 {
1382 return axum::Json(serde_json::json!({
1383 "authenticated": true,
1384 "auth_enabled": true,
1385 "email": parts[1],
1386 "name": parts[2],
1387 "username": parts[3],
1388 "dev_mode": true
1389 }))
1390 .into_response();
1391 }
1392 }
1393
1394 let access_token = if bearer.is_some() {
1396 cookie_value
1398 } else {
1399 match auth.session_manager.get_token(&cookie_value).await {
1401 Ok(token) => token,
1402 Err(e) => {
1403 tracing::debug!("Session token resolution failed for /auth/me: {}", e);
1404 return axum::Json(serde_json::json!({
1405 "authenticated": false,
1406 "auth_enabled": true
1407 }))
1408 .into_response();
1409 }
1410 }
1411 };
1412
1413 match auth.validate_token(&access_token).await {
1415 Ok(claims) => axum::Json(serde_json::json!({
1416 "authenticated": true,
1417 "auth_enabled": true,
1418 "sub": claims.sub,
1419 "email": claims.email,
1420 "name": claims.name,
1421 "username": claims.preferred_username,
1422 "dev_mode": false
1423 }))
1424 .into_response(),
1425 Err(e) => {
1426 tracing::debug!("Token validation failed for /auth/me: {}", e);
1427 axum::Json(serde_json::json!({
1428 "authenticated": false,
1429 "auth_enabled": true
1430 }))
1431 .into_response()
1432 }
1433 }
1434}
1435
1436async fn logout_handler(
1438 State(state): State<crate::ServerState>,
1439 headers: axum::http::HeaderMap,
1440) -> Response {
1441 let cookie_name = state
1442 .auth
1443 .as_ref()
1444 .map(|a| a.config.cookie_name.as_str())
1445 .unwrap_or("trustee_token");
1446
1447 if let Some(ref auth) = state.auth {
1449 if let Some(cookie_header) = headers.get(header::COOKIE).and_then(|v| v.to_str().ok()) {
1450 if let Some(session_id) = extract_token_from_cookies(cookie_header, cookie_name) {
1451 if !session_id.starts_with("dev:") {
1452 let _ = auth.session_manager.destroy_session(&session_id);
1453 }
1454 }
1455 }
1456 }
1457
1458 let cookie = Cookie::build((cookie_name.to_string(), ""))
1459 .path("/")
1460 .http_only(true)
1461 .same_site(SameSite::Lax)
1462 .max_age(TimeDuration::seconds(-1))
1463 .build();
1464
1465 Response::builder()
1466 .status(StatusCode::FOUND)
1467 .header(header::LOCATION, "/")
1468 .header(header::SET_COOKIE, cookie.to_string())
1469 .body(Body::empty())
1470 .unwrap()
1471}
1472
1473#[derive(Debug, Deserialize)]
1479pub struct McpLoginQuery {
1480 pub cred: String,
1481}
1482
1483#[derive(Debug, Deserialize)]
1485pub struct McpCallbackQuery {
1486 pub code: Option<String>,
1487 pub state: Option<String>,
1488 pub error: Option<String>,
1489 pub error_description: Option<String>,
1490}
1491
1492async fn mcp_login_handler(
1497 State(state): State<crate::ServerState>,
1498 Query(query): Query<McpLoginQuery>,
1499 headers: axum::http::HeaderMap,
1500) -> Result<Response, AuthError> {
1501 let (_cookie, _user_key) = crate::auth::check_auth(
1503 &state.auth,
1504 &headers,
1505 crate::auth::actions::UPDATE_MCP_CREDENTIALS,
1506 )
1507 .await
1508 .map_err(|_| AuthError::AuthNotConfigured)?;
1509
1510 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1511
1512 let cred_config = load_mcp_credential(&state, &query.cred).await?;
1514
1515 let (issuer_url, client_id, client_secret, scope) = match &cred_config {
1516 McpCredentialInfo::WebInteractive {
1517 issuer_url,
1518 client_id,
1519 client_secret,
1520 scope,
1521 } => (
1522 issuer_url.clone(),
1523 client_id.clone(),
1524 client_secret.clone(),
1525 scope.clone(),
1526 ),
1527 _ => {
1528 return Ok(Redirect::temporary(&format!(
1529 "/?mcp_error={}",
1530 urlencoding::encode(&format!(
1531 "Credential '{}' is not web-interactive type",
1532 query.cred
1533 ))
1534 ))
1535 .into_response());
1536 }
1537 };
1538
1539 let oidc_client = OidcClient::new();
1541 let verifier = OidcClient::generate_code_verifier();
1542 let challenge = OidcClient::generate_code_challenge(&verifier);
1543 let oauth_state = OidcClient::generate_state();
1544
1545 let mcp_redirect_uri = format!(
1547 "{}/auth/mcp/callback",
1548 auth.client_config
1549 .redirect_uri
1550 .trim_end_matches('/')
1551 .trim_end_matches("/auth/callback")
1552 );
1553
1554 let mcp_client_config = OidcClientConfig {
1555 issuer_url: issuer_url.clone(),
1556 client_id: client_id.clone(),
1557 client_secret: client_secret.clone(),
1558 redirect_uri: mcp_redirect_uri.clone(),
1559 scope: scope.clone(),
1560 code_challenge_method: "S256".to_string(),
1561 };
1562
1563 let auth_url = oidc_client
1565 .build_authorization_url(&mcp_client_config, &oauth_state, Some(&challenge))
1566 .await
1567 .map_err(|e| AuthError::OidcError(e.to_string()))?;
1568
1569 mcp_pkce()
1571 .insert(oauth_state.clone(), verifier.clone(), query.cred.clone())
1572 .await;
1573
1574 tracing::info!(
1575 "Initiating MCP browser login for credential '{}' (issuer={})",
1576 query.cred,
1577 issuer_url
1578 );
1579
1580 Ok(Response::builder()
1581 .status(StatusCode::TEMPORARY_REDIRECT)
1582 .header(header::LOCATION, &auth_url)
1583 .body(Body::empty())
1584 .unwrap())
1585}
1586
1587async fn mcp_callback_handler(
1589 State(state): State<crate::ServerState>,
1590 Query(query): Query<McpCallbackQuery>,
1591 _headers: axum::http::HeaderMap,
1592) -> Result<Response, AuthError> {
1593 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1594
1595 if let Some(error) = query.error {
1597 let desc = query.error_description.unwrap_or_default();
1598 tracing::error!("MCP OIDC error: {} - {}", error, desc);
1599 return Ok(Redirect::temporary(&format!(
1600 "/?mcp_error={}&error_description={}",
1601 urlencoding::encode(&error),
1602 urlencoding::encode(&desc)
1603 ))
1604 .into_response());
1605 }
1606
1607 let code = query.code.ok_or(AuthError::MissingCode)?;
1608 let oauth_state = query.state.ok_or(AuthError::MissingState)?;
1609
1610 let pkce_data = mcp_pkce()
1612 .take(&oauth_state)
1613 .await
1614 .ok_or(AuthError::InvalidState)?;
1615
1616 let verifier = pkce_data.verifier;
1617 let cred_name = &pkce_data.cred_name;
1618
1619 let cred_config = load_mcp_credential(&state, cred_name).await?;
1621
1622 let (issuer_url, client_id, client_secret, scope) = match &cred_config {
1623 McpCredentialInfo::WebInteractive {
1624 issuer_url,
1625 client_id,
1626 client_secret,
1627 scope,
1628 } => (
1629 issuer_url.clone(),
1630 client_id.clone(),
1631 client_secret.clone(),
1632 scope.clone(),
1633 ),
1634 _ => {
1635 return Ok(Redirect::temporary(&format!(
1636 "/?mcp_error={}",
1637 urlencoding::encode("Credential is not web-interactive type")
1638 ))
1639 .into_response());
1640 }
1641 };
1642
1643 let mcp_redirect_uri = format!(
1645 "{}/auth/mcp/callback",
1646 auth.client_config
1647 .redirect_uri
1648 .trim_end_matches('/')
1649 .trim_end_matches("/auth/callback")
1650 );
1651
1652 let mcp_client_config = OidcClientConfig {
1653 issuer_url: issuer_url.clone(),
1654 client_id: client_id.clone(),
1655 client_secret: client_secret.clone(),
1656 redirect_uri: mcp_redirect_uri,
1657 scope: scope.clone(),
1658 code_challenge_method: "S256".to_string(),
1659 };
1660
1661 tracing::info!(
1663 "Exchanging MCP authorization code for tokens (credential={})",
1664 cred_name
1665 );
1666 let oidc_client = OidcClient::new();
1667 let token_response = oidc_client
1668 .exchange_code_for_tokens(&mcp_client_config, &code, Some(&verifier))
1669 .await
1670 .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
1671
1672 let expires_at = {
1674 let now = std::time::SystemTime::now()
1675 .duration_since(std::time::UNIX_EPOCH)
1676 .unwrap_or_default()
1677 .as_secs();
1678 let expires_epoch = now + token_response.expires_in.unwrap_or(900);
1679 let days = expires_epoch / 86400;
1680 let rem = expires_epoch % 86400;
1681 let h = rem / 3600;
1682 let m = (rem % 3600) / 60;
1683 let s = rem % 60;
1684 let z = days as i64 + 719468;
1685 let era = if z >= 0 { z } else { z - 146096 } / 146097;
1686 let doe = (z - era * 146097) as u64;
1687 let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
1688 let y = yoe as i64 + era * 400;
1689 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
1690 let mp = (5 * doy + 2) / 153;
1691 let d = doy - (153 * mp + 2) / 5 + 1;
1692 let mon = if mp < 10 { mp + 3 } else { mp - 9 };
1693 let yr = if mon <= 2 { y + 1 } else { y };
1694 format!("{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z", yr, mon, d, h, m, s)
1695 };
1696
1697 use pep::{FileTokenStore, StoredToken, TokenStore};
1699
1700 let stored = StoredToken::new(
1701 &token_response.access_token,
1702 token_response.refresh_token.clone(),
1703 "Bearer",
1704 &expires_at,
1705 token_response.scope.clone(),
1706 );
1707
1708 let agent_name = state
1709 .config_toml
1710 .as_ref()
1711 .and_then(|t| {
1712 toml::from_str::<toml::Value>(t).ok().and_then(|v| {
1713 v.get("agent")
1714 .and_then(|a| a.get("name"))
1715 .and_then(|n| n.as_str())
1716 .map(String::from)
1717 })
1718 })
1719 .unwrap_or_else(|| "trustee".to_string());
1720 let token_store = FileTokenStore::new(&agent_name);
1721
1722 if let Err(e) = token_store.save(cred_name, &stored) {
1723 tracing::error!("Failed to store MCP token: {}", e);
1724 return Ok(Redirect::temporary(&format!(
1725 "/?mcp_error={}",
1726 urlencoding::encode(&format!("Failed to store token: {}", e))
1727 ))
1728 .into_response());
1729 }
1730
1731 tracing::info!(
1732 "MCP authentication successful for credential '{}' (expires {})",
1733 cred_name,
1734 expires_at
1735 );
1736
1737 Ok(Response::builder()
1738 .status(StatusCode::FOUND)
1739 .header(
1740 header::LOCATION,
1741 format!("/?mcp_connected={}", urlencoding::encode(cred_name)),
1742 )
1743 .body(Body::empty())
1744 .unwrap())
1745}
1746
1747pub(crate) async fn mcp_status_handler(
1752 State(state): State<crate::ServerState>,
1753 headers: axum::http::HeaderMap,
1754) -> Response {
1755 use pep::{FileTokenStore, TokenStore};
1756
1757 let (_cookie, _user_key) = match crate::auth::check_auth(
1759 &state.auth,
1760 &headers,
1761 crate::auth::actions::VIEW_MCP_CREDENTIALS,
1762 )
1763 .await
1764 {
1765 Ok(result) => result,
1766 Err(code) => {
1767 return (code, Json(serde_json::json!({"error": "Unauthorized"}))).into_response()
1768 }
1769 };
1770
1771 let config_toml = match &state.config_toml {
1777 Some(t) => t.clone(),
1778 None => return (StatusCode::INTERNAL_SERVER_ERROR, "Config not loaded").into_response(),
1779 };
1780
1781 let mcp_config: toml::Value = match toml::from_str(&config_toml) {
1782 Ok(v) => v,
1783 Err(_) => return Json(serde_json::json!([])).into_response(),
1784 };
1785
1786 let agent_name = {
1789 let table: toml::Table = toml::from_str(&config_toml).unwrap_or_default();
1790 table
1791 .get("agent")
1792 .and_then(|a| a.get("name"))
1793 .and_then(|n| n.as_str())
1794 .unwrap_or("trustee")
1795 .to_string()
1796 };
1797 let token_store = FileTokenStore::new(&agent_name);
1798
1799 let servers = mcp_config
1801 .get("mcp")
1802 .and_then(|m| m.get("servers"))
1803 .and_then(|s| s.as_array());
1804 let credentials = mcp_config
1805 .get("mcp")
1806 .and_then(|m| m.get("credentials"))
1807 .and_then(|c| c.as_table());
1808
1809 let mut cred_servers: std::collections::HashMap<String, Vec<String>> =
1810 std::collections::HashMap::new();
1811 if let Some(servers) = servers {
1812 for server in servers {
1813 let name = server.get("name").and_then(|n| n.as_str()).unwrap_or("");
1814 let cred_ref = server
1815 .get("credentials")
1816 .and_then(|c| c.as_str())
1817 .unwrap_or("");
1818 if !cred_ref.is_empty() {
1819 cred_servers
1820 .entry(cred_ref.to_string())
1821 .or_default()
1822 .push(name.to_string());
1823 }
1824 }
1825 }
1826
1827 let mut result = Vec::new();
1828
1829 if let Some(creds) = credentials {
1830 for (cred_name, cred_config) in creds {
1831 let cred_type = cred_config
1832 .get("type")
1833 .and_then(|t| t.as_str())
1834 .unwrap_or("unknown");
1835 let servers_using = cred_servers.get(cred_name).cloned().unwrap_or_default();
1836
1837 if cred_type == "web-session" {
1838 let connected = state.auth.is_some();
1840 result.push(serde_json::json!({
1841 "credential": cred_name,
1842 "type": cred_type,
1843 "connected": connected,
1844 "servers": servers_using,
1845 }));
1846 } else if cred_type == "service-account" {
1847 let token = cred_config
1851 .get("service_token")
1852 .and_then(|t| t.as_str())
1853 .unwrap_or("");
1854 result.push(serde_json::json!({
1855 "credential": cred_name,
1856 "type": cred_type,
1857 "connected": !token.is_empty(),
1858 "servers": servers_using,
1859 }));
1860 } else if cred_type == "static" {
1861 let token = cred_config
1863 .get("token")
1864 .and_then(|t| t.as_str())
1865 .unwrap_or("");
1866 result.push(serde_json::json!({
1867 "credential": cred_name,
1868 "type": cred_type,
1869 "connected": !token.is_empty(),
1870 "servers": servers_using,
1871 }));
1872 } else if cred_type == "web-interactive" || cred_type == "interactive" {
1873 let status = match token_store.load(cred_name) {
1875 Ok(Some(token)) => {
1876 let expired = token.is_expired();
1877 serde_json::json!({
1878 "credential": cred_name,
1879 "type": cred_type,
1880 "connected": !expired,
1881 "expires_at": token.expires_at,
1882 "servers": servers_using,
1883 })
1884 }
1885 _ => serde_json::json!({
1886 "credential": cred_name,
1887 "type": cred_type,
1888 "connected": false,
1889 "servers": servers_using,
1890 }),
1891 };
1892 result.push(status);
1893 }
1894 }
1895 }
1896
1897 Json(serde_json::Value::Array(result)).into_response()
1898}
1899
1900async fn mcp_logout_handler(
1902 State(state): State<crate::ServerState>,
1903 Query(query): Query<McpLoginQuery>,
1904 headers: axum::http::HeaderMap,
1905) -> Response {
1906 use pep::{FileTokenStore, TokenStore};
1907
1908 let (_cookie, _user_key) = match crate::auth::check_auth(
1910 &state.auth,
1911 &headers,
1912 crate::auth::actions::UPDATE_MCP_CREDENTIALS,
1913 )
1914 .await
1915 {
1916 Ok(result) => result,
1917 Err(code) => {
1918 return (code, Json(serde_json::json!({"error": "Unauthorized"}))).into_response()
1919 }
1920 };
1921
1922 let agent_name = {
1925 let config_toml = state.config_toml.clone().unwrap_or_default();
1926 let table: toml::Table = toml::from_str(&config_toml).unwrap_or_default();
1927 table
1928 .get("agent")
1929 .and_then(|a| a.get("name"))
1930 .and_then(|n| n.as_str())
1931 .unwrap_or("trustee")
1932 .to_string()
1933 };
1934 let token_store = FileTokenStore::new(&agent_name);
1935
1936 match token_store.delete(&query.cred) {
1937 Ok(()) => {
1938 tracing::info!("Removed MCP credentials for '{}'", query.cred);
1939 Json(serde_json::json!({"success": true})).into_response()
1940 }
1941 Err(e) => {
1942 tracing::error!("Failed to remove MCP credentials: {}", e);
1943 (
1944 StatusCode::INTERNAL_SERVER_ERROR,
1945 Json(serde_json::json!({"error": e.to_string()})),
1946 )
1947 .into_response()
1948 }
1949 }
1950}
1951
1952struct McpPkceStore {
1959 entries: tokio::sync::Mutex<std::collections::HashMap<String, McpPkceEntry>>,
1960}
1961
1962struct McpPkceEntry {
1963 verifier: String,
1964 cred_name: String,
1965 created_at: std::time::Instant,
1966}
1967
1968impl McpPkceStore {
1969 fn new() -> Self {
1970 Self {
1971 entries: tokio::sync::Mutex::new(std::collections::HashMap::new()),
1972 }
1973 }
1974
1975 async fn insert(&self, state: String, verifier: String, cred_name: String) {
1977 let mut map = self.entries.lock().await;
1978 let cutoff = std::time::Instant::now() - std::time::Duration::from_secs(600);
1980 map.retain(|_, v| v.created_at > cutoff);
1981 map.insert(
1982 state,
1983 McpPkceEntry {
1984 verifier,
1985 cred_name,
1986 created_at: std::time::Instant::now(),
1987 },
1988 );
1989 }
1990
1991 async fn take(&self, state: &str) -> Option<McpPkceEntry> {
1993 let mut map = self.entries.lock().await;
1994 map.remove(state)
1995 }
1996}
1997
1998static MCP_PKCE: std::sync::OnceLock<McpPkceStore> = std::sync::OnceLock::new();
2000
2001fn mcp_pkce() -> &'static McpPkceStore {
2003 MCP_PKCE.get_or_init(McpPkceStore::new)
2004}
2005
2006enum McpCredentialInfo {
2008 WebInteractive {
2009 issuer_url: String,
2010 client_id: String,
2011 client_secret: Option<String>,
2012 scope: String,
2013 },
2014 Other(String),
2015}
2016
2017async fn load_mcp_credential(
2019 state: &crate::ServerState,
2020 cred_name: &str,
2021) -> Result<McpCredentialInfo, AuthError> {
2022 let config_toml = state
2023 .config_toml
2024 .clone()
2025 .ok_or(AuthError::AuthNotConfigured)?;
2026
2027 let config: toml::Value = toml::from_str(&config_toml)
2028 .map_err(|e| AuthError::OidcError(format!("Config parse error: {}", e)))?;
2029
2030 let cred = config
2031 .get("mcp")
2032 .and_then(|m| m.get("credentials"))
2033 .and_then(|c| c.as_table())
2034 .and_then(|c| c.get(cred_name))
2035 .ok_or_else(|| AuthError::OidcError(format!("Credential '{}' not found", cred_name)))?;
2036
2037 let cred_type = cred
2038 .get("type")
2039 .and_then(|t| t.as_str())
2040 .unwrap_or("unknown");
2041
2042 match cred_type {
2043 "web-interactive" => {
2044 let issuer_url = cred
2045 .get("issuer_url")
2046 .and_then(|v| v.as_str())
2047 .ok_or_else(|| AuthError::OidcError("Missing issuer_url".into()))?
2048 .to_string();
2049 let client_id = cred
2050 .get("client_id")
2051 .and_then(|v| v.as_str())
2052 .ok_or_else(|| AuthError::OidcError("Missing client_id".into()))?
2053 .to_string();
2054 let client_secret = cred
2055 .get("client_secret")
2056 .and_then(|v| v.as_str())
2057 .map(String::from);
2058 let scope = cred
2059 .get("scope")
2060 .and_then(|v| v.as_str())
2061 .unwrap_or("openid profile email")
2062 .to_string();
2063
2064 Ok(McpCredentialInfo::WebInteractive {
2065 issuer_url,
2066 client_id,
2067 client_secret,
2068 scope,
2069 })
2070 }
2071 other => Ok(McpCredentialInfo::Other(other.to_string())),
2072 }
2073}
2074
2075fn create_auth_cookie(
2081 name: &str,
2082 value: &str,
2083 max_age: StdDuration,
2084 secure: bool,
2085) -> Cookie<'static> {
2086 Cookie::build((name.to_string(), value.to_string()))
2087 .path("/")
2088 .http_only(true)
2089 .same_site(SameSite::Lax)
2090 .secure(secure)
2091 .max_age(TimeDuration::seconds(max_age.as_secs() as i64))
2092 .build()
2093}
2094
2095#[derive(Debug)]
2101pub enum AuthError {
2102 MissingCode,
2103 MissingState,
2104 InvalidState,
2105 OidcError(String),
2106 TokenExchangeFailed(String),
2107 AuthNotConfigured,
2108}
2109
2110impl IntoResponse for AuthError {
2111 fn into_response(self) -> Response {
2112 let (_status, msg) = match self {
2113 AuthError::MissingCode => (StatusCode::BAD_REQUEST, "Missing authorization code"),
2114 AuthError::MissingState => (StatusCode::BAD_REQUEST, "Missing state parameter"),
2115 AuthError::InvalidState => (StatusCode::BAD_REQUEST, "Invalid or expired state"),
2116 AuthError::OidcError(_) => (
2117 StatusCode::SERVICE_UNAVAILABLE,
2118 "Authentication service error",
2119 ),
2120 AuthError::TokenExchangeFailed(_) => (StatusCode::BAD_REQUEST, "Token exchange failed"),
2121 AuthError::AuthNotConfigured => {
2122 (StatusCode::NOT_IMPLEMENTED, "Authentication not configured")
2123 }
2124 };
2125 Redirect::temporary(&format!("/?error={}", urlencoding::encode(msg))).into_response()
2126 }
2127}
2128
2129#[cfg(test)]
2130mod principal_tests {
2131 use super::*;
2132 use pep::oidc::types::JwtClaims;
2133
2134 fn claims_with_role(role: serde_json::Value) -> JwtClaims {
2135 let mut c = JwtClaims::default();
2136 c.sub = "sub-uuid".to_string();
2137 c.preferred_username = Some("farzan".to_string());
2138 c.extra.insert("role".to_string(), role);
2139 c
2140 }
2141
2142 #[test]
2145 fn dev_agent_token_yields_agent_namespaced_key() {
2146 assert_eq!(
2147 dev_user_key("dev:agent:farzan"),
2148 Some("agent-farzan".to_string())
2149 );
2150 assert_eq!(
2151 dev_user_key("dev:agent:paydar"),
2152 Some("agent-paydar".to_string())
2153 );
2154 }
2155
2156 #[test]
2157 fn dev_agent_token_rejects_empty_and_colon_names() {
2158 assert_eq!(dev_user_key("dev:agent:"), None);
2159 assert_eq!(dev_user_key("dev:agent: "), None);
2160 assert_eq!(
2161 dev_user_key("dev:agent:a:b"),
2162 None,
2163 "name must not contain ':'"
2164 );
2165 }
2166
2167 #[test]
2168 fn dev_human_token_format_unchanged() {
2169 assert_eq!(
2170 dev_user_key("dev:a@b.c:Some Name:someuser"),
2171 Some("dev:a@b.c".to_string())
2172 );
2173 assert_eq!(dev_user_key("dev:only:two"), None);
2174 }
2175
2176 #[test]
2179 fn role_as_string_classifies_agent() {
2180 let c = claims_with_role(serde_json::json!("agent"));
2181 assert_eq!(claim_role(&c).as_deref(), Some("agent"));
2182 assert_eq!(AuthUser::from(c).kind, PrincipalKind::Agent);
2183 }
2184
2185 #[test]
2186 fn role_as_array_takes_first_value() {
2187 let c = claims_with_role(serde_json::json!(["agent", "other"]));
2189 assert_eq!(claim_role(&c).as_deref(), Some("agent"));
2190 assert_eq!(AuthUser::from(c).kind, PrincipalKind::Agent);
2191 }
2192
2193 #[test]
2194 fn non_agent_roles_classify_human() {
2195 for role in ["user", "admin", "service"] {
2196 let c = claims_with_role(serde_json::json!(role));
2197 assert_eq!(claim_role(&c).as_deref(), Some(role));
2198 assert_eq!(AuthUser::from(c).kind, PrincipalKind::Human, "role={role}");
2199 }
2200 }
2201
2202 #[test]
2203 fn missing_or_nonstring_role_classifies_human() {
2204 let mut c = JwtClaims::default();
2205 c.sub = "sub-uuid".to_string();
2206 assert_eq!(claim_role(&c), None);
2207 assert_eq!(AuthUser::from(c.clone()).kind, PrincipalKind::Human);
2208 c.extra.insert("role".to_string(), serde_json::json!(42));
2209 assert_eq!(claim_role(&c), None, "non-string non-array role ignored");
2210 }
2211
2212 #[test]
2215 fn user_key_prefers_preferred_username() {
2216 let mut c = JwtClaims::default();
2217 c.sub = "sub-uuid".to_string();
2218 c.preferred_username = Some("farzan".to_string());
2219 c.email = Some("rebindable@example.com".to_string());
2220 assert_eq!(jwt_user_key(&c), "farzan", "email must never be the key");
2221 }
2222
2223 #[test]
2224 fn user_key_falls_back_to_sub_on_blank_username() {
2225 let mut c = JwtClaims::default();
2226 c.sub = "sub-uuid".to_string();
2227 c.preferred_username = Some(" ".to_string());
2228 assert_eq!(jwt_user_key(&c), "sub-uuid");
2229 c.preferred_username = None;
2230 assert_eq!(jwt_user_key(&c), "sub-uuid");
2231 }
2232
2233 #[test]
2234 fn user_key_agent_pinned_to_sub_even_with_username() {
2235 let mut c = claims_with_role(serde_json::json!("agent"));
2238 c.sub = "agent-sub-uuid".to_string();
2239 c.preferred_username = Some("farzan".to_string());
2240 assert_eq!(jwt_user_key(&c), "agent-sub-uuid");
2241 }
2242
2243 #[test]
2244 fn authuser_carries_role_and_kind() {
2245 let c = claims_with_role(serde_json::json!("agent"));
2246 let u = AuthUser::from(c);
2247 assert_eq!(u.role.as_deref(), Some("agent"));
2248 assert_eq!(u.kind, PrincipalKind::Agent);
2249 assert_eq!(u.username.as_deref(), Some("farzan"));
2250 }
2251}
2252
2253#[cfg(test)]
2254mod cedar_p2_tests {
2255 use super::*;
2256 use cedar_policy::{Context, Entities, EntityUid, Request};
2257 use pep::cedar::{CedarAuthorizer, CedarConfig};
2258 use std::collections::HashMap;
2259
2260 const POLICY: &str = include_str!("../policies/trustee_default.cedar");
2261 const SCHEMA: &str = include_str!("../policies/trustee_schema.cedarschema");
2262
2263 async fn authorizer() -> CedarAuthorizer {
2264 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2266 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2267 let dir = std::env::temp_dir().join(format!("trustee-cedar-p2-{}-{n}", std::process::id()));
2268 std::fs::create_dir_all(&dir).expect("temp dir");
2269 let policy_path = dir.join("trustee_default.cedar");
2270 let schema_path = dir.join("trustee_schema.cedarschema");
2271 std::fs::write(&policy_path, POLICY).expect("write policy");
2272 std::fs::write(&schema_path, SCHEMA).expect("write schema");
2273 let cfg = CedarConfig {
2274 policy_path,
2275 schema_path: Some(schema_path),
2276 entities_path: None,
2277 default_decision: pep::cedar::DefaultDecision::Deny,
2278 validate_on_load: true,
2279 policy_store_url: None,
2280 policy_store_token: None,
2281 embedded_policy: Some(POLICY),
2282 embedded_schema: Some(SCHEMA),
2283 };
2284 CedarAuthorizer::new_with_policy_store(cfg)
2285 .await
2286 .expect("Cedar init from shipped sources")
2287 }
2288
2289 fn claims_with_role(role: Option<&str>) -> JwtClaims {
2290 let mut c = JwtClaims::default();
2291 c.sub = "test-sub".to_string();
2292 if let Some(r) = role {
2293 c.extra.insert("role".to_string(), serde_json::json!(r));
2294 }
2295 c
2296 }
2297
2298 async fn allowed(auth: &CedarAuthorizer, role: Option<&str>, action: &str) -> bool {
2300 let claims = claims_with_role(role);
2301 let principal_entity = pep::cedar::build_principal_entity(&claims).unwrap();
2302 let app_entity = cedar_policy::Entity::new(
2303 EntityUid::from_str(r#"TrusteeApp::"default""#).unwrap(),
2304 HashMap::new(),
2305 std::collections::HashSet::new(),
2306 )
2307 .unwrap();
2308 let entities = Entities::from_entities(vec![principal_entity, app_entity], None).unwrap();
2309 let request = Request::new(
2310 pep::cedar::build_principal_uid(&claims).unwrap(),
2311 EntityUid::from_str(&format!("Action::\"{action}\"")).unwrap(),
2312 EntityUid::from_str(r#"TrusteeApp::"default""#).unwrap(),
2313 Context::empty(),
2314 None,
2315 )
2316 .unwrap();
2317 auth.is_allowed_with_entities(&request, &entities).allowed()
2318 }
2319
2320 #[tokio::test]
2321 async fn admin_allowed_including_destructive() {
2322 let auth = authorizer().await;
2323 for action in [
2324 actions::VIEW_SESSION,
2325 actions::COMMAND_SESSION,
2326 actions::DELETE_SESSION,
2327 actions::UPDATE_MCP_CREDENTIALS,
2328 ] {
2329 assert!(
2330 allowed(&auth, Some("admin"), action).await,
2331 "admin {action}"
2332 );
2333 }
2334 }
2335
2336 #[tokio::test]
2337 async fn user_full_session_management() {
2338 let auth = authorizer().await;
2339 for action in [
2340 actions::CREATE_SESSION,
2341 actions::COMMAND_SESSION,
2342 actions::DELETE_SESSION,
2343 actions::VIEW_HISTORY,
2344 actions::UPDATE_MCP_CREDENTIALS,
2345 ] {
2346 assert!(allowed(&auth, Some("user"), action).await, "user {action}");
2347 }
2348 }
2349
2350 #[tokio::test]
2351 async fn agent_working_set_including_delete() {
2352 let auth = authorizer().await;
2353 for action in [
2354 actions::CREATE_SESSION,
2355 actions::COMMAND_SESSION,
2356 actions::CANCEL_SESSION,
2357 actions::RESUME_SESSION,
2358 actions::VIEW_HISTORY,
2359 actions::UPDATE_MCP_CREDENTIALS,
2360 actions::DELETE_SESSION,
2365 ] {
2366 assert!(
2367 allowed(&auth, Some("agent"), action).await,
2368 "agent {action}"
2369 );
2370 }
2371 }
2372
2373 #[tokio::test]
2374 async fn service_read_only() {
2375 let auth = authorizer().await;
2376 for action in [
2377 actions::VIEW_SESSION,
2378 actions::LIST_SESSIONS,
2379 actions::VIEW_HISTORY,
2380 ] {
2381 assert!(
2382 allowed(&auth, Some("service"), action).await,
2383 "service {action}"
2384 );
2385 }
2386 for action in [actions::COMMAND_SESSION, actions::DELETE_SESSION] {
2387 assert!(
2388 !allowed(&auth, Some("service"), action).await,
2389 "service {action} denied"
2390 );
2391 }
2392 }
2393
2394 #[tokio::test]
2395 async fn missing_or_unknown_role_denied_everything() {
2396 let auth = authorizer().await;
2397 for role in [None, Some("intern"), Some("Admin")] {
2398 assert!(
2399 !allowed(&auth, role, actions::VIEW_SESSION).await,
2400 "role={role:?} must be denied (fail-closed default)"
2401 );
2402 }
2403 }
2404
2405 #[test]
2406 fn boot_decision_is_fail_closed() {
2407 assert!(crate::cedar_boot_decision(true, false, false).is_err());
2408 assert!(crate::cedar_boot_decision(true, false, true).is_ok());
2409 assert!(crate::cedar_boot_decision(true, true, false).is_ok());
2410 assert!(crate::cedar_boot_decision(false, false, false).is_ok());
2411 }
2412
2413 #[test]
2419 fn shipped_policy_files_have_zero_skew() {
2420 let unknown = crate::audit_policy_action_skew(
2421 crate::EMBEDDED_CEDAR_POLICY,
2422 crate::EMBEDDED_CEDAR_SCHEMA,
2423 );
2424 assert!(unknown.is_empty(), "the shipped pair has skew: {unknown:?}");
2425 for action in crate::policy_action_refs(crate::EMBEDDED_CEDAR_POLICY) {
2428 assert!(
2429 [
2430 actions::LIST_MODELS,
2431 actions::LIST_SESSIONS,
2432 actions::VIEW_SESSION,
2433 actions::VIEW_HISTORY,
2434 actions::CREATE_SESSION,
2435 actions::COMMAND_SESSION,
2436 actions::CANCEL_SESSION,
2437 actions::HANDOFF_SESSION,
2438 actions::RESUME_SESSION,
2439 actions::UPDATE_SESSION,
2440 actions::DELETE_SESSION,
2441 actions::VIEW_MCP_CREDENTIALS,
2442 actions::UPDATE_MCP_CREDENTIALS,
2443 ]
2444 .contains(&action.as_str()),
2445 "the shipped policy references action {action} which no route constant names"
2446 );
2447 }
2448 }
2449
2450 #[test]
2455 fn skew_audit_names_unknown_actions() {
2456 let policy = "permit(principal, action == Action::\"ViewSession\", resource);\npermit(principal, action == Action::\"BogusAction\", resource);";
2457 let schema = "action ViewSession appliesTo { principal: [User], resource: [TrusteeApp] };";
2458 let unknown = crate::audit_policy_action_skew(policy, schema);
2459 assert_eq!(unknown, vec!["BogusAction".to_string()]);
2460 assert!(!unknown.contains(&"ViewSession".to_string()));
2462 }
2463}
2464
2465#[cfg(test)]
2468mod validation_cache_tests {
2469 use super::*;
2470 use std::sync::atomic::{AtomicUsize, Ordering};
2471
2472 fn claims_with(exp_in: i64, sub: &str) -> JwtClaims {
2473 let mut c = JwtClaims::default();
2474 c.exp = unix_now() + exp_in;
2475 c.sub = sub.to_string();
2476 c.preferred_username = Some("farzan".to_string());
2477 c
2478 }
2479
2480 #[test]
2482 fn cache_hit_two_polls_one_validation() {
2483 let cache = ValidationCache::new();
2484 let validations = Arc::new(AtomicUsize::new(0));
2485 let key = validation_cache_key("token-A");
2486
2487 for _ in 0..2 {
2488 if cache.get(&key, unix_now()).is_some() {
2489 continue; }
2491 validations.fetch_add(1, Ordering::SeqCst);
2493 cache.put(key, claims_with(3600, "sub-A"), unix_now());
2494 }
2495
2496 assert_eq!(
2497 validations.load(Ordering::SeqCst),
2498 1,
2499 "second poll must hit the cache, not re-validate"
2500 );
2501 }
2502
2503 #[test]
2505 fn expired_entry_revalidates() {
2506 let cache = ValidationCache::new();
2507 let key = validation_cache_key("token-B");
2508 let now = unix_now();
2509 cache.put(key, claims_with(3600, "sub-B"), now);
2510
2511 assert!(cache.get(&key, now).is_some(), "fresh entry must hit");
2512 assert!(
2513 cache.get(&key, now + CACHE_TTL_SECS + 1).is_none(),
2514 "entry past valid_until must miss and force re-validation"
2515 );
2516 }
2517
2518 #[test]
2521 fn near_expiry_token_never_served_from_cache() {
2522 let cache = ValidationCache::new();
2523 let key = validation_cache_key("token-C");
2524 let now = unix_now();
2525 cache.put(key, claims_with(60, "sub-C"), now);
2527 assert!(
2528 cache.get(&key, now).is_none(),
2529 "entry valid_until <= now must be rejected immediately"
2530 );
2531 }
2532
2533 #[test]
2536 fn cap_evicts_instead_of_growing_unbounded() {
2537 let cache = ValidationCache::new();
2538 let now = unix_now();
2539 for i in 0..CACHE_MAX_ENTRIES {
2540 cache.put(
2541 validation_cache_key(&format!("tok-{i}")),
2542 claims_with(3600, "s"),
2543 now,
2544 );
2545 }
2546 let overflow = validation_cache_key("tok-overflow");
2548 cache.put(overflow, claims_with(3600, "s"), now);
2549
2550 let inner = cache.inner.lock().unwrap();
2551 assert!(
2552 inner.entries.len() <= CACHE_MAX_ENTRIES,
2553 "cache must stay bounded"
2554 );
2555 assert!(
2556 inner.entries.contains_key(&overflow),
2557 "newest entry must survive"
2558 );
2559 assert!(
2560 !inner.entries.contains_key(&validation_cache_key("tok-0")),
2561 "pre-overflow entries were reset, not served stale forever"
2562 );
2563 }
2564
2565 #[test]
2567 fn first_sight_fires_once_per_sub_issuer_pair() {
2568 let cache = ValidationCache::new();
2569 assert!(cache.mark_first_sight("sub-A", "https://idp.tanbal.ir"));
2570 assert!(!cache.mark_first_sight("sub-A", "https://idp.tanbal.ir"));
2571 assert!(
2572 cache.mark_first_sight("sub-A", "https://other.tanbal.ir"),
2573 "different issuer = first sight"
2574 );
2575 assert!(
2576 cache.mark_first_sight("sub-B", "https://idp.tanbal.ir"),
2577 "different sub = first sight"
2578 );
2579 }
2580
2581 #[test]
2583 fn distinct_tokens_distinct_keys() {
2584 assert_ne!(validation_cache_key("tok-1"), validation_cache_key("tok-2"));
2585 assert_eq!(
2587 validation_cache_key("secret"),
2588 validation_cache_key("secret")
2589 );
2590 }
2591}