Skip to main content

trustee_api/
auth.rs

1//! Authentication module for Trustee API.
2//!
3//! Uses PEP for OIDC/OAuth2:
4//! - `ResourceServerClient` for JWT validation (offline, cached JWKS)
5//! - `OidcClient` for authorization code + PKCE login flow
6//! - `PkceCookieManager` for stateless PKCE state (HMAC-signed cookies)
7//! - `DevConfig` for local development bypass
8//!
9//! Two deployment modes:
10//! - **Standalone**: browser hits /auth/login → IdP redirect → /auth/callback → cookie
11//! - **Centralized**: external auth app sends `Authorization: Bearer <token>` directly
12//!
13//! Token extraction order: `Authorization: Bearer` header → `trustee_token` cookie.
14
15use 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// ---------------------------------------------------------------------------
37// Configuration
38// ---------------------------------------------------------------------------
39
40/// Authentication configuration parsed from `[oidc]` and `[dev]` TOML sections.
41#[derive(Debug, Clone)]
42pub struct AuthConfig {
43    /// OIDC provider issuer URL
44    pub issuer_url: String,
45    /// OAuth2 client ID
46    pub client_id: String,
47    /// OAuth2 client secret (None → public client, PKCE only)
48    pub client_secret: Option<String>,
49    /// Redirect URI for OIDC callback
50    pub redirect_uri: String,
51    /// OAuth2 scopes
52    pub scope: String,
53    /// Token cookie name
54    pub cookie_name: String,
55    /// Development mode configuration
56    pub dev_config: DevConfig,
57    /// JWT validation options
58    pub validation_options: JwtValidationOptions,
59    /// Secret for signing PKCE state cookies
60    pub pkce_cookie_secret: String,
61}
62
63impl AuthConfig {
64    /// Parse auth config from the merged trustee TOML string.
65    ///
66    /// Reads `[oidc]` and `[dev]` sections. If neither is present, returns None
67    /// (auth disabled — all endpoints open).
68    pub fn from_toml(config_toml: &str) -> Option<Self> {
69        let table: toml::Table = toml::from_str(config_toml).ok()?;
70
71        // Check for dev mode
72        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        // Dev mode without OIDC — return early with dev-only config
82        if let Some(ref dc) = dev_config {
83            if dc.local_dev_mode {
84                // Try to get OIDC config too (for login endpoint), but it's optional in dev mode
85                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        // Production mode — requires [oidc] section
101        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    /// Parse the `[oidc]` section from a TOML table.
118    /// Returns (issuer_url, client_id, client_secret, redirect_uri, scope, validation_options, pkce_secret).
119    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    /// Build OIDC client configuration for PEP's OidcClient.
160    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/// Shared authentication state, stored in ServerState.
173#[derive(Clone)]
174pub struct AuthState {
175    /// OIDC client for login flow (authorization code + PKCE)
176    pub oidc_client: OidcClient,
177    /// Resource server client for JWT validation (lazy-initialized)
178    pub resource_server: ResourceServerClient,
179    /// OIDC client configuration
180    pub client_config: OidcClientConfig,
181    /// Auth configuration
182    pub config: AuthConfig,
183    /// Stateless PKCE cookie manager
184    pub pkce_manager: PkceCookieManager,
185    /// Web session manager (cookie session_id → server-side token with auto-refresh)
186    pub session_manager: Arc<WebSessionManager>,
187    /// Cedar authorizer for ABAC authorization (None = Cedar disabled)
188    pub cedar_authorizer: Option<Arc<pep::cedar::CedarAuthorizer>>,
189}
190
191impl AuthState {
192    /// Create new auth state from configuration.
193    pub fn new(config: AuthConfig) -> Self {
194        Self::with_cedar(config, None)
195    }
196
197    /// Create new auth state with optional Cedar authorizer.
198    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    /// Check if development mode is enabled.
225    pub fn is_dev_mode(&self) -> bool {
226        self.config.dev_config.local_dev_mode
227    }
228
229    /// Validate a JWT token using PEP's ResourceServerClient.
230    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        // Enrich with userinfo for role/groups (cached, no-op if already present)
243        let _ = self
244            .resource_server
245            .enrich_claims_with_userinfo(&mut claims, token, &self.config.issuer_url, None)
246            .await;
247
248        // PEP only merges groups/role from userinfo. If name/email are missing
249        // (Kanidm JWTs only contain sub), fetch them from userinfo directly.
250        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    /// Fetch name/email/preferred_username from the OIDC userinfo endpoint
258    /// and fill in any that are missing from the JWT claims.
259    async fn fill_userinfo_fields(&self, claims: &mut JwtClaims, token: &str) {
260        // Derive userinfo URL from issuer
261        // For Kanidm: issuer_url is the discovery endpoint,
262        // userinfo is at {issuer_url}/userinfo
263        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    /// Check Cedar authorization for the authenticated user.
307    ///
308    /// Returns Ok(()) if allowed (or if Cedar is not configured).
309    /// Returns Err(()) if denied — caller should return 403 Forbidden.
310    fn check_cedar_authorized(&self, claims: &JwtClaims) -> Result<(), ()> {
311        let Some(ref authorizer) = self.cedar_authorizer else {
312            return Ok(()); // Cedar not configured — allow
313        };
314
315        // Build principal entity from JWT claims
316        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        // Build entities set with principal + TrusteeApp resource
325        let mut entities_vec = vec![principal_entity];
326
327        // Add a TrusteeApp entity as the resource
328        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        // Build the Cedar authorization request
357        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// ---------------------------------------------------------------------------
412// Auth checking — called by protected route handlers
413// ---------------------------------------------------------------------------
414
415/// Authenticated user info extracted from the token.
416#[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
437/// Cookie max-age for session cookies (1 hour, matching the server-side idle timeout).
438const SESSION_COOKIE_MAX_AGE: StdDuration = StdDuration::from_secs(3600);
439
440/// Check authentication for a protected endpoint.
441///
442/// Returns `Ok(None)` if auth is not configured (open mode), or if a valid
443/// token is present without needing cookie renewal. Returns `Ok(Some(cookie))`
444/// if auth succeeded and the caller should include the given `Set-Cookie`
445/// header value in the response (rolling session). Returns `Err(StatusCode)`
446/// if auth is configured but no valid token is found.
447///
448/// Token sources (in order):
449/// 1. `Authorization: Bearer <token>` header (raw JWT — validated directly)
450/// 2. `trustee_token=<session_id>` cookie (looked up in WebSessionManager,
451///    auto-refreshed if near expiry)
452///
453/// Dev mode tokens use the format `dev:email:name:username`.
454pub 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); // Auth not configured — allow
460    };
461
462    // 1. Try Bearer header first (raw JWT — e.g. from API clients, Torpi proxy)
463    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        // Dev mode token — only accepted when dev mode is currently enabled
470        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    // 2. Try cookie (session_id → WebSessionManager → access token with auto-refresh)
498    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    // Dev mode token in cookie — only accepted when dev mode is currently enabled
509    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    // Session-based: look up via WebSessionManager (auto-refreshes)
523    match auth.session_manager.get_token(&session_id).await {
524        Ok(access_token) => match auth.validate_token(&access_token).await {
525            Ok(claims) => {
526                // Cedar authorization check
527                if auth.check_cedar_authorized(&claims).is_err() {
528                    return Err(StatusCode::FORBIDDEN);
529                }
530                // Roll the cookie — reset max-age so active users stay logged in
531                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                // Token was returned but JWT validation failed (e.g. ExpiredSignature
542                // due to clock skew). Force-refresh and retry once.
543                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                            // Cedar authorization check
548                            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
579/// Extract a valid access token from the request (for use by handlers that
580/// need the token itself, not just auth checking).
581///
582/// Resolves session_id cookies to actual access tokens via WebSessionManager.
583/// Bearer headers are returned as-is.
584async fn resolve_access_token(
585    auth: &AuthState,
586    headers: &axum::http::HeaderMap,
587) -> Result<String, StatusCode> {
588    // Bearer header — return as-is
589    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    // Cookie — resolve session_id → access_token
599    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
621/// Extract token value from a cookie header string.
622fn 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
632// ---------------------------------------------------------------------------
633// Auth routes: /auth/login, /auth/callback, /auth/me, /auth/logout
634// ---------------------------------------------------------------------------
635
636/// Build the auth routes as a nested Router.
637pub 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/// Query parameters for OIDC callback.
650#[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
658/// GET /auth/login — initiate OIDC login with PKCE, or create dev session.
659async 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    // Dev mode — create synthetic session
665    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    // Production — redirect to IdP with PKCE
684    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    // Set PKCE state cookie (HttpOnly, SameSite=Lax)
694    // Secure flag follows the redirect_uri scheme — HTTP localhost/LAN must not
695    // set Secure or the browser drops the cookie and PKCE state is lost.
696    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
716/// GET /auth/callback — exchange authorization code for tokens, set cookie.
717async 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    // Check for errors from IdP
725    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    // Retrieve PKCE cookie
740    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    // Verify PKCE cookie (HMAC + expiry + state match)
748    let verifier = auth
749        .pkce_manager
750        .verify(&pkce_value, &oauth_state)
751        .ok_or(AuthError::InvalidState)?;
752
753    // Exchange code for tokens
754    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    // Cookie lifetime matches server-side idle timeout (1 hour).
768    // The cookie is rolled on every successful request via check_auth().
769    let max_age = SESSION_COOKIE_MAX_AGE;
770
771    // Set auth cookie — Secure only when redirect_uri is HTTPS
772    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    // Clear PKCE cookie (single-use)
776    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
794/// GET /auth/me — return current user info.
795async 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        // Auth not configured — always authenticated (no auth required)
801        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    // Also try Authorization: Bearer header
814    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    // Dev mode token (stored directly in cookie, no session manager)
831    // Only report as authenticated when dev mode is currently enabled
832    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    // Bearer header = raw JWT; Cookie value = session_id → resolve to access token
848    let access_token = if bearer.is_some() {
849        // Already have the raw token from Bearer header
850        cookie_value
851    } else {
852        // Cookie value is a session_id — resolve via WebSessionManager
853        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    // Real JWT — validate and return claims
867    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
889/// POST /auth/logout — destroy session and clear auth cookie.
890async 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    // Destroy the session on the server side
901    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// ---------------------------------------------------------------------------
927// MCP auth routes: /auth/mcp/login, /callback, /status, /logout (C2)
928// ---------------------------------------------------------------------------
929
930/// Query parameters for MCP login initiation.
931#[derive(Debug, Deserialize)]
932pub struct McpLoginQuery {
933    pub cred: String,
934}
935
936/// Query parameters for MCP OIDC callback.
937#[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
945/// GET /auth/mcp/login?cred=<name> — initiate per-server OIDC PKCE login.
946///
947/// Reads the credential config from the session's config_toml, verifies it's
948/// `type = "web-interactive"`, then redirects to the OIDC provider.
949async 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    // Require authentication — user must be logged into trustee-web
955    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    // Parse MCP credential config from session's config_toml
962    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    // Build PKCE pair using a separate PkceCookieManager for MCP
981    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    // Build OidcClientConfig for the MCP credential's OIDC client
987    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    // Build authorization URL
1002    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    // Store PKCE state + credential name in the in-memory map
1008    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
1022/// GET /auth/mcp/callback — handle MCP OIDC callback, store tokens.
1023async 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    // Check for errors from IdP
1031    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    // Look up PKCE verifier + credential name from in-memory store
1046    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    // Parse the MCP credential config to get OIDC settings for token exchange
1053    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    // Build redirect URI (must match what was used in login)
1072    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    // Exchange code for tokens
1087    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    // Compute expires_at
1095    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    // Store via FileTokenStore (same as `trustee mcp auth`)
1120    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
1159/// GET /auth/mcp/status — return connection status for all MCP credentials.
1160async fn mcp_status_handler(
1161    State(state): State<crate::ServerState>,
1162    headers: axum::http::HeaderMap,
1163) -> Response {
1164    use pep::{FileTokenStore, TokenStore};
1165
1166    // Require auth
1167    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    // Parse MCP config from session
1172    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    // Build server → credential mapping
1196    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                // Session credentials are always "connected" if auth is enabled
1228                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                // Check token store
1237                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
1263/// POST /auth/mcp/logout?cred=<name> — remove stored MCP tokens.
1264async 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    // Require auth
1272    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
1300// ---------------------------------------------------------------------------
1301// MCP auth helpers
1302// ---------------------------------------------------------------------------
1303
1304/// In-memory store for MCP PKCE state (state token → verifier + credential name).
1305/// Entries expire after 10 minutes. Not persisted across restarts.
1306struct 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    /// Insert a PKCE entry. Cleans up entries older than 10 minutes.
1324    async fn insert(&self, state: String, verifier: String, cred_name: String) {
1325        let mut map = self.entries.lock().await;
1326        // Cleanup expired entries (older than 10 min)
1327        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    /// Take and remove a PKCE entry (single-use).
1337    async fn take(&self, state: &str) -> Option<McpPkceEntry> {
1338        let mut map = self.entries.lock().await;
1339        map.remove(state)
1340    }
1341}
1342
1343/// Global singleton PKCE store for MCP browser logins.
1344static MCP_PKCE: std::sync::OnceLock<McpPkceStore> = std::sync::OnceLock::new();
1345
1346/// Get or initialize the global MCP PKCE store.
1347fn mcp_pkce() -> &'static McpPkceStore {
1348    MCP_PKCE.get_or_init(McpPkceStore::new)
1349}
1350
1351/// Simplified MCP credential info (parsed from TOML).
1352enum 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
1362/// Load a specific MCP credential from the session's config_toml.
1363async 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
1417// ---------------------------------------------------------------------------
1418// Helpers
1419// ---------------------------------------------------------------------------
1420
1421/// Create an HttpOnly auth cookie.
1422fn 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// ---------------------------------------------------------------------------
1433// Error handling
1434// ---------------------------------------------------------------------------
1435
1436/// Authentication errors.
1437#[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}