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_uri")
129            .or_else(|| oidc.get("redirect_url")) // backward compat
130            .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    /// Build OIDC client configuration for PEP's OidcClient.
161    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/// Shared authentication state, stored in ServerState.
174#[derive(Clone)]
175pub struct AuthState {
176    /// OIDC client for login flow (authorization code + PKCE)
177    pub oidc_client: OidcClient,
178    /// Resource server client for JWT validation (lazy-initialized)
179    pub resource_server: ResourceServerClient,
180    /// OIDC client configuration
181    pub client_config: OidcClientConfig,
182    /// Auth configuration
183    pub config: AuthConfig,
184    /// Stateless PKCE cookie manager
185    pub pkce_manager: PkceCookieManager,
186    /// Web session manager (cookie session_id → server-side token with auto-refresh)
187    pub session_manager: Arc<WebSessionManager>,
188    /// Cedar authorizer for ABAC authorization (None = Cedar disabled)
189    pub cedar_authorizer: Option<Arc<pep::cedar::CedarAuthorizer>>,
190}
191
192impl AuthState {
193    /// Create new auth state from configuration.
194    pub fn new(config: AuthConfig) -> Self {
195        Self::with_cedar(config, None)
196    }
197
198    /// Create new auth state with optional Cedar authorizer.
199    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    /// Check if development mode is enabled.
226    pub fn is_dev_mode(&self) -> bool {
227        self.config.dev_config.local_dev_mode
228    }
229
230    /// Validate a JWT token using PEP's ResourceServerClient.
231    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        // Enrich with userinfo for role/groups (cached, no-op if already present)
244        let _ = self
245            .resource_server
246            .enrich_claims_with_userinfo(&mut claims, token, &self.config.issuer_url, None)
247            .await;
248
249        // PEP only merges groups/role from userinfo. If name/email are missing
250        // (Kanidm JWTs only contain sub), fetch them from userinfo directly.
251        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    /// Fetch name/email/preferred_username from the OIDC userinfo endpoint
259    /// and fill in any that are missing from the JWT claims.
260    async fn fill_userinfo_fields(&self, claims: &mut JwtClaims, token: &str) {
261        // Derive userinfo URL from issuer
262        // For Kanidm: issuer_url is the discovery endpoint,
263        // userinfo is at {issuer_url}/userinfo
264        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    /// Check Cedar authorization for the authenticated user.
308    ///
309    /// Returns Ok(()) if allowed (or if Cedar is not configured).
310    /// Returns Err(()) if denied — caller should return 403 Forbidden.
311    fn check_cedar_authorized(&self, claims: &JwtClaims) -> Result<(), ()> {
312        let Some(ref authorizer) = self.cedar_authorizer else {
313            return Ok(()); // Cedar not configured — allow
314        };
315
316        // Build principal entity from JWT claims
317        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        // Build entities set with principal + TrusteeApp resource
326        let mut entities_vec = vec![principal_entity];
327
328        // Add a TrusteeApp entity as the resource
329        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        // Build the Cedar authorization request
358        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// ---------------------------------------------------------------------------
413// Auth checking — called by protected route handlers
414// ---------------------------------------------------------------------------
415
416/// Authenticated user info extracted from the token.
417#[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
438/// Cookie max-age for session cookies (1 hour, matching the server-side idle timeout).
439const SESSION_COOKIE_MAX_AGE: StdDuration = StdDuration::from_secs(3600);
440
441/// Extract a user key from a dev-mode token string (`dev:email:name:username`).
442/// Returns `dev:{email}` on success.
443fn dev_user_key(token: &str) -> Option<String> {
444    let parts: Vec<&str> = token.splitn(4, ':').collect();
445    if parts.len() >= 4 {
446        Some(format!("dev:{}", parts[1]))
447    } else {
448        None
449    }
450}
451
452/// Check authentication for a protected endpoint.
453///
454/// Returns `Ok((None, user_key))` if auth is not configured (open mode), or if
455/// a valid token is present without needing cookie renewal. Returns
456/// `Ok((Some(cookie), user_key))` if auth succeeded and the caller should
457/// include the given `Set-Cookie` header value in the response (rolling session).
458/// Returns `Err(StatusCode)` if auth is configured but no valid token is found.
459///
460/// The returned `user_key` is the identity string used for session isolation
461/// (JWT `sub` claim for real users, `dev:{email}` for dev mode, `"default"`
462/// when auth is not configured). This avoids the need for handlers to call
463/// `resolve_user_key()` which would re-validate the JWT a second time.
464///
465/// Token sources (in order):
466/// 1. `Authorization: Bearer <token>` header (raw JWT — validated directly)
467/// 2. `trustee_token=<session_id>` cookie (looked up in WebSessionManager,
468///    auto-refreshed if near expiry)
469///
470/// Dev mode tokens use the format `dev:email:name:username`.
471pub async fn check_auth(
472    auth: &Option<Arc<AuthState>>,
473    headers: &axum::http::HeaderMap,
474) -> Result<(Option<String>, String), StatusCode> {
475    let Some(auth) = auth.as_ref() else {
476        return Ok((None, "default".to_string())); // Auth not configured — allow
477    };
478
479    // 1. Try Bearer header first (raw JWT — e.g. from API clients, Torpi proxy)
480    if let Some(token) = headers
481        .get(header::AUTHORIZATION)
482        .and_then(|v| v.to_str().ok())
483        .and_then(|v| v.strip_prefix("Bearer "))
484        .map(|s| s.to_string())
485    {
486        // Dev mode token — only accepted when dev mode is currently enabled
487        if token.starts_with("dev:") {
488            if !auth.config.dev_config.local_dev_mode {
489                tracing::warn!("Dev token presented but dev mode is disabled — rejecting");
490                return Err(StatusCode::UNAUTHORIZED);
491            }
492            return match dev_user_key(&token) {
493                Some(key) => Ok((None, key)),
494                None => Err(StatusCode::UNAUTHORIZED),
495            };
496        }
497
498        return match auth.validate_token(&token).await {
499            Ok(claims) => {
500                if auth.check_cedar_authorized(&claims).is_err() {
501                    return Err(StatusCode::FORBIDDEN);
502                }
503                Ok((None, claims.sub))
504            }
505            Err(e) => {
506                tracing::warn!("Bearer token validation failed: {}", e);
507                Err(StatusCode::UNAUTHORIZED)
508            }
509        };
510    }
511
512    // 2. Try cookie (session_id → WebSessionManager → access token with auto-refresh)
513    let session_id = headers
514        .get(header::COOKIE)
515        .and_then(|v| v.to_str().ok())
516        .and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
517
518    let Some(session_id) = session_id else {
519        tracing::warn!("No auth token found in request");
520        return Err(StatusCode::UNAUTHORIZED);
521    };
522
523    // Dev mode token in cookie — only accepted when dev mode is currently enabled
524    if session_id.starts_with("dev:") {
525        if !auth.config.dev_config.local_dev_mode {
526            tracing::warn!("Dev cookie presented but dev mode is disabled — rejecting");
527            return Err(StatusCode::UNAUTHORIZED);
528        }
529        return match dev_user_key(&session_id) {
530            Some(key) => Ok((None, key)),
531            None => Err(StatusCode::UNAUTHORIZED),
532        };
533    }
534
535    // Session-based: look up via WebSessionManager (auto-refreshes)
536    match auth.session_manager.get_token(&session_id).await {
537        Ok(access_token) => match auth.validate_token(&access_token).await {
538            Ok(claims) => {
539                // Cedar authorization check
540                if auth.check_cedar_authorized(&claims).is_err() {
541                    return Err(StatusCode::FORBIDDEN);
542                }
543                // Roll the cookie — reset max-age so active users stay logged in
544                let secure = auth.client_config.redirect_uri.starts_with("https");
545                let cookie = create_auth_cookie(
546                    &auth.config.cookie_name,
547                    &session_id,
548                    SESSION_COOKIE_MAX_AGE,
549                    secure,
550                );
551                Ok((Some(cookie.to_string()), claims.sub))
552            }
553            Err(e) => {
554                // Token was returned but JWT validation failed (e.g. ExpiredSignature
555                // due to clock skew). Force-refresh and retry once.
556                tracing::warn!("Session token validation failed: {} — attempting force-refresh", e);
557                match auth.session_manager.force_refresh(&session_id).await {
558                    Ok(new_token) => match auth.validate_token(&new_token).await {
559                        Ok(claims) => {
560                            // Cedar authorization check
561                            if auth.check_cedar_authorized(&claims).is_err() {
562                                return Err(StatusCode::FORBIDDEN);
563                            }
564                            let secure = auth.client_config.redirect_uri.starts_with("https");
565                            let cookie = create_auth_cookie(
566                                &auth.config.cookie_name,
567                                &session_id,
568                                SESSION_COOKIE_MAX_AGE,
569                                secure,
570                            );
571                            Ok((Some(cookie.to_string()), claims.sub))
572                        }
573                        Err(e2) => {
574                            tracing::warn!("Session token still invalid after force-refresh: {}", e2);
575                            Err(StatusCode::UNAUTHORIZED)
576                        }
577                    },
578                    Err(e2) => {
579                        tracing::warn!("Force-refresh failed: {}", e2);
580                        Err(StatusCode::UNAUTHORIZED)
581                    }
582                }
583            }
584        },
585        Err(e) => {
586            tracing::warn!("Session lookup/refresh failed: {}", e);
587            Err(StatusCode::UNAUTHORIZED)
588        }
589    }
590}
591
592/// Extract a valid access token from the request (for use by handlers that
593/// need the token itself, not just auth checking).
594///
595/// Resolves session_id cookies to actual access tokens via WebSessionManager.
596/// Bearer headers are returned as-is.
597async fn resolve_access_token(
598    auth: &AuthState,
599    headers: &axum::http::HeaderMap,
600) -> Result<String, StatusCode> {
601    // Bearer header — return as-is
602    if let Some(token) = headers
603        .get(header::AUTHORIZATION)
604        .and_then(|v| v.to_str().ok())
605        .and_then(|v| v.strip_prefix("Bearer "))
606        .map(|s| s.to_string())
607    {
608        return Ok(token);
609    }
610
611    // Cookie — resolve session_id → access_token
612    let session_id = headers
613        .get(header::COOKIE)
614        .and_then(|v| v.to_str().ok())
615        .and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
616
617    match session_id {
618        Some(sid) if sid.starts_with("dev:") => {
619            if !auth.config.dev_config.local_dev_mode {
620                tracing::warn!("Dev cookie in resolve_access_token but dev mode is disabled — rejecting");
621                Err(StatusCode::UNAUTHORIZED)
622            } else {
623                Ok(sid)
624            }
625        }
626        Some(sid) => auth.session_manager.get_token(&sid).await.map_err(|e| {
627            tracing::warn!("Failed to resolve session token: {}", e);
628            StatusCode::UNAUTHORIZED
629        }),
630        None => Err(StatusCode::UNAUTHORIZED),
631    }
632}
633
634/// Extract token value from a cookie header string.
635fn extract_token_from_cookies(cookie_header: &str, cookie_name: &str) -> Option<String> {
636    for cookie in cookie_header.split(';') {
637        let cookie = cookie.trim();
638        if let Some(value) = cookie.strip_prefix(&format!("{}=", cookie_name)) {
639            return Some(value.to_string());
640        }
641    }
642    None
643}
644
645// ---------------------------------------------------------------------------
646// Auth routes: /auth/login, /auth/callback, /auth/me, /auth/logout
647// ---------------------------------------------------------------------------
648
649/// Build the auth routes as a nested Router.
650pub fn auth_routes() -> axum::Router<crate::ServerState> {
651    axum::Router::new()
652        .route("/login", axum::routing::get(login_handler))
653        .route("/callback", axum::routing::get(callback_handler))
654        .route("/me", axum::routing::get(me_handler))
655        .route("/logout", axum::routing::post(logout_handler))
656        .route("/mcp/login", axum::routing::get(mcp_login_handler))
657        .route("/mcp/callback", axum::routing::get(mcp_callback_handler))
658        .route("/mcp/status", axum::routing::get(mcp_status_handler))
659        .route("/mcp/logout", axum::routing::post(mcp_logout_handler))
660}
661
662/// Query parameters for OIDC callback.
663#[derive(Debug, Deserialize)]
664pub struct CallbackQuery {
665    pub code: Option<String>,
666    pub state: Option<String>,
667    pub error: Option<String>,
668    pub error_description: Option<String>,
669}
670
671/// GET /auth/login — initiate OIDC login with PKCE, or create dev session.
672async fn login_handler(
673    State(state): State<crate::ServerState>,
674) -> Result<Response, AuthError> {
675    let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
676
677    // Dev mode — create synthetic session
678    if auth.is_dev_mode() {
679        tracing::info!("Dev mode: creating dev session");
680        let dev = &auth.config.dev_config;
681        let dev_token = format!(
682            "dev:{}:{}:{}",
683            dev.local_dev_email.as_deref().unwrap_or("dev@localhost"),
684            dev.local_dev_name.as_deref().unwrap_or("Dev User"),
685            dev.local_dev_username.as_deref().unwrap_or("dev")
686        );
687        let cookie = create_auth_cookie(&auth.config.cookie_name, &dev_token, StdDuration::from_secs(86400), false);
688        return Ok(Response::builder()
689            .status(StatusCode::FOUND)
690            .header(header::LOCATION, "/")
691            .header(header::SET_COOKIE, cookie.to_string())
692            .body(Body::empty())
693            .unwrap());
694    }
695
696    // Production — redirect to IdP with PKCE
697    let pkce_session = auth.pkce_manager.create();
698    let challenge = OidcClient::generate_code_challenge(&pkce_session.verifier);
699
700    let auth_url = auth
701        .oidc_client
702        .build_authorization_url(&auth.client_config, &pkce_session.state, Some(&challenge))
703        .await
704        .map_err(|e| AuthError::OidcError(e.to_string()))?;
705
706    // Set PKCE state cookie (HttpOnly, SameSite=Lax)
707    // Secure flag follows the redirect_uri scheme — HTTP localhost/LAN must not
708    // set Secure or the browser drops the cookie and PKCE state is lost.
709    let secure = auth.client_config.redirect_uri.starts_with("https");
710    let pkce_cookie = Cookie::build((
711        auth.pkce_manager.cookie_name().to_string(),
712        pkce_session.cookie_value,
713    ))
714        .path("/")
715        .http_only(true)
716        .same_site(SameSite::Lax)
717        .secure(secure)
718        .max_age(TimeDuration::seconds(auth.pkce_manager.ttl().as_secs() as i64))
719        .build();
720
721    Ok(Response::builder()
722        .status(StatusCode::TEMPORARY_REDIRECT)
723        .header(header::LOCATION, &auth_url)
724        .header(header::SET_COOKIE, pkce_cookie.to_string())
725        .body(Body::empty())
726        .unwrap())
727}
728
729/// GET /auth/callback — exchange authorization code for tokens, set cookie.
730async fn callback_handler(
731    State(state): State<crate::ServerState>,
732    Query(query): Query<CallbackQuery>,
733    headers: axum::http::HeaderMap,
734) -> Result<Response, AuthError> {
735    let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
736
737    // Check for errors from IdP
738    if let Some(error) = query.error {
739        let desc = query.error_description.unwrap_or_default();
740        tracing::error!("OIDC error: {} - {}", error, desc);
741        return Ok(Redirect::temporary(&format!(
742            "/?error={}&error_description={}",
743            urlencoding::encode(&error),
744            urlencoding::encode(&desc)
745        ))
746        .into_response());
747    }
748
749    let code = query.code.ok_or(AuthError::MissingCode)?;
750    let oauth_state = query.state.ok_or(AuthError::MissingState)?;
751
752    // Retrieve PKCE cookie
753    let cookie_header = headers
754        .get(header::COOKIE)
755        .and_then(|v| v.to_str().ok())
756        .unwrap_or("");
757    let pkce_value = extract_token_from_cookies(cookie_header, auth.pkce_manager.cookie_name())
758        .ok_or(AuthError::InvalidState)?;
759
760    // Verify PKCE cookie (HMAC + expiry + state match)
761    let verifier = auth
762        .pkce_manager
763        .verify(&pkce_value, &oauth_state)
764        .ok_or(AuthError::InvalidState)?;
765
766    // Exchange code for tokens
767    tracing::info!("Exchanging authorization code for tokens");
768    let token_response = auth
769        .oidc_client
770        .exchange_code_for_tokens(&auth.client_config, &code, Some(&verifier))
771        .await
772        .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
773
774    let session_id = auth
775        .session_manager
776        .create_session(&token_response)
777        .await
778        .map_err(|e| AuthError::TokenExchangeFailed(format!("Session creation failed: {}", e)))?;
779
780    // Cookie lifetime matches server-side idle timeout (1 hour).
781    // The cookie is rolled on every successful request via check_auth().
782    let max_age = SESSION_COOKIE_MAX_AGE;
783
784    // Set auth cookie — Secure only when redirect_uri is HTTPS
785    let secure = auth.client_config.redirect_uri.starts_with("https");
786    let cookie = create_auth_cookie(&auth.config.cookie_name, &session_id, max_age, secure);
787
788    // Clear PKCE cookie (single-use)
789    let clear_pkce = Cookie::build((auth.pkce_manager.cookie_name().to_string(), ""))
790        .path("/")
791        .http_only(true)
792        .same_site(SameSite::Lax)
793        .max_age(TimeDuration::seconds(-1))
794        .build();
795
796    tracing::info!("Authentication successful, redirecting to /");
797
798    Ok(Response::builder()
799        .status(StatusCode::FOUND)
800        .header(header::LOCATION, "/")
801        .header(header::SET_COOKIE, cookie.to_string())
802        .header(header::SET_COOKIE, clear_pkce.to_string())
803        .body(Body::empty())
804        .unwrap())
805}
806
807/// GET /auth/me — return current user info.
808async fn me_handler(
809    State(state): State<crate::ServerState>,
810    headers: axum::http::HeaderMap,
811) -> Response {
812    let Some(ref auth) = state.auth else {
813        // Auth not configured — always authenticated (no auth required)
814        return axum::Json(serde_json::json!({
815            "authenticated": true,
816            "auth_enabled": false
817        }))
818        .into_response();
819    };
820
821    let cookie_header = headers
822        .get(header::COOKIE)
823        .and_then(|v| v.to_str().ok())
824        .unwrap_or("");
825
826    // Also try Authorization: Bearer header
827    let bearer = headers
828        .get(header::AUTHORIZATION)
829        .and_then(|v| v.to_str().ok())
830        .and_then(|v| v.strip_prefix("Bearer "))
831        .map(String::from);
832
833    let token = bearer.clone().or_else(|| extract_token_from_cookies(cookie_header, &auth.config.cookie_name));
834
835    let Some(cookie_value) = token else {
836        return axum::Json(serde_json::json!({
837            "authenticated": false,
838            "auth_enabled": true
839        }))
840        .into_response();
841    };
842
843    // Dev mode token (stored directly in cookie, no session manager)
844    // Only report as authenticated when dev mode is currently enabled
845    if cookie_value.starts_with("dev:") && auth.config.dev_config.local_dev_mode {
846        let parts: Vec<&str> = cookie_value.splitn(4, ':').collect();
847        if parts.len() >= 4 {
848            return axum::Json(serde_json::json!({
849                "authenticated": true,
850                "auth_enabled": true,
851                "email": parts[1],
852                "name": parts[2],
853                "username": parts[3],
854                "dev_mode": true
855            }))
856            .into_response();
857        }
858    }
859
860    // Bearer header = raw JWT; Cookie value = session_id → resolve to access token
861    let access_token = if bearer.is_some() {
862        // Already have the raw token from Bearer header
863        cookie_value
864    } else {
865        // Cookie value is a session_id — resolve via WebSessionManager
866        match auth.session_manager.get_token(&cookie_value).await {
867            Ok(token) => token,
868            Err(e) => {
869                tracing::debug!("Session token resolution failed for /auth/me: {}", e);
870                return axum::Json(serde_json::json!({
871                    "authenticated": false,
872                    "auth_enabled": true
873                }))
874                .into_response();
875            }
876        }
877    };
878
879    // Real JWT — validate and return claims
880    match auth.validate_token(&access_token).await {
881        Ok(claims) => axum::Json(serde_json::json!({
882            "authenticated": true,
883            "auth_enabled": true,
884            "sub": claims.sub,
885            "email": claims.email,
886            "name": claims.name,
887            "username": claims.preferred_username,
888            "dev_mode": false
889        }))
890        .into_response(),
891        Err(e) => {
892            tracing::debug!("Token validation failed for /auth/me: {}", e);
893            axum::Json(serde_json::json!({
894                "authenticated": false,
895                "auth_enabled": true
896            }))
897            .into_response()
898        }
899    }
900}
901
902/// POST /auth/logout — destroy session and clear auth cookie.
903async fn logout_handler(
904    State(state): State<crate::ServerState>,
905    headers: axum::http::HeaderMap,
906) -> Response {
907    let cookie_name = state
908        .auth
909        .as_ref()
910        .map(|a| a.config.cookie_name.as_str())
911        .unwrap_or("trustee_token");
912
913    // Destroy the session on the server side
914    if let Some(ref auth) = state.auth {
915        if let Some(cookie_header) = headers.get(header::COOKIE).and_then(|v| v.to_str().ok()) {
916            if let Some(session_id) = extract_token_from_cookies(cookie_header, cookie_name) {
917                if !session_id.starts_with("dev:") {
918                    let _ = auth.session_manager.destroy_session(&session_id);
919                }
920            }
921        }
922    }
923
924    let cookie = Cookie::build((cookie_name.to_string(), ""))
925        .path("/")
926        .http_only(true)
927        .same_site(SameSite::Lax)
928        .max_age(TimeDuration::seconds(-1))
929        .build();
930
931    Response::builder()
932        .status(StatusCode::FOUND)
933        .header(header::LOCATION, "/")
934        .header(header::SET_COOKIE, cookie.to_string())
935        .body(Body::empty())
936        .unwrap()
937}
938
939// ---------------------------------------------------------------------------
940// MCP auth routes: /auth/mcp/login, /callback, /status, /logout (C2)
941// ---------------------------------------------------------------------------
942
943/// Query parameters for MCP login initiation.
944#[derive(Debug, Deserialize)]
945pub struct McpLoginQuery {
946    pub cred: String,
947}
948
949/// Query parameters for MCP OIDC callback.
950#[derive(Debug, Deserialize)]
951pub struct McpCallbackQuery {
952    pub code: Option<String>,
953    pub state: Option<String>,
954    pub error: Option<String>,
955    pub error_description: Option<String>,
956}
957
958/// GET /auth/mcp/login?cred=<name> — initiate per-server OIDC PKCE login.
959///
960/// Reads the credential config from the session's config_toml, verifies it's
961/// `type = "web-interactive"`, then redirects to the OIDC provider.
962async fn mcp_login_handler(
963    State(state): State<crate::ServerState>,
964    Query(query): Query<McpLoginQuery>,
965    headers: axum::http::HeaderMap,
966) -> Result<Response, AuthError> {
967    // Require authentication — user must be logged into trustee-web
968    let (_cookie, _user_key) = crate::auth::check_auth(&state.auth, &headers)
969        .await
970        .map_err(|_| AuthError::AuthNotConfigured)?;
971
972    let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
973
974    // Parse MCP credential config from session's config_toml
975    let cred_config = load_mcp_credential(&state, &query.cred).await?;
976
977    let (issuer_url, client_id, client_secret, scope) = match &cred_config {
978        McpCredentialInfo::WebInteractive {
979            issuer_url,
980            client_id,
981            client_secret,
982            scope,
983        } => (issuer_url.clone(), client_id.clone(), client_secret.clone(), scope.clone()),
984        _ => {
985            return Ok(Redirect::temporary(&format!(
986                "/?mcp_error={}",
987                urlencoding::encode(&format!("Credential '{}' is not web-interactive type", query.cred))
988            ))
989            .into_response());
990        }
991    };
992
993    // Build PKCE pair using a separate PkceCookieManager for MCP
994    let oidc_client = OidcClient::new();
995    let verifier = OidcClient::generate_code_verifier();
996    let challenge = OidcClient::generate_code_challenge(&verifier);
997    let oauth_state = OidcClient::generate_state();
998
999    // Build OidcClientConfig for the MCP credential's OIDC client
1000    let mcp_redirect_uri = format!(
1001        "{}/auth/mcp/callback",
1002        auth.client_config.redirect_uri.trim_end_matches('/').trim_end_matches("/auth/callback")
1003    );
1004
1005    let mcp_client_config = OidcClientConfig {
1006        issuer_url: issuer_url.clone(),
1007        client_id: client_id.clone(),
1008        client_secret: client_secret.clone(),
1009        redirect_uri: mcp_redirect_uri.clone(),
1010        scope: scope.clone(),
1011        code_challenge_method: "S256".to_string(),
1012    };
1013
1014    // Build authorization URL
1015    let auth_url = oidc_client
1016        .build_authorization_url(&mcp_client_config, &oauth_state, Some(&challenge))
1017        .await
1018        .map_err(|e| AuthError::OidcError(e.to_string()))?;
1019
1020    // Store PKCE state + credential name in the in-memory map
1021    mcp_pkce().insert(oauth_state.clone(), verifier.clone(), query.cred.clone()).await;
1022
1023    tracing::info!(
1024        "Initiating MCP browser login for credential '{}' (issuer={})",
1025        query.cred, issuer_url
1026    );
1027
1028    Ok(Response::builder()
1029        .status(StatusCode::TEMPORARY_REDIRECT)
1030        .header(header::LOCATION, &auth_url)
1031        .body(Body::empty())
1032        .unwrap())
1033}
1034
1035/// GET /auth/mcp/callback — handle MCP OIDC callback, store tokens.
1036async fn mcp_callback_handler(
1037    State(state): State<crate::ServerState>,
1038    Query(query): Query<McpCallbackQuery>,
1039    headers: axum::http::HeaderMap,
1040) -> Result<Response, AuthError> {
1041    let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
1042
1043    // Check for errors from IdP
1044    if let Some(error) = query.error {
1045        let desc = query.error_description.unwrap_or_default();
1046        tracing::error!("MCP OIDC error: {} - {}", error, desc);
1047        return Ok(Redirect::temporary(&format!(
1048            "/?mcp_error={}&error_description={}",
1049            urlencoding::encode(&error),
1050            urlencoding::encode(&desc)
1051        ))
1052        .into_response());
1053    }
1054
1055    let code = query.code.ok_or(AuthError::MissingCode)?;
1056    let oauth_state = query.state.ok_or(AuthError::MissingState)?;
1057
1058    // Look up PKCE verifier + credential name from in-memory store
1059    let pkce_data = mcp_pkce().take(&oauth_state).await
1060        .ok_or(AuthError::InvalidState)?;
1061
1062    let verifier = pkce_data.verifier;
1063    let cred_name = &pkce_data.cred_name;
1064
1065    // Parse the MCP credential config to get OIDC settings for token exchange
1066    let cred_config = load_mcp_credential(&state, cred_name).await?;
1067
1068    let (issuer_url, client_id, client_secret, scope) = match &cred_config {
1069        McpCredentialInfo::WebInteractive {
1070            issuer_url,
1071            client_id,
1072            client_secret,
1073            scope,
1074        } => (issuer_url.clone(), client_id.clone(), client_secret.clone(), scope.clone()),
1075        _ => {
1076            return Ok(Redirect::temporary(&format!(
1077                "/?mcp_error={}",
1078                urlencoding::encode("Credential is not web-interactive type")
1079            ))
1080            .into_response());
1081        }
1082    };
1083
1084    // Build redirect URI (must match what was used in login)
1085    let mcp_redirect_uri = format!(
1086        "{}/auth/mcp/callback",
1087        auth.client_config.redirect_uri.trim_end_matches('/').trim_end_matches("/auth/callback")
1088    );
1089
1090    let mcp_client_config = OidcClientConfig {
1091        issuer_url: issuer_url.clone(),
1092        client_id: client_id.clone(),
1093        client_secret: client_secret.clone(),
1094        redirect_uri: mcp_redirect_uri,
1095        scope: scope.clone(),
1096        code_challenge_method: "S256".to_string(),
1097    };
1098
1099    // Exchange code for tokens
1100    tracing::info!("Exchanging MCP authorization code for tokens (credential={})", cred_name);
1101    let oidc_client = OidcClient::new();
1102    let token_response = oidc_client
1103        .exchange_code_for_tokens(&mcp_client_config, &code, Some(&verifier))
1104        .await
1105        .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
1106
1107    // Compute expires_at
1108    let expires_at = {
1109        let now = std::time::SystemTime::now()
1110            .duration_since(std::time::UNIX_EPOCH)
1111            .unwrap_or_default()
1112            .as_secs();
1113        let expires_epoch = now + token_response.expires_in.unwrap_or(900);
1114        let days = expires_epoch / 86400;
1115        let rem = expires_epoch % 86400;
1116        let h = rem / 3600;
1117        let m = (rem % 3600) / 60;
1118        let s = rem % 60;
1119        let z = days as i64 + 719468;
1120        let era = if z >= 0 { z } else { z - 146096 } / 146097;
1121        let doe = (z - era * 146097) as u64;
1122        let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
1123        let y = yoe as i64 + era * 400;
1124        let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
1125        let mp = (5 * doy + 2) / 153;
1126        let d = doy - (153 * mp + 2) / 5 + 1;
1127        let mon = if mp < 10 { mp + 3 } else { mp - 9 };
1128        let yr = if mon <= 2 { y + 1 } else { y };
1129        format!("{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z", yr, mon, d, h, m, s)
1130    };
1131
1132    // Store via FileTokenStore (same as `trustee mcp auth`)
1133    use pep::{FileTokenStore, StoredToken, TokenStore};
1134
1135    let stored = StoredToken::new(
1136        &token_response.access_token,
1137        token_response.refresh_token.clone(),
1138        "Bearer",
1139        &expires_at,
1140        token_response.scope.clone(),
1141    );
1142
1143    let agent_name = state.config_toml.as_ref().and_then(|t| {
1144        toml::from_str::<toml::Value>(t).ok()
1145            .and_then(|v| v.get("agent").and_then(|a| a.get("name")).and_then(|n| n.as_str()).map(String::from))
1146    }).unwrap_or_else(|| "trustee".to_string());
1147    let token_store = FileTokenStore::new(&agent_name);
1148
1149    if let Err(e) = token_store.save(cred_name, &stored) {
1150        tracing::error!("Failed to store MCP token: {}", e);
1151        return Ok(Redirect::temporary(&format!(
1152            "/?mcp_error={}",
1153            urlencoding::encode(&format!("Failed to store token: {}", e))
1154        ))
1155        .into_response());
1156    }
1157
1158    tracing::info!(
1159        "MCP authentication successful for credential '{}' (expires {})",
1160        cred_name, expires_at
1161    );
1162
1163    Ok(Response::builder()
1164        .status(StatusCode::FOUND)
1165        .header(header::LOCATION, format!("/?mcp_connected={}", urlencoding::encode(cred_name)))
1166        .body(Body::empty())
1167        .unwrap())
1168}
1169
1170/// GET /auth/mcp/status — return connection status for all MCP credentials.
1171async fn mcp_status_handler(
1172    State(state): State<crate::ServerState>,
1173    headers: axum::http::HeaderMap,
1174) -> Response {
1175    use pep::{FileTokenStore, TokenStore};
1176
1177    // Require auth
1178    let (_cookie, user_key) = match crate::auth::check_auth(&state.auth, &headers).await {
1179        Ok(result) => result,
1180        Err(code) => return (code, Json(serde_json::json!({"error": "Unauthorized"}))).into_response(),
1181    };
1182
1183    // Parse MCP config from session
1184    let config_toml = {
1185        let (_sid, session_arc, _, _) = state.ensure_active_session(&user_key).await;
1186        let session = session_arc.lock().await;
1187        match &session.config_toml {
1188            Some(t) => t.clone(),
1189            None => return (StatusCode::INTERNAL_SERVER_ERROR, "Config not loaded").into_response(),
1190        }
1191    };
1192
1193    let mcp_config: toml::Value = match toml::from_str(&config_toml) {
1194        Ok(v) => v,
1195        Err(_) => return Json(serde_json::json!([])).into_response(),
1196    };
1197
1198    let agent_name = {
1199        let (_sid, session_arc, _, _) = state.ensure_active_session(&user_key).await;
1200        let session = session_arc.lock().await;
1201        session.agent_name.clone()
1202    };
1203    let token_store = FileTokenStore::new(&agent_name);
1204
1205    // Build server → credential mapping
1206    let servers = mcp_config
1207        .get("mcp")
1208        .and_then(|m| m.get("servers"))
1209        .and_then(|s| s.as_array());
1210    let credentials = mcp_config
1211        .get("mcp")
1212        .and_then(|m| m.get("credentials"))
1213        .and_then(|c| c.as_table());
1214
1215    let mut cred_servers: std::collections::HashMap<String, Vec<String>> = std::collections::HashMap::new();
1216    if let Some(servers) = servers {
1217        for server in servers {
1218            let name = server.get("name").and_then(|n| n.as_str()).unwrap_or("");
1219            let cred_ref = server.get("credentials").and_then(|c| c.as_str()).unwrap_or("");
1220            if !cred_ref.is_empty() {
1221                cred_servers
1222                    .entry(cred_ref.to_string())
1223                    .or_default()
1224                    .push(name.to_string());
1225            }
1226        }
1227    }
1228
1229    let mut result = Vec::new();
1230
1231    if let Some(creds) = credentials {
1232        for (cred_name, cred_config) in creds {
1233            let cred_type = cred_config.get("type").and_then(|t| t.as_str()).unwrap_or("unknown");
1234            let servers_using = cred_servers.get(cred_name).cloned().unwrap_or_default();
1235
1236            if cred_type == "web-session" {
1237                // Session credentials are always "connected" if auth is enabled
1238                let connected = state.auth.is_some();
1239                result.push(serde_json::json!({
1240                    "credential": cred_name,
1241                    "type": cred_type,
1242                    "connected": connected,
1243                    "servers": servers_using,
1244                }));
1245            } else if cred_type == "service-account" {
1246                // Long-lived service token, exchanged lazily (RFC 8693) by the
1247                // agent at runtime. "Connected" = the service_token resolved to
1248                // a non-empty value in this (already ${VAR}-substituted) config.
1249                let token = cred_config
1250                    .get("service_token")
1251                    .and_then(|t| t.as_str())
1252                    .unwrap_or("");
1253                result.push(serde_json::json!({
1254                    "credential": cred_name,
1255                    "type": cred_type,
1256                    "connected": !token.is_empty(),
1257                    "servers": servers_using,
1258                }));
1259            } else if cred_type == "static" {
1260                // Static token — connected when it resolved non-empty.
1261                let token = cred_config
1262                    .get("token")
1263                    .and_then(|t| t.as_str())
1264                    .unwrap_or("");
1265                result.push(serde_json::json!({
1266                    "credential": cred_name,
1267                    "type": cred_type,
1268                    "connected": !token.is_empty(),
1269                    "servers": servers_using,
1270                }));
1271            } else if cred_type == "web-interactive" || cred_type == "interactive" {
1272                // Check token store
1273                let status = match token_store.load(cred_name) {
1274                    Ok(Some(token)) => {
1275                        let expired = token.is_expired();
1276                        serde_json::json!({
1277                            "credential": cred_name,
1278                            "type": cred_type,
1279                            "connected": !expired,
1280                            "expires_at": token.expires_at,
1281                            "servers": servers_using,
1282                        })
1283                    }
1284                    _ => serde_json::json!({
1285                        "credential": cred_name,
1286                        "type": cred_type,
1287                        "connected": false,
1288                        "servers": servers_using,
1289                    }),
1290                };
1291                result.push(status);
1292            }
1293        }
1294    }
1295
1296    Json(serde_json::Value::Array(result)).into_response()
1297}
1298
1299/// POST /auth/mcp/logout?cred=<name> — remove stored MCP tokens.
1300async fn mcp_logout_handler(
1301    State(state): State<crate::ServerState>,
1302    Query(query): Query<McpLoginQuery>,
1303    headers: axum::http::HeaderMap,
1304) -> Response {
1305    use pep::{FileTokenStore, TokenStore};
1306
1307    // Require auth
1308    let (_cookie, user_key) = match crate::auth::check_auth(&state.auth, &headers).await {
1309        Ok(result) => result,
1310        Err(code) => return (code, Json(serde_json::json!({"error": "Unauthorized"}))).into_response(),
1311    };
1312
1313    let agent_name = {
1314        let (_sid, session_arc, _, _) = state.ensure_active_session(&user_key).await;
1315        let session = session_arc.lock().await;
1316        session.agent_name.clone()
1317    };
1318    let token_store = FileTokenStore::new(&agent_name);
1319
1320    match token_store.delete(&query.cred) {
1321        Ok(()) => {
1322            tracing::info!("Removed MCP credentials for '{}'", query.cred);
1323            Json(serde_json::json!({"success": true})).into_response()
1324        }
1325        Err(e) => {
1326            tracing::error!("Failed to remove MCP credentials: {}", e);
1327            (
1328                StatusCode::INTERNAL_SERVER_ERROR,
1329                Json(serde_json::json!({"error": e.to_string()})),
1330            )
1331                .into_response()
1332        }
1333    }
1334}
1335
1336// ---------------------------------------------------------------------------
1337// MCP auth helpers
1338// ---------------------------------------------------------------------------
1339
1340/// In-memory store for MCP PKCE state (state token → verifier + credential name).
1341/// Entries expire after 10 minutes. Not persisted across restarts.
1342struct McpPkceStore {
1343    entries: tokio::sync::Mutex<std::collections::HashMap<String, McpPkceEntry>>,
1344}
1345
1346struct McpPkceEntry {
1347    verifier: String,
1348    cred_name: String,
1349    created_at: std::time::Instant,
1350}
1351
1352impl McpPkceStore {
1353    fn new() -> Self {
1354        Self {
1355            entries: tokio::sync::Mutex::new(std::collections::HashMap::new()),
1356        }
1357    }
1358
1359    /// Insert a PKCE entry. Cleans up entries older than 10 minutes.
1360    async fn insert(&self, state: String, verifier: String, cred_name: String) {
1361        let mut map = self.entries.lock().await;
1362        // Cleanup expired entries (older than 10 min)
1363        let cutoff = std::time::Instant::now() - std::time::Duration::from_secs(600);
1364        map.retain(|_, v| v.created_at > cutoff);
1365        map.insert(state, McpPkceEntry {
1366            verifier,
1367            cred_name,
1368            created_at: std::time::Instant::now(),
1369        });
1370    }
1371
1372    /// Take and remove a PKCE entry (single-use).
1373    async fn take(&self, state: &str) -> Option<McpPkceEntry> {
1374        let mut map = self.entries.lock().await;
1375        map.remove(state)
1376    }
1377}
1378
1379/// Global singleton PKCE store for MCP browser logins.
1380static MCP_PKCE: std::sync::OnceLock<McpPkceStore> = std::sync::OnceLock::new();
1381
1382/// Get or initialize the global MCP PKCE store.
1383fn mcp_pkce() -> &'static McpPkceStore {
1384    MCP_PKCE.get_or_init(McpPkceStore::new)
1385}
1386
1387/// Simplified MCP credential info (parsed from TOML).
1388enum McpCredentialInfo {
1389    WebInteractive {
1390        issuer_url: String,
1391        client_id: String,
1392        client_secret: Option<String>,
1393        scope: String,
1394    },
1395    Other(String),
1396}
1397
1398/// Load a specific MCP credential from the session's config_toml.
1399async fn load_mcp_credential(
1400    state: &crate::ServerState,
1401    cred_name: &str,
1402) -> Result<McpCredentialInfo, AuthError> {
1403    let config_toml = state
1404        .config_toml
1405        .clone()
1406        .ok_or(AuthError::AuthNotConfigured)?;
1407
1408    let config: toml::Value = toml::from_str(&config_toml)
1409        .map_err(|e| AuthError::OidcError(format!("Config parse error: {}", e)))?;
1410
1411    let cred = config
1412        .get("mcp")
1413        .and_then(|m| m.get("credentials"))
1414        .and_then(|c| c.as_table())
1415        .and_then(|c| c.get(cred_name))
1416        .ok_or_else(|| AuthError::OidcError(format!("Credential '{}' not found", cred_name)))?;
1417
1418    let cred_type = cred.get("type").and_then(|t| t.as_str()).unwrap_or("unknown");
1419
1420    match cred_type {
1421        "web-interactive" => {
1422            let issuer_url = cred
1423                .get("issuer_url")
1424                .and_then(|v| v.as_str())
1425                .ok_or_else(|| AuthError::OidcError("Missing issuer_url".into()))?
1426                .to_string();
1427            let client_id = cred
1428                .get("client_id")
1429                .and_then(|v| v.as_str())
1430                .ok_or_else(|| AuthError::OidcError("Missing client_id".into()))?
1431                .to_string();
1432            let client_secret = cred
1433                .get("client_secret")
1434                .and_then(|v| v.as_str())
1435                .map(String::from);
1436            let scope = cred
1437                .get("scope")
1438                .and_then(|v| v.as_str())
1439                .unwrap_or("openid profile email")
1440                .to_string();
1441
1442            Ok(McpCredentialInfo::WebInteractive {
1443                issuer_url,
1444                client_id,
1445                client_secret,
1446                scope,
1447            })
1448        }
1449        other => Ok(McpCredentialInfo::Other(other.to_string())),
1450    }
1451}
1452
1453// ---------------------------------------------------------------------------
1454// Helpers
1455// ---------------------------------------------------------------------------
1456
1457/// Create an HttpOnly auth cookie.
1458fn create_auth_cookie(name: &str, value: &str, max_age: StdDuration, secure: bool) -> Cookie<'static> {
1459    Cookie::build((name.to_string(), value.to_string()))
1460        .path("/")
1461        .http_only(true)
1462        .same_site(SameSite::Lax)
1463        .secure(secure)
1464        .max_age(TimeDuration::seconds(max_age.as_secs() as i64))
1465        .build()
1466}
1467
1468// ---------------------------------------------------------------------------
1469// Error handling
1470// ---------------------------------------------------------------------------
1471
1472/// Authentication errors.
1473#[derive(Debug)]
1474pub enum AuthError {
1475    MissingCode,
1476    MissingState,
1477    InvalidState,
1478    OidcError(String),
1479    TokenExchangeFailed(String),
1480    AuthNotConfigured,
1481}
1482
1483impl IntoResponse for AuthError {
1484    fn into_response(self) -> Response {
1485        let (_status, msg) = match self {
1486            AuthError::MissingCode => (StatusCode::BAD_REQUEST, "Missing authorization code"),
1487            AuthError::MissingState => (StatusCode::BAD_REQUEST, "Missing state parameter"),
1488            AuthError::InvalidState => (StatusCode::BAD_REQUEST, "Invalid or expired state"),
1489            AuthError::OidcError(_) => (StatusCode::SERVICE_UNAVAILABLE, "Authentication service error"),
1490            AuthError::TokenExchangeFailed(_) => (StatusCode::BAD_REQUEST, "Token exchange failed"),
1491            AuthError::AuthNotConfigured => (StatusCode::NOT_IMPLEMENTED, "Authentication not configured"),
1492        };
1493        Redirect::temporary(&format!("/?error={}", urlencoding::encode(msg))).into_response()
1494    }
1495}