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