1use std::sync::Arc;
16use std::time::Duration as StdDuration;
17
18use axum::{
19 body::Body,
20 extract::{Query, State},
21 http::{header, StatusCode},
22 response::{IntoResponse, Json, Redirect, Response},
23};
24use axum_extra::extract::cookie::{Cookie, SameSite};
25use pep::oidc_client::OidcClient;
26use pep::oidc_resource_server::ResourceServerClient;
27use pep::oidc::pkce_cookie::PkceCookieManager;
28use pep::session_manager::WebSessionManager;
29use pep::{DevConfig, JwtClaims, JwtValidationOptions, OidcClientConfig};
30use serde::Deserialize;
31use time::Duration as TimeDuration;
32
33use cedar_policy::{Context, Entities, EntityUid, Request};
34use std::str::FromStr;
35
36#[derive(Debug, Clone)]
42pub struct AuthConfig {
43 pub issuer_url: String,
45 pub client_id: String,
47 pub client_secret: Option<String>,
49 pub redirect_uri: String,
51 pub scope: String,
53 pub cookie_name: String,
55 pub dev_config: DevConfig,
57 pub validation_options: JwtValidationOptions,
59 pub pkce_cookie_secret: String,
61}
62
63impl AuthConfig {
64 pub fn from_toml(config_toml: &str) -> Option<Self> {
69 let table: toml::Table = toml::from_str(config_toml).ok()?;
70
71 let dev_config = table.get("dev").and_then(|d| d.as_table()).map(|d| {
73 DevConfig {
74 local_dev_mode: d.get("local_dev_mode").and_then(|v| v.as_bool()).unwrap_or(false),
75 local_dev_email: d.get("local_dev_email").and_then(|v| v.as_str()).map(String::from),
76 local_dev_name: d.get("local_dev_name").and_then(|v| v.as_str()).map(String::from),
77 local_dev_username: d.get("local_dev_username").and_then(|v| v.as_str()).map(String::from),
78 }
79 });
80
81 if let Some(ref dc) = dev_config {
83 if dc.local_dev_mode {
84 let oidc = Self::parse_oidc_section(&table);
86 return Some(Self {
87 issuer_url: oidc.as_ref().map(|o| o.0.clone()).unwrap_or_else(|| "https://auth.example.com".into()),
88 client_id: oidc.as_ref().map(|o| o.1.clone()).unwrap_or_else(|| "trustee".into()),
89 client_secret: oidc.as_ref().and_then(|o| o.2.clone()),
90 redirect_uri: oidc.as_ref().map(|o| o.3.clone()).unwrap_or_else(|| "http://localhost:3000/auth/callback".into()),
91 scope: oidc.as_ref().map(|o| o.4.clone()).unwrap_or_else(|| "openid profile email".into()),
92 cookie_name: "trustee_token".into(),
93 dev_config: dc.clone(),
94 validation_options: JwtValidationOptions::default(),
95 pkce_cookie_secret: oidc.as_ref().map(|o| o.6.clone()).unwrap_or_else(|| "trustee-default-pkce-secret-change-me".into()),
96 });
97 }
98 }
99
100 let (issuer_url, client_id, client_secret, redirect_uri, scope, validation_options, pkce_secret) =
102 Self::parse_oidc_section(&table)?;
103
104 Some(Self {
105 issuer_url,
106 client_id,
107 client_secret,
108 redirect_uri,
109 scope,
110 cookie_name: "trustee_token".into(),
111 dev_config: dev_config.unwrap_or_default(),
112 validation_options,
113 pkce_cookie_secret: pkce_secret,
114 })
115 }
116
117 fn parse_oidc_section(
120 table: &toml::Table,
121 ) -> Option<(String, String, Option<String>, String, String, JwtValidationOptions, String)> {
122 let oidc = table.get("oidc")?.as_table()?;
123
124 let issuer_url = oidc.get("issuer_url")?.as_str()?.to_string();
125 let client_id = oidc.get("client_id")?.as_str()?.to_string();
126 let client_secret = oidc.get("client_secret").and_then(|v| v.as_str()).map(String::from);
127 let redirect_uri = oidc
128 .get("redirect_url")
129 .and_then(|v| v.as_str())
130 .unwrap_or("http://localhost:3000/auth/callback")
131 .to_string();
132 let scope = oidc
133 .get("scope")
134 .and_then(|v| v.as_str())
135 .unwrap_or("openid profile email")
136 .to_string();
137
138 let mut validation_options = JwtValidationOptions::default();
139 if let Some(skip) = oidc.get("skip_issuer_validation").and_then(|v| v.as_bool()) {
140 validation_options.skip_issuer_validation = skip;
141 }
142 if let Some(skip) = oidc.get("skip_audience_validation").and_then(|v| v.as_bool()) {
143 validation_options.skip_audience_validation = skip;
144 }
145 validation_options.expected_audience = oidc
146 .get("expected_audience")
147 .and_then(|v| v.as_str())
148 .map(String::from);
149
150 let pkce_secret = oidc
151 .get("pkce_cookie_secret")
152 .and_then(|v| v.as_str())
153 .unwrap_or("trustee-default-pkce-secret-change-me")
154 .to_string();
155
156 Some((issuer_url, client_id, client_secret, redirect_uri, scope, validation_options, pkce_secret))
157 }
158
159 pub fn oidc_client_config(&self) -> OidcClientConfig {
161 OidcClientConfig {
162 issuer_url: self.issuer_url.clone(),
163 client_id: self.client_id.clone(),
164 client_secret: self.client_secret.clone(),
165 redirect_uri: self.redirect_uri.clone(),
166 scope: self.scope.clone(),
167 code_challenge_method: "S256".to_string(),
168 }
169 }
170}
171
172#[derive(Clone)]
174pub struct AuthState {
175 pub oidc_client: OidcClient,
177 pub resource_server: ResourceServerClient,
179 pub client_config: OidcClientConfig,
181 pub config: AuthConfig,
183 pub pkce_manager: PkceCookieManager,
185 pub session_manager: Arc<WebSessionManager>,
187 pub cedar_authorizer: Option<Arc<pep::cedar::CedarAuthorizer>>,
189}
190
191impl AuthState {
192 pub fn new(config: AuthConfig) -> Self {
194 Self::with_cedar(config, None)
195 }
196
197 pub fn with_cedar(config: AuthConfig, cedar_authorizer: Option<Arc<pep::cedar::CedarAuthorizer>>) -> Self {
199 let pkce_manager = PkceCookieManager::new(
200 config.pkce_cookie_secret.as_bytes(),
201 "trustee_pkce_state",
202 StdDuration::from_secs(600),
203 );
204
205 let session_manager = Arc::new(WebSessionManager::new(
206 OidcClient::new(),
207 config.issuer_url.clone(),
208 config.client_id.clone(),
209 config.client_secret.clone(),
210 config.scope.clone(),
211 ));
212
213 Self {
214 oidc_client: OidcClient::new(),
215 resource_server: ResourceServerClient::new(),
216 client_config: config.oidc_client_config(),
217 pkce_manager,
218 session_manager,
219 config,
220 cedar_authorizer,
221 }
222 }
223
224 pub fn is_dev_mode(&self) -> bool {
226 self.config.dev_config.local_dev_mode
227 }
228
229 pub async fn validate_token(&self, token: &str) -> anyhow::Result<JwtClaims> {
231 let mut claims = self
232 .resource_server
233 .validate_jwt_with_options(
234 token,
235 &self.config.issuer_url,
236 &self.config.client_id,
237 &self.config.validation_options,
238 )
239 .await
240 .map_err(|e| anyhow::anyhow!("Token validation failed: {}", e))?;
241
242 let _ = self
244 .resource_server
245 .enrich_claims_with_userinfo(&mut claims, token, &self.config.issuer_url, None)
246 .await;
247
248 if claims.name.is_none() || claims.email.is_none() {
251 self.fill_userinfo_fields(&mut claims, token).await;
252 }
253
254 Ok(claims)
255 }
256
257 async fn fill_userinfo_fields(&self, claims: &mut JwtClaims, token: &str) {
260 let userinfo_url = format!("{}/userinfo", self.config.issuer_url.trim_end_matches('/'));
264
265 let client = reqwest::Client::new();
266 let resp = client
267 .get(&userinfo_url)
268 .header("Authorization", format!("Bearer {}", token))
269 .header("Accept", "application/json")
270 .send()
271 .await;
272
273 let Ok(resp) = resp else {
274 tracing::debug!("Userinfo request failed for name/email enrichment");
275 return;
276 };
277
278 if !resp.status().is_success() {
279 tracing::debug!("Userinfo returned {} for name/email enrichment", resp.status());
280 return;
281 }
282
283 let Ok(userinfo): Result<serde_json::Map<String, serde_json::Value>, _> = resp.json().await else {
284 return;
285 };
286
287 tracing::debug!("Userinfo keys: {:?}", userinfo.keys().collect::<Vec<_>>());
288
289 if claims.name.is_none() {
290 if let Some(name) = userinfo.get("name").and_then(|v| v.as_str()) {
291 claims.name = Some(name.to_string());
292 }
293 }
294 if claims.email.is_none() {
295 if let Some(email) = userinfo.get("email").and_then(|v| v.as_str()) {
296 claims.email = Some(email.to_string());
297 }
298 }
299 if claims.preferred_username.is_none() {
300 if let Some(uname) = userinfo.get("preferred_username").and_then(|v| v.as_str()) {
301 claims.preferred_username = Some(uname.to_string());
302 }
303 }
304 }
305
306 fn check_cedar_authorized(&self, claims: &JwtClaims) -> Result<(), ()> {
311 let Some(ref authorizer) = self.cedar_authorizer else {
312 return Ok(()); };
314
315 let principal_entity = match pep::cedar::build_principal_entity(claims) {
317 Ok(e) => e,
318 Err(e) => {
319 tracing::error!("Cedar: failed to build principal entity: {}", e);
320 return Err(());
321 }
322 };
323
324 let mut entities_vec = vec![principal_entity];
326
327 let app_uid = match EntityUid::from_str(r#"TrusteeApp::"default""#) {
329 Ok(uid) => uid,
330 Err(e) => {
331 tracing::error!("Cedar: failed to build TrusteeApp uid: {}", e);
332 return Err(());
333 }
334 };
335 let app_entity = match cedar_policy::Entity::new(
336 app_uid,
337 std::collections::HashMap::new(),
338 std::collections::HashSet::new(),
339 ) {
340 Ok(e) => e,
341 Err(e) => {
342 tracing::error!("Cedar: failed to build TrusteeApp entity: {}", e);
343 return Err(());
344 }
345 };
346 entities_vec.push(app_entity);
347
348 let entities = match Entities::from_entities(entities_vec, None) {
349 Ok(e) => e,
350 Err(e) => {
351 tracing::error!("Cedar: failed to build entities set: {}", e);
352 return Err(());
353 }
354 };
355
356 let principal_uid = match pep::cedar::build_principal_uid(claims) {
358 Ok(uid) => uid,
359 Err(e) => {
360 tracing::error!("Cedar: failed to build principal uid: {}", e);
361 return Err(());
362 }
363 };
364
365 let action_uid = match EntityUid::from_str(r#"Action::"Access""#) {
366 Ok(uid) => uid,
367 Err(e) => {
368 tracing::error!("Cedar: failed to build action uid: {}", e);
369 return Err(());
370 }
371 };
372
373 let resource_uid = match EntityUid::from_str(r#"TrusteeApp::"default""#) {
374 Ok(uid) => uid,
375 Err(e) => {
376 tracing::error!("Cedar: failed to build resource uid: {}", e);
377 return Err(());
378 }
379 };
380
381 let request = match Request::new(principal_uid, action_uid, resource_uid, Context::empty(), None) {
382 Ok(r) => r,
383 Err(e) => {
384 tracing::error!("Cedar: failed to build request: {}", e);
385 return Err(());
386 }
387 };
388
389 let response = authorizer.is_allowed_with_entities(&request, &entities);
390
391 if response.allowed() {
392 tracing::debug!(
393 "Cedar: authorized user {} (sub={})",
394 claims.email.as_deref().unwrap_or("unknown"),
395 claims.sub
396 );
397 Ok(())
398 } else {
399 tracing::warn!(
400 "Cedar: DENIED user {} (sub={}) — matched policies: {:?}, errors: {:?}",
401 claims.email.as_deref().unwrap_or("unknown"),
402 claims.sub,
403 response.matched_policies(),
404 response.errors()
405 );
406 Err(())
407 }
408 }
409}
410
411#[derive(Debug, Clone)]
417pub struct AuthUser {
418 pub sub: String,
419 pub email: Option<String>,
420 pub name: Option<String>,
421 pub username: Option<String>,
422 pub is_dev: bool,
423}
424
425impl From<JwtClaims> for AuthUser {
426 fn from(claims: JwtClaims) -> Self {
427 Self {
428 sub: claims.sub,
429 email: claims.email,
430 name: claims.name,
431 username: claims.preferred_username,
432 is_dev: false,
433 }
434 }
435}
436
437const SESSION_COOKIE_MAX_AGE: StdDuration = StdDuration::from_secs(3600);
439
440pub async fn check_auth(
455 auth: &Option<Arc<AuthState>>,
456 headers: &axum::http::HeaderMap,
457) -> Result<Option<String>, StatusCode> {
458 let Some(auth) = auth.as_ref() else {
459 return Ok(None); };
461
462 if let Some(token) = headers
464 .get(header::AUTHORIZATION)
465 .and_then(|v| v.to_str().ok())
466 .and_then(|v| v.strip_prefix("Bearer "))
467 .map(|s| s.to_string())
468 {
469 if token.starts_with("dev:") {
471 if !auth.config.dev_config.local_dev_mode {
472 tracing::warn!("Dev token presented but dev mode is disabled — rejecting");
473 return Err(StatusCode::UNAUTHORIZED);
474 }
475 let parts: Vec<&str> = token.splitn(4, ':').collect();
476 return if parts.len() >= 4 {
477 Ok(None)
478 } else {
479 Err(StatusCode::UNAUTHORIZED)
480 };
481 }
482
483 return match auth.validate_token(&token).await {
484 Ok(claims) => {
485 if auth.check_cedar_authorized(&claims).is_err() {
486 return Err(StatusCode::FORBIDDEN);
487 }
488 Ok(None)
489 }
490 Err(e) => {
491 tracing::warn!("Bearer token validation failed: {}", e);
492 Err(StatusCode::UNAUTHORIZED)
493 }
494 };
495 }
496
497 let session_id = headers
499 .get(header::COOKIE)
500 .and_then(|v| v.to_str().ok())
501 .and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
502
503 let Some(session_id) = session_id else {
504 tracing::warn!("No auth token found in request");
505 return Err(StatusCode::UNAUTHORIZED);
506 };
507
508 if session_id.starts_with("dev:") {
510 if !auth.config.dev_config.local_dev_mode {
511 tracing::warn!("Dev cookie presented but dev mode is disabled — rejecting");
512 return Err(StatusCode::UNAUTHORIZED);
513 }
514 let parts: Vec<&str> = session_id.splitn(4, ':').collect();
515 return if parts.len() >= 4 {
516 Ok(None)
517 } else {
518 Err(StatusCode::UNAUTHORIZED)
519 };
520 }
521
522 match auth.session_manager.get_token(&session_id).await {
524 Ok(access_token) => match auth.validate_token(&access_token).await {
525 Ok(claims) => {
526 if auth.check_cedar_authorized(&claims).is_err() {
528 return Err(StatusCode::FORBIDDEN);
529 }
530 let secure = auth.client_config.redirect_uri.starts_with("https");
532 let cookie = create_auth_cookie(
533 &auth.config.cookie_name,
534 &session_id,
535 SESSION_COOKIE_MAX_AGE,
536 secure,
537 );
538 Ok(Some(cookie.to_string()))
539 }
540 Err(e) => {
541 tracing::warn!("Session token validation failed: {} — attempting force-refresh", e);
544 match auth.session_manager.force_refresh(&session_id).await {
545 Ok(new_token) => match auth.validate_token(&new_token).await {
546 Ok(claims) => {
547 if auth.check_cedar_authorized(&claims).is_err() {
549 return Err(StatusCode::FORBIDDEN);
550 }
551 let secure = auth.client_config.redirect_uri.starts_with("https");
552 let cookie = create_auth_cookie(
553 &auth.config.cookie_name,
554 &session_id,
555 SESSION_COOKIE_MAX_AGE,
556 secure,
557 );
558 Ok(Some(cookie.to_string()))
559 }
560 Err(e2) => {
561 tracing::warn!("Session token still invalid after force-refresh: {}", e2);
562 Err(StatusCode::UNAUTHORIZED)
563 }
564 },
565 Err(e2) => {
566 tracing::warn!("Force-refresh failed: {}", e2);
567 Err(StatusCode::UNAUTHORIZED)
568 }
569 }
570 }
571 },
572 Err(e) => {
573 tracing::warn!("Session lookup/refresh failed: {}", e);
574 Err(StatusCode::UNAUTHORIZED)
575 }
576 }
577}
578
579async fn resolve_access_token(
585 auth: &AuthState,
586 headers: &axum::http::HeaderMap,
587) -> Result<String, StatusCode> {
588 if let Some(token) = headers
590 .get(header::AUTHORIZATION)
591 .and_then(|v| v.to_str().ok())
592 .and_then(|v| v.strip_prefix("Bearer "))
593 .map(|s| s.to_string())
594 {
595 return Ok(token);
596 }
597
598 let session_id = headers
600 .get(header::COOKIE)
601 .and_then(|v| v.to_str().ok())
602 .and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
603
604 match session_id {
605 Some(sid) if sid.starts_with("dev:") => {
606 if !auth.config.dev_config.local_dev_mode {
607 tracing::warn!("Dev cookie in resolve_access_token but dev mode is disabled — rejecting");
608 Err(StatusCode::UNAUTHORIZED)
609 } else {
610 Ok(sid)
611 }
612 }
613 Some(sid) => auth.session_manager.get_token(&sid).await.map_err(|e| {
614 tracing::warn!("Failed to resolve session token: {}", e);
615 StatusCode::UNAUTHORIZED
616 }),
617 None => Err(StatusCode::UNAUTHORIZED),
618 }
619}
620
621fn extract_token_from_cookies(cookie_header: &str, cookie_name: &str) -> Option<String> {
623 for cookie in cookie_header.split(';') {
624 let cookie = cookie.trim();
625 if let Some(value) = cookie.strip_prefix(&format!("{}=", cookie_name)) {
626 return Some(value.to_string());
627 }
628 }
629 None
630}
631
632pub fn auth_routes() -> axum::Router<crate::ServerState> {
638 axum::Router::new()
639 .route("/login", axum::routing::get(login_handler))
640 .route("/callback", axum::routing::get(callback_handler))
641 .route("/me", axum::routing::get(me_handler))
642 .route("/logout", axum::routing::post(logout_handler))
643 .route("/mcp/login", axum::routing::get(mcp_login_handler))
644 .route("/mcp/callback", axum::routing::get(mcp_callback_handler))
645 .route("/mcp/status", axum::routing::get(mcp_status_handler))
646 .route("/mcp/logout", axum::routing::post(mcp_logout_handler))
647}
648
649#[derive(Debug, Deserialize)]
651pub struct CallbackQuery {
652 pub code: Option<String>,
653 pub state: Option<String>,
654 pub error: Option<String>,
655 pub error_description: Option<String>,
656}
657
658async fn login_handler(
660 State(state): State<crate::ServerState>,
661) -> Result<Response, AuthError> {
662 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
663
664 if auth.is_dev_mode() {
666 tracing::info!("Dev mode: creating dev session");
667 let dev = &auth.config.dev_config;
668 let dev_token = format!(
669 "dev:{}:{}:{}",
670 dev.local_dev_email.as_deref().unwrap_or("dev@localhost"),
671 dev.local_dev_name.as_deref().unwrap_or("Dev User"),
672 dev.local_dev_username.as_deref().unwrap_or("dev")
673 );
674 let cookie = create_auth_cookie(&auth.config.cookie_name, &dev_token, StdDuration::from_secs(86400), false);
675 return Ok(Response::builder()
676 .status(StatusCode::FOUND)
677 .header(header::LOCATION, "/")
678 .header(header::SET_COOKIE, cookie.to_string())
679 .body(Body::empty())
680 .unwrap());
681 }
682
683 let pkce_session = auth.pkce_manager.create();
685 let challenge = OidcClient::generate_code_challenge(&pkce_session.verifier);
686
687 let auth_url = auth
688 .oidc_client
689 .build_authorization_url(&auth.client_config, &pkce_session.state, Some(&challenge))
690 .await
691 .map_err(|e| AuthError::OidcError(e.to_string()))?;
692
693 let secure = auth.client_config.redirect_uri.starts_with("https");
697 let pkce_cookie = Cookie::build((
698 auth.pkce_manager.cookie_name().to_string(),
699 pkce_session.cookie_value,
700 ))
701 .path("/")
702 .http_only(true)
703 .same_site(SameSite::Lax)
704 .secure(secure)
705 .max_age(TimeDuration::seconds(auth.pkce_manager.ttl().as_secs() as i64))
706 .build();
707
708 Ok(Response::builder()
709 .status(StatusCode::TEMPORARY_REDIRECT)
710 .header(header::LOCATION, &auth_url)
711 .header(header::SET_COOKIE, pkce_cookie.to_string())
712 .body(Body::empty())
713 .unwrap())
714}
715
716async fn callback_handler(
718 State(state): State<crate::ServerState>,
719 Query(query): Query<CallbackQuery>,
720 headers: axum::http::HeaderMap,
721) -> Result<Response, AuthError> {
722 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
723
724 if let Some(error) = query.error {
726 let desc = query.error_description.unwrap_or_default();
727 tracing::error!("OIDC error: {} - {}", error, desc);
728 return Ok(Redirect::temporary(&format!(
729 "/?error={}&error_description={}",
730 urlencoding::encode(&error),
731 urlencoding::encode(&desc)
732 ))
733 .into_response());
734 }
735
736 let code = query.code.ok_or(AuthError::MissingCode)?;
737 let oauth_state = query.state.ok_or(AuthError::MissingState)?;
738
739 let cookie_header = headers
741 .get(header::COOKIE)
742 .and_then(|v| v.to_str().ok())
743 .unwrap_or("");
744 let pkce_value = extract_token_from_cookies(cookie_header, auth.pkce_manager.cookie_name())
745 .ok_or(AuthError::InvalidState)?;
746
747 let verifier = auth
749 .pkce_manager
750 .verify(&pkce_value, &oauth_state)
751 .ok_or(AuthError::InvalidState)?;
752
753 tracing::info!("Exchanging authorization code for tokens");
755 let token_response = auth
756 .oidc_client
757 .exchange_code_for_tokens(&auth.client_config, &code, Some(&verifier))
758 .await
759 .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
760
761 let session_id = auth
762 .session_manager
763 .create_session(&token_response)
764 .await
765 .map_err(|e| AuthError::TokenExchangeFailed(format!("Session creation failed: {}", e)))?;
766
767 let max_age = SESSION_COOKIE_MAX_AGE;
770
771 let secure = auth.client_config.redirect_uri.starts_with("https");
773 let cookie = create_auth_cookie(&auth.config.cookie_name, &session_id, max_age, secure);
774
775 let clear_pkce = Cookie::build((auth.pkce_manager.cookie_name().to_string(), ""))
777 .path("/")
778 .http_only(true)
779 .same_site(SameSite::Lax)
780 .max_age(TimeDuration::seconds(-1))
781 .build();
782
783 tracing::info!("Authentication successful, redirecting to /");
784
785 Ok(Response::builder()
786 .status(StatusCode::FOUND)
787 .header(header::LOCATION, "/")
788 .header(header::SET_COOKIE, cookie.to_string())
789 .header(header::SET_COOKIE, clear_pkce.to_string())
790 .body(Body::empty())
791 .unwrap())
792}
793
794async fn me_handler(
796 State(state): State<crate::ServerState>,
797 headers: axum::http::HeaderMap,
798) -> Response {
799 let Some(ref auth) = state.auth else {
800 return axum::Json(serde_json::json!({
802 "authenticated": true,
803 "auth_enabled": false
804 }))
805 .into_response();
806 };
807
808 let cookie_header = headers
809 .get(header::COOKIE)
810 .and_then(|v| v.to_str().ok())
811 .unwrap_or("");
812
813 let bearer = headers
815 .get(header::AUTHORIZATION)
816 .and_then(|v| v.to_str().ok())
817 .and_then(|v| v.strip_prefix("Bearer "))
818 .map(String::from);
819
820 let token = bearer.clone().or_else(|| extract_token_from_cookies(cookie_header, &auth.config.cookie_name));
821
822 let Some(cookie_value) = token else {
823 return axum::Json(serde_json::json!({
824 "authenticated": false,
825 "auth_enabled": true
826 }))
827 .into_response();
828 };
829
830 if cookie_value.starts_with("dev:") && auth.config.dev_config.local_dev_mode {
833 let parts: Vec<&str> = cookie_value.splitn(4, ':').collect();
834 if parts.len() >= 4 {
835 return axum::Json(serde_json::json!({
836 "authenticated": true,
837 "auth_enabled": true,
838 "email": parts[1],
839 "name": parts[2],
840 "username": parts[3],
841 "dev_mode": true
842 }))
843 .into_response();
844 }
845 }
846
847 let access_token = if bearer.is_some() {
849 cookie_value
851 } else {
852 match auth.session_manager.get_token(&cookie_value).await {
854 Ok(token) => token,
855 Err(e) => {
856 tracing::debug!("Session token resolution failed for /auth/me: {}", e);
857 return axum::Json(serde_json::json!({
858 "authenticated": false,
859 "auth_enabled": true
860 }))
861 .into_response();
862 }
863 }
864 };
865
866 match auth.validate_token(&access_token).await {
868 Ok(claims) => axum::Json(serde_json::json!({
869 "authenticated": true,
870 "auth_enabled": true,
871 "sub": claims.sub,
872 "email": claims.email,
873 "name": claims.name,
874 "username": claims.preferred_username,
875 "dev_mode": false
876 }))
877 .into_response(),
878 Err(e) => {
879 tracing::debug!("Token validation failed for /auth/me: {}", e);
880 axum::Json(serde_json::json!({
881 "authenticated": false,
882 "auth_enabled": true
883 }))
884 .into_response()
885 }
886 }
887}
888
889async fn logout_handler(
891 State(state): State<crate::ServerState>,
892 headers: axum::http::HeaderMap,
893) -> Response {
894 let cookie_name = state
895 .auth
896 .as_ref()
897 .map(|a| a.config.cookie_name.as_str())
898 .unwrap_or("trustee_token");
899
900 if let Some(ref auth) = state.auth {
902 if let Some(cookie_header) = headers.get(header::COOKIE).and_then(|v| v.to_str().ok()) {
903 if let Some(session_id) = extract_token_from_cookies(cookie_header, cookie_name) {
904 if !session_id.starts_with("dev:") {
905 let _ = auth.session_manager.destroy_session(&session_id);
906 }
907 }
908 }
909 }
910
911 let cookie = Cookie::build((cookie_name.to_string(), ""))
912 .path("/")
913 .http_only(true)
914 .same_site(SameSite::Lax)
915 .max_age(TimeDuration::seconds(-1))
916 .build();
917
918 Response::builder()
919 .status(StatusCode::FOUND)
920 .header(header::LOCATION, "/")
921 .header(header::SET_COOKIE, cookie.to_string())
922 .body(Body::empty())
923 .unwrap()
924}
925
926#[derive(Debug, Deserialize)]
932pub struct McpLoginQuery {
933 pub cred: String,
934}
935
936#[derive(Debug, Deserialize)]
938pub struct McpCallbackQuery {
939 pub code: Option<String>,
940 pub state: Option<String>,
941 pub error: Option<String>,
942 pub error_description: Option<String>,
943}
944
945async fn mcp_login_handler(
950 State(state): State<crate::ServerState>,
951 Query(query): Query<McpLoginQuery>,
952 headers: axum::http::HeaderMap,
953) -> Result<Response, AuthError> {
954 crate::auth::check_auth(&state.auth, &headers)
956 .await
957 .map_err(|_| AuthError::AuthNotConfigured)?;
958
959 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
960
961 let cred_config = load_mcp_credential(&state, &query.cred).await?;
963
964 let (issuer_url, client_id, client_secret, scope) = match &cred_config {
965 McpCredentialInfo::WebInteractive {
966 issuer_url,
967 client_id,
968 client_secret,
969 scope,
970 } => (issuer_url.clone(), client_id.clone(), client_secret.clone(), scope.clone()),
971 _ => {
972 return Ok(Redirect::temporary(&format!(
973 "/?mcp_error={}",
974 urlencoding::encode(&format!("Credential '{}' is not web-interactive type", query.cred))
975 ))
976 .into_response());
977 }
978 };
979
980 let oidc_client = OidcClient::new();
982 let verifier = OidcClient::generate_code_verifier();
983 let challenge = OidcClient::generate_code_challenge(&verifier);
984 let oauth_state = OidcClient::generate_state();
985
986 let mcp_redirect_uri = format!(
988 "{}/auth/mcp/callback",
989 auth.client_config.redirect_uri.trim_end_matches('/').trim_end_matches("/auth/callback")
990 );
991
992 let mcp_client_config = OidcClientConfig {
993 issuer_url: issuer_url.clone(),
994 client_id: client_id.clone(),
995 client_secret: client_secret.clone(),
996 redirect_uri: mcp_redirect_uri.clone(),
997 scope: scope.clone(),
998 code_challenge_method: "S256".to_string(),
999 };
1000
1001 let auth_url = oidc_client
1003 .build_authorization_url(&mcp_client_config, &oauth_state, Some(&challenge))
1004 .await
1005 .map_err(|e| AuthError::OidcError(e.to_string()))?;
1006
1007 mcp_pkce().insert(oauth_state.clone(), verifier.clone(), query.cred.clone()).await;
1009
1010 tracing::info!(
1011 "Initiating MCP browser login for credential '{}' (issuer={})",
1012 query.cred, issuer_url
1013 );
1014
1015 Ok(Response::builder()
1016 .status(StatusCode::TEMPORARY_REDIRECT)
1017 .header(header::LOCATION, &auth_url)
1018 .body(Body::empty())
1019 .unwrap())
1020}
1021
1022async fn mcp_callback_handler(
1024 State(state): State<crate::ServerState>,
1025 Query(query): Query<McpCallbackQuery>,
1026 headers: axum::http::HeaderMap,
1027) -> Result<Response, AuthError> {
1028 let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1029
1030 if let Some(error) = query.error {
1032 let desc = query.error_description.unwrap_or_default();
1033 tracing::error!("MCP OIDC error: {} - {}", error, desc);
1034 return Ok(Redirect::temporary(&format!(
1035 "/?mcp_error={}&error_description={}",
1036 urlencoding::encode(&error),
1037 urlencoding::encode(&desc)
1038 ))
1039 .into_response());
1040 }
1041
1042 let code = query.code.ok_or(AuthError::MissingCode)?;
1043 let oauth_state = query.state.ok_or(AuthError::MissingState)?;
1044
1045 let pkce_data = mcp_pkce().take(&oauth_state).await
1047 .ok_or(AuthError::InvalidState)?;
1048
1049 let verifier = pkce_data.verifier;
1050 let cred_name = &pkce_data.cred_name;
1051
1052 let cred_config = load_mcp_credential(&state, cred_name).await?;
1054
1055 let (issuer_url, client_id, client_secret, scope) = match &cred_config {
1056 McpCredentialInfo::WebInteractive {
1057 issuer_url,
1058 client_id,
1059 client_secret,
1060 scope,
1061 } => (issuer_url.clone(), client_id.clone(), client_secret.clone(), scope.clone()),
1062 _ => {
1063 return Ok(Redirect::temporary(&format!(
1064 "/?mcp_error={}",
1065 urlencoding::encode("Credential is not web-interactive type")
1066 ))
1067 .into_response());
1068 }
1069 };
1070
1071 let mcp_redirect_uri = format!(
1073 "{}/auth/mcp/callback",
1074 auth.client_config.redirect_uri.trim_end_matches('/').trim_end_matches("/auth/callback")
1075 );
1076
1077 let mcp_client_config = OidcClientConfig {
1078 issuer_url: issuer_url.clone(),
1079 client_id: client_id.clone(),
1080 client_secret: client_secret.clone(),
1081 redirect_uri: mcp_redirect_uri,
1082 scope: scope.clone(),
1083 code_challenge_method: "S256".to_string(),
1084 };
1085
1086 tracing::info!("Exchanging MCP authorization code for tokens (credential={})", cred_name);
1088 let oidc_client = OidcClient::new();
1089 let token_response = oidc_client
1090 .exchange_code_for_tokens(&mcp_client_config, &code, Some(&verifier))
1091 .await
1092 .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
1093
1094 let expires_at = {
1096 let now = std::time::SystemTime::now()
1097 .duration_since(std::time::UNIX_EPOCH)
1098 .unwrap_or_default()
1099 .as_secs();
1100 let expires_epoch = now + token_response.expires_in.unwrap_or(900);
1101 let days = expires_epoch / 86400;
1102 let rem = expires_epoch % 86400;
1103 let h = rem / 3600;
1104 let m = (rem % 3600) / 60;
1105 let s = rem % 60;
1106 let z = days as i64 + 719468;
1107 let era = if z >= 0 { z } else { z - 146096 } / 146097;
1108 let doe = (z - era * 146097) as u64;
1109 let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
1110 let y = yoe as i64 + era * 400;
1111 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
1112 let mp = (5 * doy + 2) / 153;
1113 let d = doy - (153 * mp + 2) / 5 + 1;
1114 let mon = if mp < 10 { mp + 3 } else { mp - 9 };
1115 let yr = if mon <= 2 { y + 1 } else { y };
1116 format!("{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z", yr, mon, d, h, m, s)
1117 };
1118
1119 use pep::{FileTokenStore, StoredToken, TokenStore};
1121
1122 let stored = StoredToken::new(
1123 &token_response.access_token,
1124 token_response.refresh_token.clone(),
1125 "Bearer",
1126 &expires_at,
1127 token_response.scope.clone(),
1128 );
1129
1130 let agent_name = {
1131 let user_key = state.resolve_user_key(&headers).await;
1132 let (_sid, session_arc, _, _) = state.ensure_active_session(&user_key).await;
1133 let session = session_arc.lock().await;
1134 session.agent_name.clone()
1135 };
1136 let token_store = FileTokenStore::new(&agent_name);
1137
1138 if let Err(e) = token_store.save(cred_name, &stored) {
1139 tracing::error!("Failed to store MCP token: {}", e);
1140 return Ok(Redirect::temporary(&format!(
1141 "/?mcp_error={}",
1142 urlencoding::encode(&format!("Failed to store token: {}", e))
1143 ))
1144 .into_response());
1145 }
1146
1147 tracing::info!(
1148 "MCP authentication successful for credential '{}' (expires {})",
1149 cred_name, expires_at
1150 );
1151
1152 Ok(Response::builder()
1153 .status(StatusCode::FOUND)
1154 .header(header::LOCATION, format!("/?mcp_connected={}", urlencoding::encode(cred_name)))
1155 .body(Body::empty())
1156 .unwrap())
1157}
1158
1159async fn mcp_status_handler(
1161 State(state): State<crate::ServerState>,
1162 headers: axum::http::HeaderMap,
1163) -> Response {
1164 use pep::{FileTokenStore, TokenStore};
1165
1166 if let Err(code) = crate::auth::check_auth(&state.auth, &headers).await {
1168 return (code, Json(serde_json::json!({"error": "Unauthorized"}))).into_response();
1169 }
1170
1171 let config_toml = {
1173 let user_key = state.resolve_user_key(&headers).await;
1174 let (_sid, session_arc, _, _) = state.ensure_active_session(&user_key).await;
1175 let session = session_arc.lock().await;
1176 match &session.config_toml {
1177 Some(t) => t.clone(),
1178 None => return (StatusCode::INTERNAL_SERVER_ERROR, "Config not loaded").into_response(),
1179 }
1180 };
1181
1182 let mcp_config: toml::Value = match toml::from_str(&config_toml) {
1183 Ok(v) => v,
1184 Err(_) => return Json(serde_json::json!([])).into_response(),
1185 };
1186
1187 let agent_name = {
1188 let user_key = state.resolve_user_key(&headers).await;
1189 let (_sid, session_arc, _, _) = state.ensure_active_session(&user_key).await;
1190 let session = session_arc.lock().await;
1191 session.agent_name.clone()
1192 };
1193 let token_store = FileTokenStore::new(&agent_name);
1194
1195 let servers = mcp_config
1197 .get("mcp")
1198 .and_then(|m| m.get("servers"))
1199 .and_then(|s| s.as_array());
1200 let credentials = mcp_config
1201 .get("mcp")
1202 .and_then(|m| m.get("credentials"))
1203 .and_then(|c| c.as_table());
1204
1205 let mut cred_servers: std::collections::HashMap<String, Vec<String>> = std::collections::HashMap::new();
1206 if let Some(servers) = servers {
1207 for server in servers {
1208 let name = server.get("name").and_then(|n| n.as_str()).unwrap_or("");
1209 let cred_ref = server.get("credentials").and_then(|c| c.as_str()).unwrap_or("");
1210 if !cred_ref.is_empty() {
1211 cred_servers
1212 .entry(cred_ref.to_string())
1213 .or_default()
1214 .push(name.to_string());
1215 }
1216 }
1217 }
1218
1219 let mut result = Vec::new();
1220
1221 if let Some(creds) = credentials {
1222 for (cred_name, cred_config) in creds {
1223 let cred_type = cred_config.get("type").and_then(|t| t.as_str()).unwrap_or("unknown");
1224 let servers_using = cred_servers.get(cred_name).cloned().unwrap_or_default();
1225
1226 if cred_type == "web-session" {
1227 let connected = state.auth.is_some();
1229 result.push(serde_json::json!({
1230 "credential": cred_name,
1231 "type": cred_type,
1232 "connected": connected,
1233 "servers": servers_using,
1234 }));
1235 } else if cred_type == "web-interactive" || cred_type == "interactive" {
1236 let status = match token_store.load(cred_name) {
1238 Ok(Some(token)) => {
1239 let expired = token.is_expired();
1240 serde_json::json!({
1241 "credential": cred_name,
1242 "type": cred_type,
1243 "connected": !expired,
1244 "expires_at": token.expires_at,
1245 "servers": servers_using,
1246 })
1247 }
1248 _ => serde_json::json!({
1249 "credential": cred_name,
1250 "type": cred_type,
1251 "connected": false,
1252 "servers": servers_using,
1253 }),
1254 };
1255 result.push(status);
1256 }
1257 }
1258 }
1259
1260 Json(serde_json::Value::Array(result)).into_response()
1261}
1262
1263async fn mcp_logout_handler(
1265 State(state): State<crate::ServerState>,
1266 Query(query): Query<McpLoginQuery>,
1267 headers: axum::http::HeaderMap,
1268) -> Response {
1269 use pep::{FileTokenStore, TokenStore};
1270
1271 if let Err(code) = crate::auth::check_auth(&state.auth, &headers).await {
1273 return (code, Json(serde_json::json!({"error": "Unauthorized"}))).into_response();
1274 }
1275
1276 let agent_name = {
1277 let user_key = state.resolve_user_key(&headers).await;
1278 let (_sid, session_arc, _, _) = state.ensure_active_session(&user_key).await;
1279 let session = session_arc.lock().await;
1280 session.agent_name.clone()
1281 };
1282 let token_store = FileTokenStore::new(&agent_name);
1283
1284 match token_store.delete(&query.cred) {
1285 Ok(()) => {
1286 tracing::info!("Removed MCP credentials for '{}'", query.cred);
1287 Json(serde_json::json!({"success": true})).into_response()
1288 }
1289 Err(e) => {
1290 tracing::error!("Failed to remove MCP credentials: {}", e);
1291 (
1292 StatusCode::INTERNAL_SERVER_ERROR,
1293 Json(serde_json::json!({"error": e.to_string()})),
1294 )
1295 .into_response()
1296 }
1297 }
1298}
1299
1300struct McpPkceStore {
1307 entries: tokio::sync::Mutex<std::collections::HashMap<String, McpPkceEntry>>,
1308}
1309
1310struct McpPkceEntry {
1311 verifier: String,
1312 cred_name: String,
1313 created_at: std::time::Instant,
1314}
1315
1316impl McpPkceStore {
1317 fn new() -> Self {
1318 Self {
1319 entries: tokio::sync::Mutex::new(std::collections::HashMap::new()),
1320 }
1321 }
1322
1323 async fn insert(&self, state: String, verifier: String, cred_name: String) {
1325 let mut map = self.entries.lock().await;
1326 let cutoff = std::time::Instant::now() - std::time::Duration::from_secs(600);
1328 map.retain(|_, v| v.created_at > cutoff);
1329 map.insert(state, McpPkceEntry {
1330 verifier,
1331 cred_name,
1332 created_at: std::time::Instant::now(),
1333 });
1334 }
1335
1336 async fn take(&self, state: &str) -> Option<McpPkceEntry> {
1338 let mut map = self.entries.lock().await;
1339 map.remove(state)
1340 }
1341}
1342
1343static MCP_PKCE: std::sync::OnceLock<McpPkceStore> = std::sync::OnceLock::new();
1345
1346fn mcp_pkce() -> &'static McpPkceStore {
1348 MCP_PKCE.get_or_init(McpPkceStore::new)
1349}
1350
1351enum McpCredentialInfo {
1353 WebInteractive {
1354 issuer_url: String,
1355 client_id: String,
1356 client_secret: Option<String>,
1357 scope: String,
1358 },
1359 Other(String),
1360}
1361
1362async fn load_mcp_credential(
1364 state: &crate::ServerState,
1365 cred_name: &str,
1366) -> Result<McpCredentialInfo, AuthError> {
1367 let config_toml = state
1368 .config_toml
1369 .clone()
1370 .ok_or(AuthError::AuthNotConfigured)?;
1371
1372 let config: toml::Value = toml::from_str(&config_toml)
1373 .map_err(|e| AuthError::OidcError(format!("Config parse error: {}", e)))?;
1374
1375 let cred = config
1376 .get("mcp")
1377 .and_then(|m| m.get("credentials"))
1378 .and_then(|c| c.as_table())
1379 .and_then(|c| c.get(cred_name))
1380 .ok_or_else(|| AuthError::OidcError(format!("Credential '{}' not found", cred_name)))?;
1381
1382 let cred_type = cred.get("type").and_then(|t| t.as_str()).unwrap_or("unknown");
1383
1384 match cred_type {
1385 "web-interactive" => {
1386 let issuer_url = cred
1387 .get("issuer_url")
1388 .and_then(|v| v.as_str())
1389 .ok_or_else(|| AuthError::OidcError("Missing issuer_url".into()))?
1390 .to_string();
1391 let client_id = cred
1392 .get("client_id")
1393 .and_then(|v| v.as_str())
1394 .ok_or_else(|| AuthError::OidcError("Missing client_id".into()))?
1395 .to_string();
1396 let client_secret = cred
1397 .get("client_secret")
1398 .and_then(|v| v.as_str())
1399 .map(String::from);
1400 let scope = cred
1401 .get("scope")
1402 .and_then(|v| v.as_str())
1403 .unwrap_or("openid profile email")
1404 .to_string();
1405
1406 Ok(McpCredentialInfo::WebInteractive {
1407 issuer_url,
1408 client_id,
1409 client_secret,
1410 scope,
1411 })
1412 }
1413 other => Ok(McpCredentialInfo::Other(other.to_string())),
1414 }
1415}
1416
1417fn create_auth_cookie(name: &str, value: &str, max_age: StdDuration, secure: bool) -> Cookie<'static> {
1423 Cookie::build((name.to_string(), value.to_string()))
1424 .path("/")
1425 .http_only(true)
1426 .same_site(SameSite::Lax)
1427 .secure(secure)
1428 .max_age(TimeDuration::seconds(max_age.as_secs() as i64))
1429 .build()
1430}
1431
1432#[derive(Debug)]
1438pub enum AuthError {
1439 MissingCode,
1440 MissingState,
1441 InvalidState,
1442 OidcError(String),
1443 TokenExchangeFailed(String),
1444 AuthNotConfigured,
1445}
1446
1447impl IntoResponse for AuthError {
1448 fn into_response(self) -> Response {
1449 let (_status, msg) = match self {
1450 AuthError::MissingCode => (StatusCode::BAD_REQUEST, "Missing authorization code"),
1451 AuthError::MissingState => (StatusCode::BAD_REQUEST, "Missing state parameter"),
1452 AuthError::InvalidState => (StatusCode::BAD_REQUEST, "Invalid or expired state"),
1453 AuthError::OidcError(_) => (StatusCode::SERVICE_UNAVAILABLE, "Authentication service error"),
1454 AuthError::TokenExchangeFailed(_) => (StatusCode::BAD_REQUEST, "Token exchange failed"),
1455 AuthError::AuthNotConfigured => (StatusCode::NOT_IMPLEMENTED, "Authentication not configured"),
1456 };
1457 Redirect::temporary(&format!("/?error={}", urlencoding::encode(msg))).into_response()
1458 }
1459}