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

//! AuthnRequest construction + HTTP-Redirect binding encoding.
//!
//! For SP-initiated login we build an `<AuthnRequest>`, then encode it for the
//! HTTP-Redirect binding: raw DEFLATE (RFC 1951) → base64 → percent-encode,
//! appended as the `SAMLRequest` query parameter. AuthnRequests are sent
//! unsigned in this release (the IdP must not require client signatures).

use std::io::Write as _;

use base64::Engine;
use flate2::write::DeflateEncoder;
use flate2::Compression;
use time::format_description::well_known::Rfc3339;
use time::OffsetDateTime;

use super::{SamlError, SpParams};
use crate::encoding::{pct_encode, xml_escape};

const NS_PROTOCOL: &str = "urn:oasis:names:tc:SAML:2.0:protocol";
const NS_ASSERTION: &str = "urn:oasis:names:tc:SAML:2.0:assertion";
const NAMEID_EMAIL: &str = "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress";
const BINDING_POST: &str = "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST";

/// A built AuthnRequest, ready to redirect the browser to the IdP.
#[derive(Debug, Clone)]
pub struct AuthnRequest {
    /// The request `ID` — the caller records this so the matching response's
    /// `InResponseTo` can be correlated (replay/CSRF protection).
    pub id: String,
    /// The full IdP URL to 302 the browser to (includes `SAMLRequest` and,
    /// when supplied, `RelayState`).
    pub location: String,
}

/// Build an AuthnRequest targeting `idp_sso_url` and encode it for the
/// HTTP-Redirect binding. `relay_state`, if given, round-trips back to us via
/// the response (we use it to send the operator to their intended page).
pub fn build(
    sp: &SpParams,
    idp_sso_url: &str,
    relay_state: Option<&str>,
) -> Result<AuthnRequest, SamlError> {
    let id = format!("_{}", uuid::Uuid::new_v4().simple());
    let issue_instant = OffsetDateTime::now_utc()
        .replace_nanosecond(0)
        .unwrap_or_else(|_| OffsetDateTime::now_utc())
        .format(&Rfc3339)
        .map_err(|e| SamlError::Timestamp(e.to_string()))?;

    let xml = format!(
        r#"<samlp:AuthnRequest xmlns:samlp="{NS_PROTOCOL}" xmlns:saml="{NS_ASSERTION}" ID="{id}" Version="2.0" IssueInstant="{instant}" Destination="{dest}" ProtocolBinding="{BINDING_POST}" AssertionConsumerServiceURL="{acs}"><saml:Issuer>{issuer}</saml:Issuer><samlp:NameIDPolicy Format="{NAMEID_EMAIL}" AllowCreate="true"/></samlp:AuthnRequest>"#,
        instant = issue_instant,
        dest = xml_escape(idp_sso_url),
        acs = xml_escape(&sp.acs_url),
        issuer = xml_escape(&sp.entity_id),
    );

    let encoded = deflate_base64(&xml)?;

    let sep = if idp_sso_url.contains('?') { '&' } else { '?' };
    let mut location = format!("{idp_sso_url}{sep}SAMLRequest={}", pct_encode(&encoded));
    if let Some(rs) = relay_state {
        location.push_str("&RelayState=");
        location.push_str(&pct_encode(rs));
    }

    Ok(AuthnRequest { id, location })
}

/// Raw-DEFLATE then base64 — the HTTP-Redirect binding's `SAMLRequest` payload.
fn deflate_base64(xml: &str) -> Result<String, SamlError> {
    let mut enc = DeflateEncoder::new(Vec::new(), Compression::default());
    enc.write_all(xml.as_bytes())
        .and_then(|()| enc.try_finish())
        .map_err(|e| SamlError::Xml(format!("deflate: {e}")))?;
    let compressed = enc
        .finish()
        .map_err(|e| SamlError::Xml(format!("deflate: {e}")))?;
    Ok(base64::engine::general_purpose::STANDARD.encode(compressed))
}

// `pct_encode` + `xml_escape` now live in `openpxe_core::encoding` (v0.5.4)
// — imported above.

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::read::DeflateDecoder;
    use std::io::Read;

    fn sp() -> SpParams {
        SpParams {
            entity_id: "https://pxe.example.com".into(),
            acs_url: "https://pxe.example.com/api/sso/acs".into(),
        }
    }

    fn pct_decode(s: &str) -> Vec<u8> {
        let bytes = s.as_bytes();
        let mut out = Vec::with_capacity(bytes.len());
        let mut i = 0;
        while i < bytes.len() {
            if bytes[i] == b'%' && i + 2 < bytes.len() {
                let hi = (bytes[i + 1] as char).to_digit(16).unwrap();
                let lo = (bytes[i + 2] as char).to_digit(16).unwrap();
                out.push((hi * 16 + lo) as u8);
                i += 3;
            } else {
                out.push(bytes[i]);
                i += 1;
            }
        }
        out
    }

    #[test]
    fn id_is_ncname_and_location_has_request() {
        let req = build(&sp(), "https://idp.example.com/sso", Some("/dashboard")).unwrap();
        assert!(req.id.starts_with('_'));
        assert!(req
            .location
            .starts_with("https://idp.example.com/sso?SAMLRequest="));
        assert!(req.location.contains("&RelayState=%2Fdashboard"));
    }

    #[test]
    fn redirect_payload_round_trips_to_our_authn_request() {
        let req = build(&sp(), "https://idp.example.com/sso", None).unwrap();
        // Pull SAMLRequest value out of the query string.
        let q = req.location.split("SAMLRequest=").nth(1).unwrap();
        let val = q.split('&').next().unwrap();
        let compressed = base64::engine::general_purpose::STANDARD
            .decode(pct_decode(val))
            .unwrap();
        let mut inflate = DeflateDecoder::new(&compressed[..]);
        let mut xml = String::new();
        inflate.read_to_string(&mut xml).unwrap();

        let doc = roxmltree::Document::parse(&xml).unwrap();
        let root = doc.root_element();
        assert_eq!(root.tag_name().name(), "AuthnRequest");
        assert_eq!(root.attribute("ID").unwrap(), req.id);
        assert_eq!(
            root.attribute("AssertionConsumerServiceURL").unwrap(),
            "https://pxe.example.com/api/sso/acs"
        );
        let issuer = root
            .descendants()
            .find(|n| n.tag_name().name() == "Issuer")
            .unwrap();
        assert_eq!(issuer.text().unwrap(), "https://pxe.example.com");
    }

    #[test]
    fn existing_query_uses_ampersand_separator() {
        let req = build(&sp(), "https://idp.example.com/sso?foo=bar", None).unwrap();
        assert!(req.location.contains("?foo=bar&SAMLRequest="));
    }
}