OpenPXE/crates/core/src/saml/response.rs

//! SAMLResponse consumption: signature verification + SP-side validation.
//!
//! [`consume`] is intentionally **stateless** — it verifies the XML signature
//! against the IdP's pinned certificate(s) and enforces every check that can
//! be made from the response alone (Status, Destination, Issuer, Audience,
//! time bounds, signature scope). It then returns the `assertion_id` and
//! `in_response_to` so the HTTP layer can perform the *stateful* checks it
//! owns: replay rejection, correlating the request we issued, and gating
//! IdP-initiated login.

use roxmltree::{Document, Node};
use time::format_description::well_known::Rfc3339;
use time::{Duration, OffsetDateTime};

use super::metadata::IdpMetadata;
use super::{SamlError, SpParams};

const STATUS_SUCCESS: &str = "urn:oasis:names:tc:SAML:2.0:status:Success";

/// The verified subject of a SAML assertion. OpenPXE is single-tier, so this
/// is all an operator session needs.
#[derive(Debug, Clone)]
pub struct VerifiedPrincipal {
    /// The `<NameID>` value (an email, per our requested NameID format).
    pub name_id: String,
    /// Email used as the session identity. Equals `name_id` for the
    /// emailAddress NameID format.
    pub email: String,
    /// Human-readable display name, if the IdP sent one as an attribute.
    pub display_name: Option<String>,
}

/// Result of a successful [`consume`]. The IDs/expiry feed the HTTP layer's
/// stateful checks.
#[derive(Debug, Clone)]
pub struct VerifiedResponse {
    pub principal: VerifiedPrincipal,
    /// `InResponseTo` from the response, if present. `None` = unsolicited
    /// (IdP-initiated) — the HTTP layer only accepts that when the operator
    /// enabled it.
    pub in_response_to: Option<String>,
    /// The assertion's `ID` — used by the caller as the replay-guard key.
    pub assertion_id: String,
    /// The assertion's expiry (`Conditions/@NotOnOrAfter`) — the replay
    /// guard can drop the consumed ID after this instant.
    pub assertion_expiry: OffsetDateTime,
    /// `AuthnStatement/@SessionIndex`, if present (useful for future SLO).
    pub session_index: Option<String>,
}

/// Verify and validate a decoded `SAMLResponse` XML document.
pub fn consume(
    xml: &str,
    sp: &SpParams,
    idp: &IdpMetadata,
    now: OffsetDateTime,
    clock_skew: Duration,
) -> Result<VerifiedResponse, SamlError> {
    // 1. Cryptographically verify the signature against the pinned IdP cert(s).
    //    `trusted_keys_only` ignores any cert embedded in the document's
    //    KeyInfo, so an attacker can't substitute their own key.
    let verified_uris = verify_signature(xml, &idp.signing_certs_der)?;

    // 2. Parse for semantic validation.
    let doc = Document::parse(xml).map_err(|e| SamlError::Xml(e.to_string()))?;
    let root = doc.root_element();
    if root.tag_name().name() != "Response" {
        return Err(SamlError::MissingElement("Response".into()));
    }
    let response_id = root.attribute("ID").map(str::to_owned);
    let in_response_to = root.attribute("InResponseTo").map(str::to_owned);

    // 3. Status must be Success.
    let status_value = root
        .descendants()
        .find(|n| n.is_element() && n.tag_name().name() == "StatusCode")
        .and_then(|sc| sc.attribute("Value"))
        .unwrap_or("");
    if status_value != STATUS_SUCCESS && !status_value.ends_with(":Success") {
        return Err(SamlError::Status(status_value.to_owned()));
    }

    // 4. Destination (if the IdP set one) must be our ACS.
    if let Some(dest) = root.attribute("Destination") {
        if !urls_equal(dest, &sp.acs_url) {
            return Err(SamlError::DestinationMismatch);
        }
    }

    // 5. Exactly one (unencrypted) Assertion.
    if root
        .descendants()
        .any(|n| n.is_element() && n.tag_name().name() == "EncryptedAssertion")
    {
        return Err(SamlError::EncryptedAssertionUnsupported);
    }
    let assertions: Vec<Node<'_, '_>> = root
        .children()
        .filter(|c| c.is_element() && c.tag_name().name() == "Assertion")
        .collect();
    if assertions.len() != 1 {
        return Err(SamlError::AssertionCount(assertions.len()));
    }
    let assertion = assertions[0];
    let assertion_id = assertion
        .attribute("ID")
        .map(str::to_owned)
        .ok_or_else(|| SamlError::MissingElement("Assertion/@ID".into()))?;

    // 6. The signature must actually cover the assertion we're about to trust:
    //    either the assertion itself, the enclosing response, or the whole
    //    document. (bergshamra's strict_verification already constrains where
    //    the signed element may sit; this ties it to *our* assertion.)
    let covers_assertion = verified_uris.iter().any(|u| {
        u.is_empty()
            || u == &format!("#{assertion_id}")
            || response_id
                .as_ref()
                .is_some_and(|rid| u == &format!("#{rid}"))
    });
    if !covers_assertion {
        return Err(SamlError::SignatureScope);
    }

    // 7. Issuer must be the configured IdP.
    let issuer = first_child(assertion, "Issuer")
        .map(text_of)
        .unwrap_or_default();
    if !idp.entity_id.is_empty() && issuer != idp.entity_id {
        return Err(SamlError::IssuerMismatch);
    }

    // 8. Subject → NameID + SubjectConfirmationData time/recipient checks.
    let subject = first_child(assertion, "Subject")
        .ok_or_else(|| SamlError::MissingElement("Subject".into()))?;
    let name_id = first_child(subject, "NameID")
        .map(text_of)
        .filter(|s| !s.is_empty())
        .ok_or_else(|| SamlError::MissingElement("NameID".into()))?;
    if let Some(scd) = subject
        .descendants()
        .find(|n| n.is_element() && n.tag_name().name() == "SubjectConfirmationData")
    {
        if let Some(recipient) = scd.attribute("Recipient") {
            if !urls_equal(recipient, &sp.acs_url) {
                return Err(SamlError::DestinationMismatch);
            }
        }
        if let Some(noa) = scd.attribute("NotOnOrAfter") {
            let noa = parse_instant(noa)?;
            if now >= noa + clock_skew {
                return Err(SamlError::TimeBounds);
            }
        }
    }

    // 9. Conditions: time window + audience.
    let conditions = first_child(assertion, "Conditions");
    if let Some(cond) = conditions {
        if let Some(nb) = cond.attribute("NotBefore") {
            let nb = parse_instant(nb)?;
            if now < nb - clock_skew {
                return Err(SamlError::TimeBounds);
            }
        }
    }
    let assertion_expiry = conditions
        .and_then(|c| c.attribute("NotOnOrAfter"))
        .map(parse_instant)
        .transpose()?
        .ok_or_else(|| SamlError::MissingElement("Conditions/@NotOnOrAfter".into()))?;
    if now >= assertion_expiry + clock_skew {
        return Err(SamlError::TimeBounds);
    }

    let audience_ok = conditions.is_some_and(|c| {
        c.descendants()
            .filter(|n| n.is_element() && n.tag_name().name() == "Audience")
            .any(|a| text_of(a) == sp.entity_id)
    });
    if !audience_ok {
        return Err(SamlError::AudienceMismatch);
    }

    // 10. Optional: SessionIndex + display-name attribute.
    let session_index = assertion
        .descendants()
        .find(|n| n.is_element() && n.tag_name().name() == "AuthnStatement")
        .and_then(|a| a.attribute("SessionIndex"))
        .map(str::to_owned);

    let display_name = extract_display_name(assertion);

    Ok(VerifiedResponse {
        principal: VerifiedPrincipal {
            email: name_id.clone(),
            name_id,
            display_name,
        },
        in_response_to,
        assertion_id,
        assertion_expiry,
        session_index,
    })
}

/// Verify the document's XML-DSig against each pinned IdP cert in turn
/// (handles key rotation), returning the verified `<Reference>` URIs.
fn verify_signature(xml: &str, certs_der: &[Vec<u8>]) -> Result<Vec<String>, SamlError> {
    let mut last_err = String::from("no signing certificate matched");
    for der in certs_der {
        let key = match bergshamra::keys::loader::load_x509_cert_der(der) {
            Ok(k) => k,
            Err(e) => {
                last_err = e.to_string();
                continue;
            }
        };
        let mut km = bergshamra::keys::KeysManager::new();
        km.add_key(key);
        // trusted_keys_only: only ever trust the pinned IdP key, never an
        // inline KeyInfo cert. strict_verification: XSW positional defense.
        let ctx = bergshamra::DsigContext::new(km)
            .with_trusted_keys_only(true)
            .with_strict_verification(true);
        match bergshamra::verify(&ctx, xml) {
            Ok(bergshamra::VerifyResult::Valid { references, .. }) => {
                return Ok(references.into_iter().map(|r| r.uri).collect());
            }
            Ok(bergshamra::VerifyResult::Invalid { reason }) => last_err = reason,
            Err(e) => last_err = e.to_string(),
        }
    }
    Err(SamlError::Signature(last_err))
}

/// Pull a display name from the assertion's attribute statement, trying the
/// common attribute names IdPs use (FleetDM checks the same set).
fn extract_display_name(assertion: Node<'_, '_>) -> Option<String> {
    const WANTED: &[&str] = &[
        "name",
        "displayname",
        "cn",
        "urn:oid:2.5.4.3",
        "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/name",
    ];
    for attr in assertion
        .descendants()
        .filter(|n| n.is_element() && n.tag_name().name() == "Attribute")
    {
        let key = attr
            .attribute("Name")
            .or_else(|| attr.attribute("FriendlyName"))
            .unwrap_or("")
            .to_ascii_lowercase();
        if WANTED.contains(&key.as_str()) {
            if let Some(val) = attr
                .descendants()
                .find(|n| n.is_element() && n.tag_name().name() == "AttributeValue")
            {
                let v = text_of(val);
                if !v.is_empty() {
                    return Some(v);
                }
            }
        }
    }
    None
}

fn first_child<'a, 'i>(n: Node<'a, 'i>, local: &str) -> Option<Node<'a, 'i>> {
    n.children()
        .find(|c| c.is_element() && c.tag_name().name() == local)
}

fn text_of(n: Node<'_, '_>) -> String {
    n.children()
        .filter(Node::is_text)
        .filter_map(|c| c.text())
        .collect::<String>()
        .trim()
        .to_owned()
}

/// Parse an `xs:dateTime` (e.g. `2026-05-31T12:00:00.000Z`).
fn parse_instant(s: &str) -> Result<OffsetDateTime, SamlError> {
    OffsetDateTime::parse(s.trim(), &Rfc3339).map_err(|e| SamlError::Timestamp(format!("{s}: {e}")))
}

/// Compare two URLs for SAML endpoint-matching purposes: exact, or differing
/// only by a single trailing slash.
fn urls_equal(a: &str, b: &str) -> bool {
    a == b || a.trim_end_matches('/') == b.trim_end_matches('/')
}