OpenPXE/crates/http-api/src/saml_routes.rs

//! SAML 2.0 Service Provider HTTP endpoints (v0.5.1).
//!
//! * `GET  /api/sso/login`    — SP-initiated: build an AuthnRequest, record its
//!   ID, and 302 the browser to the IdP.
//! * `POST /api/sso/acs`      — Assertion Consumer Service: verify + validate
//!   the IdP's `SAMLResponse`, perform the stateful checks (InResponseTo
//!   correlation, IdP-initiated gating, assertion replay), mint an operator
//!   session, and 302 to the dashboard. (Mirrors FleetDM's `/sso/callback`.)
//! * `GET  /api/sso/metadata` — serve our SP metadata XML for IdP import.
//!
//! Stateless crypto + semantic validation live in `openpxe_core::saml`; this
//! module owns only the HTTP glue and the in-memory state the SP needs.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration as StdDuration, Instant};

use axum::{
    body::Body,
    extract::{Form, Query, State},
    http::{header, StatusCode},
    response::{IntoResponse, Response},
};
use base64::Engine;
use parking_lot::Mutex;
use serde::Deserialize;
use time::{Duration, OffsetDateTime};

use openpxe_core::saml::{self, metadata::IdpMetadata, SamlError, SpParams};
use openpxe_core::SsoConfig;

use crate::auth;
use crate::state::AppState;

/// Outstanding AuthnRequest IDs live at most this long before a matching
/// response is considered stale (covers a slow human at the IdP login form).
const REQUEST_TTL: StdDuration = StdDuration::from_mins(10);
/// How long we fetch-cache IdP metadata loaded from a URL.
const METADATA_FETCH_TIMEOUT: StdDuration = StdDuration::from_secs(10);

/// In-memory SAML runtime state. Cheap to clone (Arc-shared).
#[derive(Clone, Default)]
pub struct SamlRuntime {
    /// request_id → issued_at. Correlates a response's `InResponseTo` to a
    /// request *we* actually sent (replay / CSRF defense for SP-initiated).
    outstanding: Arc<Mutex<HashMap<String, Instant>>>,
    /// assertion_id → expiry. A consumed assertion may not be replayed.
    consumed: Arc<Mutex<HashMap<String, Instant>>>,
    /// Cache of IdP metadata fetched from a URL: (url, parsed).
    metadata_cache: Arc<Mutex<Option<(String, IdpMetadata)>>>,
}

impl SamlRuntime {
    /// Record an AuthnRequest we just sent.
    pub fn register_request(&self, id: &str) {
        let mut g = self.outstanding.lock();
        prune(&mut g);
        g.insert(id.to_owned(), Instant::now());
    }

    /// Consume an outstanding request ID, returning `true` if it was present
    /// and still fresh. A miss means the response doesn't correlate to any
    /// live request we issued.
    pub fn take_request(&self, id: &str) -> bool {
        let mut g = self.outstanding.lock();
        prune(&mut g);
        g.remove(id).is_some()
    }

    /// Record a consumed assertion. Returns `false` if it was already
    /// consumed (a replay) — in which case the caller must reject.
    pub fn record_assertion(&self, id: &str, expiry: OffsetDateTime) -> bool {
        let mut g = self.consumed.lock();
        prune(&mut g);
        if g.contains_key(id) {
            return false;
        }
        let ttl = (expiry - OffsetDateTime::now_utc())
            .max(Duration::ZERO)
            .unsigned_abs();
        g.insert(id.to_owned(), Instant::now() + ttl);
        true
    }

    fn cached_metadata(&self, url: &str) -> Option<IdpMetadata> {
        let g = self.metadata_cache.lock();
        match &*g {
            Some((cached_url, md)) if cached_url == url => Some(md.clone()),
            _ => None,
        }
    }

    fn cache_metadata(&self, url: String, md: IdpMetadata) {
        *self.metadata_cache.lock() = Some((url, md));
    }
}

/// Drop expired entries so neither map grows unbounded.
fn prune(map: &mut HashMap<String, Instant>) {
    let now = Instant::now();
    // For the request map this over-prunes (entries store issued_at, not
    // expiry), so cap by REQUEST_TTL; the consumed map stores absolute
    // expiry instants. Using saturating logic keeps both correct: request
    // entries older than REQUEST_TTL go, consumed entries past expiry go.
    map.retain(|_, &mut t| now.saturating_duration_since(t) < REQUEST_TTL || t > now);
}

// ─── GET /api/sso/login ───────────────────────────────────────────────────

#[derive(Debug, Deserialize)]
pub struct LoginQuery {
    /// Optional local path to return to after login (becomes RelayState).
    #[serde(default)]
    pub next: Option<String>,
}

pub async fn sso_login(State(state): State<AppState>, Query(q): Query<LoginQuery>) -> Response {
    let cfg = state.sso.snapshot();
    if !cfg.is_usable() {
        return redirect("/?sso_error=unavailable");
    }
    let idp = match resolve_idp_metadata(&state, &cfg).await {
        Ok(m) => m,
        Err(e) => {
            tracing::warn!(target: "openpxe::saml", "sso_login: metadata unavailable: {e}");
            return redirect("/?sso_error=metadata");
        }
    };
    let Some(dest) = idp.sso_destination().map(str::to_owned) else {
        tracing::warn!(target: "openpxe::saml", "sso_login: IdP metadata has no SSO endpoint");
        return redirect("/?sso_error=metadata");
    };
    let sp = sp_params(&state, &cfg);
    let relay = safe_local_path(q.next.as_deref());
    match saml::authn_request::build(&sp, &dest, Some(&relay)) {
        Ok(req) => {
            state.saml.register_request(&req.id);
            redirect(&req.location)
        }
        Err(e) => {
            tracing::warn!(target: "openpxe::saml", "sso_login: build AuthnRequest failed: {e}");
            redirect("/?sso_error=request")
        }
    }
}

// ─── POST /api/sso/acs ──────────────────────────────────────────────────────

#[derive(Debug, Deserialize)]
pub struct AcsForm {
    #[serde(rename = "SAMLResponse")]
    pub saml_response: String,
    #[serde(rename = "RelayState", default)]
    pub relay_state: Option<String>,
}

pub async fn sso_acs(State(state): State<AppState>, Form(form): Form<AcsForm>) -> Response {
    let cfg = state.sso.snapshot();
    if !cfg.is_usable() {
        return redirect("/?sso_error=unavailable");
    }
    let xml = match base64::engine::general_purpose::STANDARD.decode(form.saml_response.as_bytes())
    {
        Ok(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
        Err(e) => {
            tracing::warn!(target: "openpxe::saml", "acs: base64 decode failed: {e}");
            return redirect("/?sso_error=1");
        }
    };
    let idp = match resolve_idp_metadata(&state, &cfg).await {
        Ok(m) => m,
        Err(e) => {
            tracing::warn!(target: "openpxe::saml", "acs: metadata unavailable: {e}");
            return redirect("/?sso_error=metadata");
        }
    };
    let sp = sp_params(&state, &cfg);

    // Signature verification + semantic checks are CPU-bound — keep them off
    // the async executor.
    let now = OffsetDateTime::now_utc();
    let skew = Duration::seconds(saml::DEFAULT_CLOCK_SKEW_SECS);
    let verify = {
        let xml = xml.clone();
        let sp = sp.clone();
        tokio::task::spawn_blocking(move || saml::response::consume(&xml, &sp, &idp, now, skew))
            .await
    };
    let verified = match verify {
        Ok(Ok(v)) => v,
        Ok(Err(e)) => {
            // Never leak which specific check failed to the browser.
            tracing::warn!(target: "openpxe::saml", "acs: response rejected: {e}");
            return redirect("/?sso_error=1");
        }
        Err(join) => {
            tracing::error!(target: "openpxe::saml", "acs: verify task panicked: {join}");
            return redirect("/?sso_error=1");
        }
    };

    // Stateful checks the core deliberately left to us.
    match &verified.in_response_to {
        Some(id) => {
            if !state.saml.take_request(id) {
                tracing::warn!(target: "openpxe::saml", "acs: InResponseTo matches no live request");
                return redirect("/?sso_error=1");
            }
        }
        None => {
            if !cfg.allow_idp_initiated {
                tracing::warn!(target: "openpxe::saml", "acs: IdP-initiated login is disabled");
                return redirect("/?sso_error=idp_initiated");
            }
        }
    }
    if !state
        .saml
        .record_assertion(&verified.assertion_id, verified.assertion_expiry)
    {
        tracing::warn!(target: "openpxe::saml", "acs: assertion replay rejected");
        return redirect("/?sso_error=1");
    }

    // Success → mint an operator session keyed to the verified email.
    let session = state.sessions.create(&verified.principal.email);
    tracing::info!(
        target: "openpxe::saml",
        email = %verified.principal.email,
        idp_initiated = verified.in_response_to.is_none(),
        "SAML SSO sign-in"
    );
    // safe_local_path already maps None / unsafe values to "/".
    let relay = safe_local_path(form.relay_state.as_deref());
    redirect_with_session(&relay, &session)
}

// ─── GET /api/sso/metadata ──────────────────────────────────────────────────

pub async fn sso_metadata(State(state): State<AppState>) -> Response {
    let cfg = state.sso.snapshot();
    let sp = sp_params(&state, &cfg);
    let xml = saml::metadata::build_sp_metadata(&sp);
    (
        StatusCode::OK,
        [(header::CONTENT_TYPE, "application/samlmetadata+xml")],
        xml,
    )
        .into_response()
}

// ─── helpers ────────────────────────────────────────────────────────────────

/// Derive runtime SP parameters from config + the advertised public base URL.
fn sp_params(state: &AppState, cfg: &SsoConfig) -> SpParams {
    let base = state.public_base_url.trim_end_matches('/');
    let entity_id = if cfg.entity_id.trim().is_empty() {
        base.to_owned()
    } else {
        cfg.entity_id.trim().to_owned()
    };
    SpParams {
        entity_id,
        acs_url: format!("{base}/api/sso/acs"),
    }
}

/// Resolve the IdP metadata: prefer the metadata URL (fetched + cached) per
/// the "URL wins" rule, else parse the pasted XML.
async fn resolve_idp_metadata(state: &AppState, cfg: &SsoConfig) -> Result<IdpMetadata, SamlError> {
    let url = cfg.metadata_url.trim();
    if !url.is_empty() {
        if let Some(md) = state.saml.cached_metadata(url) {
            return Ok(md);
        }
        let body = fetch_metadata(url).await?;
        let md = IdpMetadata::parse(&body)?;
        state.saml.cache_metadata(url.to_owned(), md.clone());
        return Ok(md);
    }
    if !cfg.metadata.trim().is_empty() {
        return IdpMetadata::parse(&cfg.metadata);
    }
    Err(SamlError::Metadata("no metadata source configured".into()))
}

async fn fetch_metadata(url: &str) -> Result<String, SamlError> {
    let client = reqwest::Client::builder()
        .timeout(METADATA_FETCH_TIMEOUT)
        .build()
        .map_err(|e| SamlError::Metadata(format!("http client: {e}")))?;
    let resp = client
        .get(url)
        .send()
        .await
        .map_err(|e| SamlError::Metadata(format!("fetch {url}: {e}")))?;
    if !resp.status().is_success() {
        return Err(SamlError::Metadata(format!(
            "fetch {url}: HTTP {}",
            resp.status()
        )));
    }
    resp.text()
        .await
        .map_err(|e| SamlError::Metadata(format!("read {url}: {e}")))
}

/// Only permit a same-site path (single leading slash) as a redirect target —
/// blocks open-redirect / protocol-relative (`//evil.com`) abuse of RelayState.
fn safe_local_path(p: Option<&str>) -> String {
    match p {
        Some(p) if p.starts_with('/') && !p.starts_with("//") => p.to_owned(),
        _ => "/".to_owned(),
    }
}

fn redirect(location: &str) -> Response {
    Response::builder()
        .status(StatusCode::FOUND)
        .header(header::LOCATION, location)
        .body(Body::empty())
        .map_or_else(
            |_| StatusCode::INTERNAL_SERVER_ERROR.into_response(),
            IntoResponse::into_response,
        )
}

fn redirect_with_session(location: &str, session: &str) -> Response {
    Response::builder()
        .status(StatusCode::FOUND)
        .header(header::LOCATION, location)
        .header(header::SET_COOKIE, auth::session_cookie(session))
        .body(Body::empty())
        .map_or_else(
            |_| StatusCode::INTERNAL_SERVER_ERROR.into_response(),
            IntoResponse::into_response,
        )
}

#[cfg(test)]
mod tests {
    use super::*;
    use wiremock::matchers::method;
    use wiremock::{Mock, MockServer, ResponseTemplate};

    // v0.5.4: exercise the SAML metadata-URL fetch against a mock server —
    // previously this path did a real network GET and had no coverage.
    #[tokio::test]
    async fn fetch_metadata_returns_body_on_200() {
        let server = MockServer::start().await;
        let xml = "<EntityDescriptor>idp</EntityDescriptor>";
        Mock::given(method("GET"))
            .respond_with(ResponseTemplate::new(200).set_body_string(xml))
            .mount(&server)
            .await;
        let got = fetch_metadata(&server.uri()).await.expect("fetch ok");
        assert_eq!(got, xml);
    }

    #[tokio::test]
    async fn fetch_metadata_errors_on_non_2xx() {
        let server = MockServer::start().await;
        Mock::given(method("GET"))
            .respond_with(ResponseTemplate::new(503))
            .mount(&server)
            .await;
        let err = fetch_metadata(&server.uri()).await.unwrap_err();
        assert!(matches!(err, SamlError::Metadata(_)), "got {err:?}");
    }
}