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 {
867 kind == PrincipalKind::Human && role == Some("admin")
868}
869
870pub async fn check_dispatch_admin(
882 auth: &Option<Arc<AuthState>>,
883 headers: &axum::http::HeaderMap,
884) -> Result<(), StatusCode> {
885 let Some(auth) = auth.as_ref() else {
886 return Ok(()); };
888 let Some(token) = headers
889 .get(header::AUTHORIZATION)
890 .and_then(|v| v.to_str().ok())
891 .and_then(|v| v.strip_prefix("Bearer "))
892 .map(|s| s.to_string())
893 else {
894 return Err(StatusCode::UNAUTHORIZED);
895 };
896 if token.starts_with("dev:") {
897 tracing::warn!("xagent dispatch rejected: dev tokens cannot dispatch agents");
898 return Err(StatusCode::FORBIDDEN);
899 }
900 let claims = auth.validate_token_flexible(&token).await.map_err(|e| {
901 tracing::warn!("xagent dispatch: caller token validation failed: {}", e);
902 StatusCode::UNAUTHORIZED
903 })?;
904 let user = AuthUser::from(claims);
905 if !dispatch_allowed(user.kind, user.role.as_deref()) {
906 tracing::warn!(
907 "xagent dispatch rejected: principal kind={:?} role={:?} — admin required",
908 user.kind,
909 user.role
910 );
911 return Err(StatusCode::FORBIDDEN);
912 }
913 Ok(())
914}
915
916impl AuthState {
917 pub async fn exchange_agent_token(
929 &self,
930 issuer_url: &str,
931 service_token: &str,
932 ) -> Result<(String, u64), StatusCode> {
933 let tr = self
934 .oidc_client
935 .exchange_token(
936 issuer_url,
937 &self.config.client_id,
938 None, service_token,
940 &self.config.client_id,
941 Some("openid groups"),
942 )
943 .await
944 .map_err(|e| {
945 tracing::error!(
946 "xagent dispatch: agent token exchange FAILED at {issuer_url}: {} — impersonation unavailable",
947 e
948 );
949 StatusCode::BAD_GATEWAY
950 })?;
951 Ok((tr.access_token, tr.expires_in.unwrap_or(900)))
952 }
953}
954
955pub async fn check_auth(
976 auth: &Option<Arc<AuthState>>,
977 headers: &axum::http::HeaderMap,
978 action: &str,
979) -> Result<(Option<String>, String), StatusCode> {
980 let Some(auth) = auth.as_ref() else {
981 return Ok((None, "default".to_string())); };
983
984 if let Some(token) = headers
986 .get(header::AUTHORIZATION)
987 .and_then(|v| v.to_str().ok())
988 .and_then(|v| v.strip_prefix("Bearer "))
989 .map(|s| s.to_string())
990 {
991 if token.starts_with("dev:") {
993 if !auth.config.dev_config.local_dev_mode {
994 tracing::warn!("Dev token presented but dev mode is disabled — rejecting");
995 return Err(StatusCode::UNAUTHORIZED);
996 }
997 return match dev_user_key(&token) {
998 Some(key) => Ok((None, key)),
999 None => Err(StatusCode::UNAUTHORIZED),
1000 };
1001 }
1002
1003 return match auth.validate_token_flexible(&token).await {
1004 Ok(claims) => {
1005 if auth.check_cedar_authorized(&claims, action).is_err() {
1006 return Err(StatusCode::FORBIDDEN);
1007 }
1008 Ok((None, auth.resolve_user_key(&claims)))
1009 }
1010 Err(e) => {
1011 tracing::warn!("Bearer token validation failed: {}", e);
1012 Err(StatusCode::UNAUTHORIZED)
1013 }
1014 };
1015 }
1016
1017 let session_id = headers
1019 .get(header::COOKIE)
1020 .and_then(|v| v.to_str().ok())
1021 .and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
1022
1023 let Some(session_id) = session_id else {
1024 tracing::warn!("No auth token found in request");
1025 return Err(StatusCode::UNAUTHORIZED);
1026 };
1027
1028 if session_id.starts_with("dev:") {
1030 if !auth.config.dev_config.local_dev_mode {
1031 tracing::warn!("Dev cookie presented but dev mode is disabled — rejecting");
1032 return Err(StatusCode::UNAUTHORIZED);
1033 }
1034 return match dev_user_key(&session_id) {
1035 Some(key) => Ok((None, key)),
1036 None => Err(StatusCode::UNAUTHORIZED),
1037 };
1038 }
1039
1040 match auth.session_manager.get_token(&session_id).await {
1042 Ok(access_token) => match auth.validate_token(&access_token).await {
1043 Ok(claims) => {
1044 if auth.check_cedar_authorized(&claims, action).is_err() {
1046 return Err(StatusCode::FORBIDDEN);
1047 }
1048 let secure = auth.client_config.redirect_uri.starts_with("https");
1050 let cookie = create_auth_cookie(
1051 &auth.config.cookie_name,
1052 &session_id,
1053 SESSION_COOKIE_MAX_AGE,
1054 secure,
1055 );
1056 Ok((Some(cookie.to_string()), auth.resolve_user_key(&claims)))
1057 }
1058 Err(e) => {
1059 tracing::warn!(
1062 "Session token validation failed: {} — attempting force-refresh",
1063 e
1064 );
1065 match auth.session_manager.force_refresh(&session_id).await {
1066 Ok(new_token) => match auth.validate_token(&new_token).await {
1067 Ok(claims) => {
1068 if auth.check_cedar_authorized(&claims, action).is_err() {
1070 return Err(StatusCode::FORBIDDEN);
1071 }
1072 let secure = auth.client_config.redirect_uri.starts_with("https");
1073 let cookie = create_auth_cookie(
1074 &auth.config.cookie_name,
1075 &session_id,
1076 SESSION_COOKIE_MAX_AGE,
1077 secure,
1078 );
1079 Ok((Some(cookie.to_string()), auth.resolve_user_key(&claims)))
1080 }
1081 Err(e2) => {
1082 tracing::warn!(
1083 "Session token still invalid after force-refresh: {}",
1084 e2
1085 );
1086 Err(StatusCode::UNAUTHORIZED)
1087 }
1088 },
1089 Err(e2) => {
1090 tracing::warn!("Force-refresh failed: {}", e2);
1091 Err(StatusCode::UNAUTHORIZED)
1092 }
1093 }
1094 }
1095 },
1096 Err(e) => {
1097 tracing::warn!("Session lookup/refresh failed: {}", e);
1098 Err(StatusCode::UNAUTHORIZED)
1099 }
1100 }
1101}
1102
1103async fn resolve_access_token(
1109 auth: &AuthState,
1110 headers: &axum::http::HeaderMap,
1111) -> Result<String, StatusCode> {
1112 if let Some(token) = headers
1114 .get(header::AUTHORIZATION)
1115 .and_then(|v| v.to_str().ok())
1116 .and_then(|v| v.strip_prefix("Bearer "))
1117 .map(|s| s.to_string())
1118 {
1119 return Ok(token);
1120 }
1121
1122 let session_id = headers
1124 .get(header::COOKIE)
1125 .and_then(|v| v.to_str().ok())
1126 .and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
1127
1128 match session_id {
1129 Some(sid) if sid.starts_with("dev:") => {
1130 if !auth.config.dev_config.local_dev_mode {
1131 tracing::warn!(
1132 "Dev cookie in resolve_access_token but dev mode is disabled — rejecting"
1133 );
1134 Err(StatusCode::UNAUTHORIZED)
1135 } else {
1136 Ok(sid)
1137 }
1138 }
1139 Some(sid) => auth.session_manager.get_token(&sid).await.map_err(|e| {
1140 tracing::warn!("Failed to resolve session token: {}", e);
1141 StatusCode::UNAUTHORIZED
1142 }),
1143 None => Err(StatusCode::UNAUTHORIZED),
1144 }
1145}
1146
1147fn extract_token_from_cookies(cookie_header: &str, cookie_name: &str) -> Option<String> {
1149 for cookie in cookie_header.split(';') {
1150 let cookie = cookie.trim();
1151 if let Some(value) = cookie.strip_prefix(&format!("{}=", cookie_name)) {
1152 return Some(value.to_string());
1153 }
1154 }
1155 None
1156}
1157
1158pub fn auth_routes() -> axum::Router<crate::ServerState> {
1164 axum::Router::new()
1165 .route("/login", axum::routing::get(login_handler))
1166 .route("/callback", axum::routing::get(callback_handler))
1167 .route("/me", axum::routing::get(me_handler))
1168 .route("/logout", axum::routing::post(logout_handler))
1169 .route("/mcp/login", axum::routing::get(mcp_login_handler))
1170 .route("/mcp/callback", axum::routing::get(mcp_callback_handler))
1171 .route("/mcp/status", axum::routing::get(mcp_status_handler))
1172 .route("/mcp/logout", axum::routing::post(mcp_logout_handler))
1173}
1174
1175#[derive(Debug, Deserialize)]
1177pub struct CallbackQuery {
1178 pub code: Option<String>,
1179 pub state: Option<String>,
1180 pub error: Option<String>,
1181 pub error_description: Option<String>,
1182}
1183
1184async fn login_handler(State(state): State<crate::ServerState>) -> Result<Response, AuthError> {
1186 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1187
1188 if auth.is_dev_mode() {
1190 tracing::info!("Dev mode: creating dev session");
1191 let dev = &auth.config.dev_config;
1192 let dev_token = format!(
1193 "dev:{}:{}:{}",
1194 dev.local_dev_email.as_deref().unwrap_or("dev@localhost"),
1195 dev.local_dev_name.as_deref().unwrap_or("Dev User"),
1196 dev.local_dev_username.as_deref().unwrap_or("dev")
1197 );
1198 let cookie = create_auth_cookie(
1199 &auth.config.cookie_name,
1200 &dev_token,
1201 StdDuration::from_secs(86400),
1202 false,
1203 );
1204 return Ok(Response::builder()
1205 .status(StatusCode::FOUND)
1206 .header(header::LOCATION, "/")
1207 .header(header::SET_COOKIE, cookie.to_string())
1208 .body(Body::empty())
1209 .unwrap());
1210 }
1211
1212 let pkce_session = auth.pkce_manager.create();
1214 let challenge = OidcClient::generate_code_challenge(&pkce_session.verifier);
1215
1216 let auth_url = auth
1217 .oidc_client
1218 .build_authorization_url(&auth.client_config, &pkce_session.state, Some(&challenge))
1219 .await
1220 .map_err(|e| AuthError::OidcError(e.to_string()))?;
1221
1222 let secure = auth.client_config.redirect_uri.starts_with("https");
1226 let pkce_cookie = Cookie::build((
1227 auth.pkce_manager.cookie_name().to_string(),
1228 pkce_session.cookie_value,
1229 ))
1230 .path("/")
1231 .http_only(true)
1232 .same_site(SameSite::Lax)
1233 .secure(secure)
1234 .max_age(TimeDuration::seconds(
1235 auth.pkce_manager.ttl().as_secs() as i64
1236 ))
1237 .build();
1238
1239 Ok(Response::builder()
1240 .status(StatusCode::TEMPORARY_REDIRECT)
1241 .header(header::LOCATION, &auth_url)
1242 .header(header::SET_COOKIE, pkce_cookie.to_string())
1243 .body(Body::empty())
1244 .unwrap())
1245}
1246
1247async fn callback_handler(
1249 State(state): State<crate::ServerState>,
1250 Query(query): Query<CallbackQuery>,
1251 headers: axum::http::HeaderMap,
1252) -> Result<Response, AuthError> {
1253 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1254
1255 if let Some(error) = query.error {
1257 let desc = query.error_description.unwrap_or_default();
1258 tracing::error!("OIDC error: {} - {}", error, desc);
1259 return Ok(Redirect::temporary(&format!(
1260 "/?error={}&error_description={}",
1261 urlencoding::encode(&error),
1262 urlencoding::encode(&desc)
1263 ))
1264 .into_response());
1265 }
1266
1267 let code = query.code.ok_or(AuthError::MissingCode)?;
1268 let oauth_state = query.state.ok_or(AuthError::MissingState)?;
1269
1270 let cookie_header = headers
1272 .get(header::COOKIE)
1273 .and_then(|v| v.to_str().ok())
1274 .unwrap_or("");
1275 let pkce_value = extract_token_from_cookies(cookie_header, auth.pkce_manager.cookie_name())
1276 .ok_or(AuthError::InvalidState)?;
1277
1278 let verifier = auth
1280 .pkce_manager
1281 .verify(&pkce_value, &oauth_state)
1282 .ok_or(AuthError::InvalidState)?;
1283
1284 tracing::info!("Exchanging authorization code for tokens");
1286 let token_response = auth
1287 .oidc_client
1288 .exchange_code_for_tokens(&auth.client_config, &code, Some(&verifier))
1289 .await
1290 .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
1291
1292 let session_id = auth
1293 .session_manager
1294 .create_session(&token_response)
1295 .await
1296 .map_err(|e| AuthError::TokenExchangeFailed(format!("Session creation failed: {}", e)))?;
1297
1298 let max_age = SESSION_COOKIE_MAX_AGE;
1301
1302 let secure = auth.client_config.redirect_uri.starts_with("https");
1304 let cookie = create_auth_cookie(&auth.config.cookie_name, &session_id, max_age, secure);
1305
1306 let clear_pkce = Cookie::build((auth.pkce_manager.cookie_name().to_string(), ""))
1308 .path("/")
1309 .http_only(true)
1310 .same_site(SameSite::Lax)
1311 .max_age(TimeDuration::seconds(-1))
1312 .build();
1313
1314 tracing::info!("Authentication successful, redirecting to /");
1315
1316 Ok(Response::builder()
1317 .status(StatusCode::FOUND)
1318 .header(header::LOCATION, "/")
1319 .header(header::SET_COOKIE, cookie.to_string())
1320 .header(header::SET_COOKIE, clear_pkce.to_string())
1321 .body(Body::empty())
1322 .unwrap())
1323}
1324
1325async fn me_handler(
1327 State(state): State<crate::ServerState>,
1328 headers: axum::http::HeaderMap,
1329) -> Response {
1330 let Some(ref auth) = state.auth else {
1331 return axum::Json(serde_json::json!({
1333 "authenticated": true,
1334 "auth_enabled": false
1335 }))
1336 .into_response();
1337 };
1338
1339 let cookie_header = headers
1340 .get(header::COOKIE)
1341 .and_then(|v| v.to_str().ok())
1342 .unwrap_or("");
1343
1344 let bearer = headers
1346 .get(header::AUTHORIZATION)
1347 .and_then(|v| v.to_str().ok())
1348 .and_then(|v| v.strip_prefix("Bearer "))
1349 .map(String::from);
1350
1351 let token = bearer
1352 .clone()
1353 .or_else(|| extract_token_from_cookies(cookie_header, &auth.config.cookie_name));
1354
1355 let Some(cookie_value) = token else {
1356 return axum::Json(serde_json::json!({
1357 "authenticated": false,
1358 "auth_enabled": true
1359 }))
1360 .into_response();
1361 };
1362
1363 if cookie_value.starts_with("dev:") && auth.config.dev_config.local_dev_mode {
1366 let parts: Vec<&str> = cookie_value.splitn(4, ':').collect();
1367 if parts.len() >= 4 {
1368 return axum::Json(serde_json::json!({
1369 "authenticated": true,
1370 "auth_enabled": true,
1371 "email": parts[1],
1372 "name": parts[2],
1373 "username": parts[3],
1374 "dev_mode": true
1375 }))
1376 .into_response();
1377 }
1378 }
1379
1380 let access_token = if bearer.is_some() {
1382 cookie_value
1384 } else {
1385 match auth.session_manager.get_token(&cookie_value).await {
1387 Ok(token) => token,
1388 Err(e) => {
1389 tracing::debug!("Session token resolution failed for /auth/me: {}", e);
1390 return axum::Json(serde_json::json!({
1391 "authenticated": false,
1392 "auth_enabled": true
1393 }))
1394 .into_response();
1395 }
1396 }
1397 };
1398
1399 match auth.validate_token(&access_token).await {
1401 Ok(claims) => axum::Json(serde_json::json!({
1402 "authenticated": true,
1403 "auth_enabled": true,
1404 "sub": claims.sub,
1405 "email": claims.email,
1406 "name": claims.name,
1407 "username": claims.preferred_username,
1408 "dev_mode": false
1409 }))
1410 .into_response(),
1411 Err(e) => {
1412 tracing::debug!("Token validation failed for /auth/me: {}", e);
1413 axum::Json(serde_json::json!({
1414 "authenticated": false,
1415 "auth_enabled": true
1416 }))
1417 .into_response()
1418 }
1419 }
1420}
1421
1422async fn logout_handler(
1424 State(state): State<crate::ServerState>,
1425 headers: axum::http::HeaderMap,
1426) -> Response {
1427 let cookie_name = state
1428 .auth
1429 .as_ref()
1430 .map(|a| a.config.cookie_name.as_str())
1431 .unwrap_or("trustee_token");
1432
1433 if let Some(ref auth) = state.auth {
1435 if let Some(cookie_header) = headers.get(header::COOKIE).and_then(|v| v.to_str().ok()) {
1436 if let Some(session_id) = extract_token_from_cookies(cookie_header, cookie_name) {
1437 if !session_id.starts_with("dev:") {
1438 let _ = auth.session_manager.destroy_session(&session_id);
1439 }
1440 }
1441 }
1442 }
1443
1444 let cookie = Cookie::build((cookie_name.to_string(), ""))
1445 .path("/")
1446 .http_only(true)
1447 .same_site(SameSite::Lax)
1448 .max_age(TimeDuration::seconds(-1))
1449 .build();
1450
1451 Response::builder()
1452 .status(StatusCode::FOUND)
1453 .header(header::LOCATION, "/")
1454 .header(header::SET_COOKIE, cookie.to_string())
1455 .body(Body::empty())
1456 .unwrap()
1457}
1458
1459#[derive(Debug, Deserialize)]
1465pub struct McpLoginQuery {
1466 pub cred: String,
1467}
1468
1469#[derive(Debug, Deserialize)]
1471pub struct McpCallbackQuery {
1472 pub code: Option<String>,
1473 pub state: Option<String>,
1474 pub error: Option<String>,
1475 pub error_description: Option<String>,
1476}
1477
1478async fn mcp_login_handler(
1483 State(state): State<crate::ServerState>,
1484 Query(query): Query<McpLoginQuery>,
1485 headers: axum::http::HeaderMap,
1486) -> Result<Response, AuthError> {
1487 let (_cookie, _user_key) = crate::auth::check_auth(
1489 &state.auth,
1490 &headers,
1491 crate::auth::actions::UPDATE_MCP_CREDENTIALS,
1492 )
1493 .await
1494 .map_err(|_| AuthError::AuthNotConfigured)?;
1495
1496 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1497
1498 let cred_config = load_mcp_credential(&state, &query.cred).await?;
1500
1501 let (issuer_url, client_id, client_secret, scope) = match &cred_config {
1502 McpCredentialInfo::WebInteractive {
1503 issuer_url,
1504 client_id,
1505 client_secret,
1506 scope,
1507 } => (
1508 issuer_url.clone(),
1509 client_id.clone(),
1510 client_secret.clone(),
1511 scope.clone(),
1512 ),
1513 _ => {
1514 return Ok(Redirect::temporary(&format!(
1515 "/?mcp_error={}",
1516 urlencoding::encode(&format!(
1517 "Credential '{}' is not web-interactive type",
1518 query.cred
1519 ))
1520 ))
1521 .into_response());
1522 }
1523 };
1524
1525 let oidc_client = OidcClient::new();
1527 let verifier = OidcClient::generate_code_verifier();
1528 let challenge = OidcClient::generate_code_challenge(&verifier);
1529 let oauth_state = OidcClient::generate_state();
1530
1531 let mcp_redirect_uri = format!(
1533 "{}/auth/mcp/callback",
1534 auth.client_config
1535 .redirect_uri
1536 .trim_end_matches('/')
1537 .trim_end_matches("/auth/callback")
1538 );
1539
1540 let mcp_client_config = OidcClientConfig {
1541 issuer_url: issuer_url.clone(),
1542 client_id: client_id.clone(),
1543 client_secret: client_secret.clone(),
1544 redirect_uri: mcp_redirect_uri.clone(),
1545 scope: scope.clone(),
1546 code_challenge_method: "S256".to_string(),
1547 };
1548
1549 let auth_url = oidc_client
1551 .build_authorization_url(&mcp_client_config, &oauth_state, Some(&challenge))
1552 .await
1553 .map_err(|e| AuthError::OidcError(e.to_string()))?;
1554
1555 mcp_pkce()
1557 .insert(oauth_state.clone(), verifier.clone(), query.cred.clone())
1558 .await;
1559
1560 tracing::info!(
1561 "Initiating MCP browser login for credential '{}' (issuer={})",
1562 query.cred,
1563 issuer_url
1564 );
1565
1566 Ok(Response::builder()
1567 .status(StatusCode::TEMPORARY_REDIRECT)
1568 .header(header::LOCATION, &auth_url)
1569 .body(Body::empty())
1570 .unwrap())
1571}
1572
1573async fn mcp_callback_handler(
1575 State(state): State<crate::ServerState>,
1576 Query(query): Query<McpCallbackQuery>,
1577 _headers: axum::http::HeaderMap,
1578) -> Result<Response, AuthError> {
1579 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1580
1581 if let Some(error) = query.error {
1583 let desc = query.error_description.unwrap_or_default();
1584 tracing::error!("MCP OIDC error: {} - {}", error, desc);
1585 return Ok(Redirect::temporary(&format!(
1586 "/?mcp_error={}&error_description={}",
1587 urlencoding::encode(&error),
1588 urlencoding::encode(&desc)
1589 ))
1590 .into_response());
1591 }
1592
1593 let code = query.code.ok_or(AuthError::MissingCode)?;
1594 let oauth_state = query.state.ok_or(AuthError::MissingState)?;
1595
1596 let pkce_data = mcp_pkce()
1598 .take(&oauth_state)
1599 .await
1600 .ok_or(AuthError::InvalidState)?;
1601
1602 let verifier = pkce_data.verifier;
1603 let cred_name = &pkce_data.cred_name;
1604
1605 let cred_config = load_mcp_credential(&state, cred_name).await?;
1607
1608 let (issuer_url, client_id, client_secret, scope) = match &cred_config {
1609 McpCredentialInfo::WebInteractive {
1610 issuer_url,
1611 client_id,
1612 client_secret,
1613 scope,
1614 } => (
1615 issuer_url.clone(),
1616 client_id.clone(),
1617 client_secret.clone(),
1618 scope.clone(),
1619 ),
1620 _ => {
1621 return Ok(Redirect::temporary(&format!(
1622 "/?mcp_error={}",
1623 urlencoding::encode("Credential is not web-interactive type")
1624 ))
1625 .into_response());
1626 }
1627 };
1628
1629 let mcp_redirect_uri = format!(
1631 "{}/auth/mcp/callback",
1632 auth.client_config
1633 .redirect_uri
1634 .trim_end_matches('/')
1635 .trim_end_matches("/auth/callback")
1636 );
1637
1638 let mcp_client_config = OidcClientConfig {
1639 issuer_url: issuer_url.clone(),
1640 client_id: client_id.clone(),
1641 client_secret: client_secret.clone(),
1642 redirect_uri: mcp_redirect_uri,
1643 scope: scope.clone(),
1644 code_challenge_method: "S256".to_string(),
1645 };
1646
1647 tracing::info!(
1649 "Exchanging MCP authorization code for tokens (credential={})",
1650 cred_name
1651 );
1652 let oidc_client = OidcClient::new();
1653 let token_response = oidc_client
1654 .exchange_code_for_tokens(&mcp_client_config, &code, Some(&verifier))
1655 .await
1656 .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
1657
1658 let expires_at = {
1660 let now = std::time::SystemTime::now()
1661 .duration_since(std::time::UNIX_EPOCH)
1662 .unwrap_or_default()
1663 .as_secs();
1664 let expires_epoch = now + token_response.expires_in.unwrap_or(900);
1665 let days = expires_epoch / 86400;
1666 let rem = expires_epoch % 86400;
1667 let h = rem / 3600;
1668 let m = (rem % 3600) / 60;
1669 let s = rem % 60;
1670 let z = days as i64 + 719468;
1671 let era = if z >= 0 { z } else { z - 146096 } / 146097;
1672 let doe = (z - era * 146097) as u64;
1673 let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
1674 let y = yoe as i64 + era * 400;
1675 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
1676 let mp = (5 * doy + 2) / 153;
1677 let d = doy - (153 * mp + 2) / 5 + 1;
1678 let mon = if mp < 10 { mp + 3 } else { mp - 9 };
1679 let yr = if mon <= 2 { y + 1 } else { y };
1680 format!("{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z", yr, mon, d, h, m, s)
1681 };
1682
1683 use pep::{FileTokenStore, StoredToken, TokenStore};
1685
1686 let stored = StoredToken::new(
1687 &token_response.access_token,
1688 token_response.refresh_token.clone(),
1689 "Bearer",
1690 &expires_at,
1691 token_response.scope.clone(),
1692 );
1693
1694 let agent_name = state
1695 .config_toml
1696 .as_ref()
1697 .and_then(|t| {
1698 toml::from_str::<toml::Value>(t).ok().and_then(|v| {
1699 v.get("agent")
1700 .and_then(|a| a.get("name"))
1701 .and_then(|n| n.as_str())
1702 .map(String::from)
1703 })
1704 })
1705 .unwrap_or_else(|| "trustee".to_string());
1706 let token_store = FileTokenStore::new(&agent_name);
1707
1708 if let Err(e) = token_store.save(cred_name, &stored) {
1709 tracing::error!("Failed to store MCP token: {}", e);
1710 return Ok(Redirect::temporary(&format!(
1711 "/?mcp_error={}",
1712 urlencoding::encode(&format!("Failed to store token: {}", e))
1713 ))
1714 .into_response());
1715 }
1716
1717 tracing::info!(
1718 "MCP authentication successful for credential '{}' (expires {})",
1719 cred_name,
1720 expires_at
1721 );
1722
1723 Ok(Response::builder()
1724 .status(StatusCode::FOUND)
1725 .header(
1726 header::LOCATION,
1727 format!("/?mcp_connected={}", urlencoding::encode(cred_name)),
1728 )
1729 .body(Body::empty())
1730 .unwrap())
1731}
1732
1733pub(crate) async fn mcp_status_handler(
1738 State(state): State<crate::ServerState>,
1739 headers: axum::http::HeaderMap,
1740) -> Response {
1741 use pep::{FileTokenStore, TokenStore};
1742
1743 let (_cookie, _user_key) = match crate::auth::check_auth(
1745 &state.auth,
1746 &headers,
1747 crate::auth::actions::VIEW_MCP_CREDENTIALS,
1748 )
1749 .await
1750 {
1751 Ok(result) => result,
1752 Err(code) => {
1753 return (code, Json(serde_json::json!({"error": "Unauthorized"}))).into_response()
1754 }
1755 };
1756
1757 let config_toml = match &state.config_toml {
1763 Some(t) => t.clone(),
1764 None => return (StatusCode::INTERNAL_SERVER_ERROR, "Config not loaded").into_response(),
1765 };
1766
1767 let mcp_config: toml::Value = match toml::from_str(&config_toml) {
1768 Ok(v) => v,
1769 Err(_) => return Json(serde_json::json!([])).into_response(),
1770 };
1771
1772 let agent_name = {
1775 let table: toml::Table = toml::from_str(&config_toml).unwrap_or_default();
1776 table
1777 .get("agent")
1778 .and_then(|a| a.get("name"))
1779 .and_then(|n| n.as_str())
1780 .unwrap_or("trustee")
1781 .to_string()
1782 };
1783 let token_store = FileTokenStore::new(&agent_name);
1784
1785 let servers = mcp_config
1787 .get("mcp")
1788 .and_then(|m| m.get("servers"))
1789 .and_then(|s| s.as_array());
1790 let credentials = mcp_config
1791 .get("mcp")
1792 .and_then(|m| m.get("credentials"))
1793 .and_then(|c| c.as_table());
1794
1795 let mut cred_servers: std::collections::HashMap<String, Vec<String>> =
1796 std::collections::HashMap::new();
1797 if let Some(servers) = servers {
1798 for server in servers {
1799 let name = server.get("name").and_then(|n| n.as_str()).unwrap_or("");
1800 let cred_ref = server
1801 .get("credentials")
1802 .and_then(|c| c.as_str())
1803 .unwrap_or("");
1804 if !cred_ref.is_empty() {
1805 cred_servers
1806 .entry(cred_ref.to_string())
1807 .or_default()
1808 .push(name.to_string());
1809 }
1810 }
1811 }
1812
1813 let mut result = Vec::new();
1814
1815 if let Some(creds) = credentials {
1816 for (cred_name, cred_config) in creds {
1817 let cred_type = cred_config
1818 .get("type")
1819 .and_then(|t| t.as_str())
1820 .unwrap_or("unknown");
1821 let servers_using = cred_servers.get(cred_name).cloned().unwrap_or_default();
1822
1823 if cred_type == "web-session" {
1824 let connected = state.auth.is_some();
1826 result.push(serde_json::json!({
1827 "credential": cred_name,
1828 "type": cred_type,
1829 "connected": connected,
1830 "servers": servers_using,
1831 }));
1832 } else if cred_type == "service-account" {
1833 let token = cred_config
1837 .get("service_token")
1838 .and_then(|t| t.as_str())
1839 .unwrap_or("");
1840 result.push(serde_json::json!({
1841 "credential": cred_name,
1842 "type": cred_type,
1843 "connected": !token.is_empty(),
1844 "servers": servers_using,
1845 }));
1846 } else if cred_type == "static" {
1847 let token = cred_config
1849 .get("token")
1850 .and_then(|t| t.as_str())
1851 .unwrap_or("");
1852 result.push(serde_json::json!({
1853 "credential": cred_name,
1854 "type": cred_type,
1855 "connected": !token.is_empty(),
1856 "servers": servers_using,
1857 }));
1858 } else if cred_type == "web-interactive" || cred_type == "interactive" {
1859 let status = match token_store.load(cred_name) {
1861 Ok(Some(token)) => {
1862 let expired = token.is_expired();
1863 serde_json::json!({
1864 "credential": cred_name,
1865 "type": cred_type,
1866 "connected": !expired,
1867 "expires_at": token.expires_at,
1868 "servers": servers_using,
1869 })
1870 }
1871 _ => serde_json::json!({
1872 "credential": cred_name,
1873 "type": cred_type,
1874 "connected": false,
1875 "servers": servers_using,
1876 }),
1877 };
1878 result.push(status);
1879 }
1880 }
1881 }
1882
1883 Json(serde_json::Value::Array(result)).into_response()
1884}
1885
1886async fn mcp_logout_handler(
1888 State(state): State<crate::ServerState>,
1889 Query(query): Query<McpLoginQuery>,
1890 headers: axum::http::HeaderMap,
1891) -> Response {
1892 use pep::{FileTokenStore, TokenStore};
1893
1894 let (_cookie, _user_key) = match crate::auth::check_auth(
1896 &state.auth,
1897 &headers,
1898 crate::auth::actions::UPDATE_MCP_CREDENTIALS,
1899 )
1900 .await
1901 {
1902 Ok(result) => result,
1903 Err(code) => {
1904 return (code, Json(serde_json::json!({"error": "Unauthorized"}))).into_response()
1905 }
1906 };
1907
1908 let agent_name = {
1911 let config_toml = state.config_toml.clone().unwrap_or_default();
1912 let table: toml::Table = toml::from_str(&config_toml).unwrap_or_default();
1913 table
1914 .get("agent")
1915 .and_then(|a| a.get("name"))
1916 .and_then(|n| n.as_str())
1917 .unwrap_or("trustee")
1918 .to_string()
1919 };
1920 let token_store = FileTokenStore::new(&agent_name);
1921
1922 match token_store.delete(&query.cred) {
1923 Ok(()) => {
1924 tracing::info!("Removed MCP credentials for '{}'", query.cred);
1925 Json(serde_json::json!({"success": true})).into_response()
1926 }
1927 Err(e) => {
1928 tracing::error!("Failed to remove MCP credentials: {}", e);
1929 (
1930 StatusCode::INTERNAL_SERVER_ERROR,
1931 Json(serde_json::json!({"error": e.to_string()})),
1932 )
1933 .into_response()
1934 }
1935 }
1936}
1937
1938struct McpPkceStore {
1945 entries: tokio::sync::Mutex<std::collections::HashMap<String, McpPkceEntry>>,
1946}
1947
1948struct McpPkceEntry {
1949 verifier: String,
1950 cred_name: String,
1951 created_at: std::time::Instant,
1952}
1953
1954impl McpPkceStore {
1955 fn new() -> Self {
1956 Self {
1957 entries: tokio::sync::Mutex::new(std::collections::HashMap::new()),
1958 }
1959 }
1960
1961 async fn insert(&self, state: String, verifier: String, cred_name: String) {
1963 let mut map = self.entries.lock().await;
1964 let cutoff = std::time::Instant::now() - std::time::Duration::from_secs(600);
1966 map.retain(|_, v| v.created_at > cutoff);
1967 map.insert(
1968 state,
1969 McpPkceEntry {
1970 verifier,
1971 cred_name,
1972 created_at: std::time::Instant::now(),
1973 },
1974 );
1975 }
1976
1977 async fn take(&self, state: &str) -> Option<McpPkceEntry> {
1979 let mut map = self.entries.lock().await;
1980 map.remove(state)
1981 }
1982}
1983
1984static MCP_PKCE: std::sync::OnceLock<McpPkceStore> = std::sync::OnceLock::new();
1986
1987fn mcp_pkce() -> &'static McpPkceStore {
1989 MCP_PKCE.get_or_init(McpPkceStore::new)
1990}
1991
1992enum McpCredentialInfo {
1994 WebInteractive {
1995 issuer_url: String,
1996 client_id: String,
1997 client_secret: Option<String>,
1998 scope: String,
1999 },
2000 Other(String),
2001}
2002
2003async fn load_mcp_credential(
2005 state: &crate::ServerState,
2006 cred_name: &str,
2007) -> Result<McpCredentialInfo, AuthError> {
2008 let config_toml = state
2009 .config_toml
2010 .clone()
2011 .ok_or(AuthError::AuthNotConfigured)?;
2012
2013 let config: toml::Value = toml::from_str(&config_toml)
2014 .map_err(|e| AuthError::OidcError(format!("Config parse error: {}", e)))?;
2015
2016 let cred = config
2017 .get("mcp")
2018 .and_then(|m| m.get("credentials"))
2019 .and_then(|c| c.as_table())
2020 .and_then(|c| c.get(cred_name))
2021 .ok_or_else(|| AuthError::OidcError(format!("Credential '{}' not found", cred_name)))?;
2022
2023 let cred_type = cred
2024 .get("type")
2025 .and_then(|t| t.as_str())
2026 .unwrap_or("unknown");
2027
2028 match cred_type {
2029 "web-interactive" => {
2030 let issuer_url = cred
2031 .get("issuer_url")
2032 .and_then(|v| v.as_str())
2033 .ok_or_else(|| AuthError::OidcError("Missing issuer_url".into()))?
2034 .to_string();
2035 let client_id = cred
2036 .get("client_id")
2037 .and_then(|v| v.as_str())
2038 .ok_or_else(|| AuthError::OidcError("Missing client_id".into()))?
2039 .to_string();
2040 let client_secret = cred
2041 .get("client_secret")
2042 .and_then(|v| v.as_str())
2043 .map(String::from);
2044 let scope = cred
2045 .get("scope")
2046 .and_then(|v| v.as_str())
2047 .unwrap_or("openid profile email")
2048 .to_string();
2049
2050 Ok(McpCredentialInfo::WebInteractive {
2051 issuer_url,
2052 client_id,
2053 client_secret,
2054 scope,
2055 })
2056 }
2057 other => Ok(McpCredentialInfo::Other(other.to_string())),
2058 }
2059}
2060
2061fn create_auth_cookie(
2067 name: &str,
2068 value: &str,
2069 max_age: StdDuration,
2070 secure: bool,
2071) -> Cookie<'static> {
2072 Cookie::build((name.to_string(), value.to_string()))
2073 .path("/")
2074 .http_only(true)
2075 .same_site(SameSite::Lax)
2076 .secure(secure)
2077 .max_age(TimeDuration::seconds(max_age.as_secs() as i64))
2078 .build()
2079}
2080
2081#[derive(Debug)]
2087pub enum AuthError {
2088 MissingCode,
2089 MissingState,
2090 InvalidState,
2091 OidcError(String),
2092 TokenExchangeFailed(String),
2093 AuthNotConfigured,
2094}
2095
2096impl IntoResponse for AuthError {
2097 fn into_response(self) -> Response {
2098 let (_status, msg) = match self {
2099 AuthError::MissingCode => (StatusCode::BAD_REQUEST, "Missing authorization code"),
2100 AuthError::MissingState => (StatusCode::BAD_REQUEST, "Missing state parameter"),
2101 AuthError::InvalidState => (StatusCode::BAD_REQUEST, "Invalid or expired state"),
2102 AuthError::OidcError(_) => (
2103 StatusCode::SERVICE_UNAVAILABLE,
2104 "Authentication service error",
2105 ),
2106 AuthError::TokenExchangeFailed(_) => (StatusCode::BAD_REQUEST, "Token exchange failed"),
2107 AuthError::AuthNotConfigured => {
2108 (StatusCode::NOT_IMPLEMENTED, "Authentication not configured")
2109 }
2110 };
2111 Redirect::temporary(&format!("/?error={}", urlencoding::encode(msg))).into_response()
2112 }
2113}
2114
2115#[cfg(test)]
2116mod principal_tests {
2117 use super::*;
2118 use pep::oidc::types::JwtClaims;
2119
2120 fn claims_with_role(role: serde_json::Value) -> JwtClaims {
2121 let mut c = JwtClaims::default();
2122 c.sub = "sub-uuid".to_string();
2123 c.preferred_username = Some("farzan".to_string());
2124 c.extra.insert("role".to_string(), role);
2125 c
2126 }
2127
2128 #[test]
2131 fn dev_agent_token_yields_agent_namespaced_key() {
2132 assert_eq!(
2133 dev_user_key("dev:agent:farzan"),
2134 Some("agent-farzan".to_string())
2135 );
2136 assert_eq!(
2137 dev_user_key("dev:agent:paydar"),
2138 Some("agent-paydar".to_string())
2139 );
2140 }
2141
2142 #[test]
2143 fn dev_agent_token_rejects_empty_and_colon_names() {
2144 assert_eq!(dev_user_key("dev:agent:"), None);
2145 assert_eq!(dev_user_key("dev:agent: "), None);
2146 assert_eq!(
2147 dev_user_key("dev:agent:a:b"),
2148 None,
2149 "name must not contain ':'"
2150 );
2151 }
2152
2153 #[test]
2154 fn dev_human_token_format_unchanged() {
2155 assert_eq!(
2156 dev_user_key("dev:a@b.c:Some Name:someuser"),
2157 Some("dev:a@b.c".to_string())
2158 );
2159 assert_eq!(dev_user_key("dev:only:two"), None);
2160 }
2161
2162 #[test]
2165 fn role_as_string_classifies_agent() {
2166 let c = claims_with_role(serde_json::json!("agent"));
2167 assert_eq!(claim_role(&c).as_deref(), Some("agent"));
2168 assert_eq!(AuthUser::from(c).kind, PrincipalKind::Agent);
2169 }
2170
2171 #[test]
2172 fn role_as_array_takes_first_value() {
2173 let c = claims_with_role(serde_json::json!(["agent", "other"]));
2175 assert_eq!(claim_role(&c).as_deref(), Some("agent"));
2176 assert_eq!(AuthUser::from(c).kind, PrincipalKind::Agent);
2177 }
2178
2179 #[test]
2180 fn non_agent_roles_classify_human() {
2181 for role in ["user", "admin", "service"] {
2182 let c = claims_with_role(serde_json::json!(role));
2183 assert_eq!(claim_role(&c).as_deref(), Some(role));
2184 assert_eq!(AuthUser::from(c).kind, PrincipalKind::Human, "role={role}");
2185 }
2186 }
2187
2188 #[test]
2189 fn missing_or_nonstring_role_classifies_human() {
2190 let mut c = JwtClaims::default();
2191 c.sub = "sub-uuid".to_string();
2192 assert_eq!(claim_role(&c), None);
2193 assert_eq!(AuthUser::from(c.clone()).kind, PrincipalKind::Human);
2194 c.extra.insert("role".to_string(), serde_json::json!(42));
2195 assert_eq!(claim_role(&c), None, "non-string non-array role ignored");
2196 }
2197
2198 #[test]
2201 fn user_key_prefers_preferred_username() {
2202 let mut c = JwtClaims::default();
2203 c.sub = "sub-uuid".to_string();
2204 c.preferred_username = Some("farzan".to_string());
2205 c.email = Some("rebindable@example.com".to_string());
2206 assert_eq!(jwt_user_key(&c), "farzan", "email must never be the key");
2207 }
2208
2209 #[test]
2210 fn user_key_falls_back_to_sub_on_blank_username() {
2211 let mut c = JwtClaims::default();
2212 c.sub = "sub-uuid".to_string();
2213 c.preferred_username = Some(" ".to_string());
2214 assert_eq!(jwt_user_key(&c), "sub-uuid");
2215 c.preferred_username = None;
2216 assert_eq!(jwt_user_key(&c), "sub-uuid");
2217 }
2218
2219 #[test]
2220 fn user_key_agent_pinned_to_sub_even_with_username() {
2221 let mut c = claims_with_role(serde_json::json!("agent"));
2224 c.sub = "agent-sub-uuid".to_string();
2225 c.preferred_username = Some("farzan".to_string());
2226 assert_eq!(jwt_user_key(&c), "agent-sub-uuid");
2227 }
2228
2229 #[test]
2230 fn authuser_carries_role_and_kind() {
2231 let c = claims_with_role(serde_json::json!("agent"));
2232 let u = AuthUser::from(c);
2233 assert_eq!(u.role.as_deref(), Some("agent"));
2234 assert_eq!(u.kind, PrincipalKind::Agent);
2235 assert_eq!(u.username.as_deref(), Some("farzan"));
2236 }
2237}
2238
2239#[cfg(test)]
2240mod cedar_p2_tests {
2241 use super::*;
2242 use cedar_policy::{Context, Entities, EntityUid, Request};
2243 use pep::cedar::{CedarAuthorizer, CedarConfig};
2244 use std::collections::HashMap;
2245
2246 const POLICY: &str = include_str!("../policies/trustee_default.cedar");
2247 const SCHEMA: &str = include_str!("../policies/trustee_schema.cedarschema");
2248
2249 async fn authorizer() -> CedarAuthorizer {
2250 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2252 let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2253 let dir = std::env::temp_dir().join(format!("trustee-cedar-p2-{}-{n}", std::process::id()));
2254 std::fs::create_dir_all(&dir).expect("temp dir");
2255 let policy_path = dir.join("trustee_default.cedar");
2256 let schema_path = dir.join("trustee_schema.cedarschema");
2257 std::fs::write(&policy_path, POLICY).expect("write policy");
2258 std::fs::write(&schema_path, SCHEMA).expect("write schema");
2259 let cfg = CedarConfig {
2260 policy_path,
2261 schema_path: Some(schema_path),
2262 entities_path: None,
2263 default_decision: pep::cedar::DefaultDecision::Deny,
2264 validate_on_load: true,
2265 policy_store_url: None,
2266 policy_store_token: None,
2267 embedded_policy: Some(POLICY),
2268 embedded_schema: Some(SCHEMA),
2269 };
2270 CedarAuthorizer::new_with_policy_store(cfg)
2271 .await
2272 .expect("Cedar init from shipped sources")
2273 }
2274
2275 fn claims_with_role(role: Option<&str>) -> JwtClaims {
2276 let mut c = JwtClaims::default();
2277 c.sub = "test-sub".to_string();
2278 if let Some(r) = role {
2279 c.extra.insert("role".to_string(), serde_json::json!(r));
2280 }
2281 c
2282 }
2283
2284 async fn allowed(auth: &CedarAuthorizer, role: Option<&str>, action: &str) -> bool {
2286 let claims = claims_with_role(role);
2287 let principal_entity = pep::cedar::build_principal_entity(&claims).unwrap();
2288 let app_entity = cedar_policy::Entity::new(
2289 EntityUid::from_str(r#"TrusteeApp::"default""#).unwrap(),
2290 HashMap::new(),
2291 std::collections::HashSet::new(),
2292 )
2293 .unwrap();
2294 let entities = Entities::from_entities(vec![principal_entity, app_entity], None).unwrap();
2295 let request = Request::new(
2296 pep::cedar::build_principal_uid(&claims).unwrap(),
2297 EntityUid::from_str(&format!("Action::\"{action}\"")).unwrap(),
2298 EntityUid::from_str(r#"TrusteeApp::"default""#).unwrap(),
2299 Context::empty(),
2300 None,
2301 )
2302 .unwrap();
2303 auth.is_allowed_with_entities(&request, &entities).allowed()
2304 }
2305
2306 #[tokio::test]
2307 async fn admin_allowed_including_destructive() {
2308 let auth = authorizer().await;
2309 for action in [
2310 actions::VIEW_SESSION,
2311 actions::COMMAND_SESSION,
2312 actions::DELETE_SESSION,
2313 actions::UPDATE_MCP_CREDENTIALS,
2314 ] {
2315 assert!(
2316 allowed(&auth, Some("admin"), action).await,
2317 "admin {action}"
2318 );
2319 }
2320 }
2321
2322 #[tokio::test]
2323 async fn user_full_session_management() {
2324 let auth = authorizer().await;
2325 for action in [
2326 actions::CREATE_SESSION,
2327 actions::COMMAND_SESSION,
2328 actions::DELETE_SESSION,
2329 actions::VIEW_HISTORY,
2330 actions::UPDATE_MCP_CREDENTIALS,
2331 ] {
2332 assert!(allowed(&auth, Some("user"), action).await, "user {action}");
2333 }
2334 }
2335
2336 #[tokio::test]
2337 async fn agent_working_set_including_delete() {
2338 let auth = authorizer().await;
2339 for action in [
2340 actions::CREATE_SESSION,
2341 actions::COMMAND_SESSION,
2342 actions::CANCEL_SESSION,
2343 actions::RESUME_SESSION,
2344 actions::VIEW_HISTORY,
2345 actions::UPDATE_MCP_CREDENTIALS,
2346 actions::DELETE_SESSION,
2351 ] {
2352 assert!(
2353 allowed(&auth, Some("agent"), action).await,
2354 "agent {action}"
2355 );
2356 }
2357 }
2358
2359 #[tokio::test]
2360 async fn service_read_only() {
2361 let auth = authorizer().await;
2362 for action in [
2363 actions::VIEW_SESSION,
2364 actions::LIST_SESSIONS,
2365 actions::VIEW_HISTORY,
2366 ] {
2367 assert!(
2368 allowed(&auth, Some("service"), action).await,
2369 "service {action}"
2370 );
2371 }
2372 for action in [actions::COMMAND_SESSION, actions::DELETE_SESSION] {
2373 assert!(
2374 !allowed(&auth, Some("service"), action).await,
2375 "service {action} denied"
2376 );
2377 }
2378 }
2379
2380 #[tokio::test]
2381 async fn missing_or_unknown_role_denied_everything() {
2382 let auth = authorizer().await;
2383 for role in [None, Some("intern"), Some("Admin")] {
2384 assert!(
2385 !allowed(&auth, role, actions::VIEW_SESSION).await,
2386 "role={role:?} must be denied (fail-closed default)"
2387 );
2388 }
2389 }
2390
2391 #[test]
2392 fn boot_decision_is_fail_closed() {
2393 assert!(crate::cedar_boot_decision(true, false, false).is_err());
2394 assert!(crate::cedar_boot_decision(true, false, true).is_ok());
2395 assert!(crate::cedar_boot_decision(true, true, false).is_ok());
2396 assert!(crate::cedar_boot_decision(false, false, false).is_ok());
2397 }
2398
2399 #[test]
2405 fn shipped_policy_files_have_zero_skew() {
2406 let unknown = crate::audit_policy_action_skew(
2407 crate::EMBEDDED_CEDAR_POLICY,
2408 crate::EMBEDDED_CEDAR_SCHEMA,
2409 );
2410 assert!(unknown.is_empty(), "the shipped pair has skew: {unknown:?}");
2411 for action in crate::policy_action_refs(crate::EMBEDDED_CEDAR_POLICY) {
2414 assert!(
2415 [
2416 actions::LIST_MODELS,
2417 actions::LIST_SESSIONS,
2418 actions::VIEW_SESSION,
2419 actions::VIEW_HISTORY,
2420 actions::CREATE_SESSION,
2421 actions::COMMAND_SESSION,
2422 actions::CANCEL_SESSION,
2423 actions::HANDOFF_SESSION,
2424 actions::RESUME_SESSION,
2425 actions::UPDATE_SESSION,
2426 actions::DELETE_SESSION,
2427 actions::VIEW_MCP_CREDENTIALS,
2428 actions::UPDATE_MCP_CREDENTIALS,
2429 ]
2430 .contains(&action.as_str()),
2431 "the shipped policy references action {action} which no route constant names"
2432 );
2433 }
2434 }
2435
2436 #[test]
2441 fn skew_audit_names_unknown_actions() {
2442 let policy = "permit(principal, action == Action::\"ViewSession\", resource);\npermit(principal, action == Action::\"BogusAction\", resource);";
2443 let schema = "action ViewSession appliesTo { principal: [User], resource: [TrusteeApp] };";
2444 let unknown = crate::audit_policy_action_skew(policy, schema);
2445 assert_eq!(unknown, vec!["BogusAction".to_string()]);
2446 assert!(!unknown.contains(&"ViewSession".to_string()));
2448 }
2449}
2450
2451#[cfg(test)]
2454mod validation_cache_tests {
2455 use super::*;
2456 use std::sync::atomic::{AtomicUsize, Ordering};
2457
2458 fn claims_with(exp_in: i64, sub: &str) -> JwtClaims {
2459 let mut c = JwtClaims::default();
2460 c.exp = unix_now() + exp_in;
2461 c.sub = sub.to_string();
2462 c.preferred_username = Some("farzan".to_string());
2463 c
2464 }
2465
2466 #[test]
2468 fn cache_hit_two_polls_one_validation() {
2469 let cache = ValidationCache::new();
2470 let validations = Arc::new(AtomicUsize::new(0));
2471 let key = validation_cache_key("token-A");
2472
2473 for _ in 0..2 {
2474 if cache.get(&key, unix_now()).is_some() {
2475 continue; }
2477 validations.fetch_add(1, Ordering::SeqCst);
2479 cache.put(key, claims_with(3600, "sub-A"), unix_now());
2480 }
2481
2482 assert_eq!(
2483 validations.load(Ordering::SeqCst),
2484 1,
2485 "second poll must hit the cache, not re-validate"
2486 );
2487 }
2488
2489 #[test]
2491 fn expired_entry_revalidates() {
2492 let cache = ValidationCache::new();
2493 let key = validation_cache_key("token-B");
2494 let now = unix_now();
2495 cache.put(key, claims_with(3600, "sub-B"), now);
2496
2497 assert!(cache.get(&key, now).is_some(), "fresh entry must hit");
2498 assert!(
2499 cache.get(&key, now + CACHE_TTL_SECS + 1).is_none(),
2500 "entry past valid_until must miss and force re-validation"
2501 );
2502 }
2503
2504 #[test]
2507 fn near_expiry_token_never_served_from_cache() {
2508 let cache = ValidationCache::new();
2509 let key = validation_cache_key("token-C");
2510 let now = unix_now();
2511 cache.put(key, claims_with(60, "sub-C"), now);
2513 assert!(
2514 cache.get(&key, now).is_none(),
2515 "entry valid_until <= now must be rejected immediately"
2516 );
2517 }
2518
2519 #[test]
2522 fn cap_evicts_instead_of_growing_unbounded() {
2523 let cache = ValidationCache::new();
2524 let now = unix_now();
2525 for i in 0..CACHE_MAX_ENTRIES {
2526 cache.put(
2527 validation_cache_key(&format!("tok-{i}")),
2528 claims_with(3600, "s"),
2529 now,
2530 );
2531 }
2532 let overflow = validation_cache_key("tok-overflow");
2534 cache.put(overflow, claims_with(3600, "s"), now);
2535
2536 let inner = cache.inner.lock().unwrap();
2537 assert!(
2538 inner.entries.len() <= CACHE_MAX_ENTRIES,
2539 "cache must stay bounded"
2540 );
2541 assert!(
2542 inner.entries.contains_key(&overflow),
2543 "newest entry must survive"
2544 );
2545 assert!(
2546 !inner.entries.contains_key(&validation_cache_key("tok-0")),
2547 "pre-overflow entries were reset, not served stale forever"
2548 );
2549 }
2550
2551 #[test]
2553 fn first_sight_fires_once_per_sub_issuer_pair() {
2554 let cache = ValidationCache::new();
2555 assert!(cache.mark_first_sight("sub-A", "https://idp.tanbal.ir"));
2556 assert!(!cache.mark_first_sight("sub-A", "https://idp.tanbal.ir"));
2557 assert!(
2558 cache.mark_first_sight("sub-A", "https://other.tanbal.ir"),
2559 "different issuer = first sight"
2560 );
2561 assert!(
2562 cache.mark_first_sight("sub-B", "https://idp.tanbal.ir"),
2563 "different sub = first sight"
2564 );
2565 }
2566
2567 #[test]
2569 fn distinct_tokens_distinct_keys() {
2570 assert_ne!(validation_cache_key("tok-1"), validation_cache_key("tok-2"));
2571 assert_eq!(
2573 validation_cache_key("secret"),
2574 validation_cache_key("secret")
2575 );
2576 }
2577}