pdk_xml_validator_lib/

lib.rs

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! # PDK XML Validator Library
//!
//! Provides XML validation functionality with configurable structural constraints,
//! validating XML payloads incrementally from the request body stream.
//!
//! ## Highlights
//!
//! - Stream-first validation API for policy usage
//! - Configurable limits for depth, attributes, children, text and comments
//!
//! ## Primary types
//!
//! - [`XmlValidatorBuilder`]: builder to configure validation constraints
//! - [`XmlValidator`]: validator instance used to validate XML payload streams
//! - [`XmlValidator::validate`]: validate from an async byte source
//! - [`XmlValidator::validate_stream`]: validate a PDK body stream
//! - [`XmlValidationError`]: error type for structural violations and malformed XML

use classy::hl::{BodyStream, BodyStreamAsyncReader};
use quick_xml::events::BytesStart;
use quick_xml::{events::Event, reader::NsReader};
use std::error::Error;
use thiserror::Error;
use tokio::io::{AsyncRead, BufReader};
use tokio_util::compat::FuturesAsyncReadCompatExt;

/// Errors that can occur during XML validation.
#[non_exhaustive]
#[derive(Debug, Error)]
pub enum XmlValidationError {
    /// The maximum allowed depth of the XML document has been exceeded.
    #[error("Node depth count has been exceeded. Maximum allowed is: {0}")]
    MaxDepthExceeded(usize),
    /// The maximum allowed number of attributes in an element has been exceeded.
    #[error("Attribute count has been exceeded. Maximum allowed is: {0}")]
    MaxAttributeCountExceeded(usize),
    /// The maximum allowed length of the text in an element has been exceeded.
    #[error("Text length has been exceeded. Maximum allowed is: {0}")]
    MaxTextLengthExceeded(usize),
    /// The maximum allowed length of the comment in an element has been exceeded.
    #[error("Comment length has been exceeded. Maximum allowed is: {0}")]
    MaxCommentLengthExceeded(usize),
    /// The maximum allowed length of the attribute value has been exceeded.
    #[error("Attribute value length has been exceeded. Maximum allowed is: {0}")]
    MaxAttributeLengthExceeded(usize),
    /// The maximum allowed number of children in an element has been exceeded.
    #[error("Child count has been exceeded. Maximum allowed is: {0}")]
    MaxChildCountExceeded(usize),
    /// The attribute cannot be read.
    #[error("Can not read attributes: {0}")]
    AttributeError(Box<dyn Error>),
    /// The XML document is malformed.
    #[error("Malformed document: {0}")]
    MalformedDocument(Box<dyn Error>),
    /// The XML document structure is invalid.
    #[error("XML Validation Error: {0}")]
    InvalidStructure(String),
}

#[derive(Debug, Default)]
pub struct XmlValidatorBuilder {
    max_depth: Option<usize>,
    max_attribute_count: Option<usize>,
    max_child_count: Option<usize>,
    max_text_length: Option<usize>,
    max_attribute_length: Option<usize>,
    max_comment_length: Option<usize>,
}

impl XmlValidatorBuilder {
    /// Creates a new builder without restrictions.
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the maximum allowed depth of the XML document.
    pub fn with_max_depth(mut self, depth: usize) -> Self {
        self.max_depth = Some(depth);
        self
    }

    /// Sets the maximum allowed number of attributes in an element.
    pub fn with_max_attribute_count(mut self, count: usize) -> Self {
        self.max_attribute_count = Some(count);
        self
    }

    /// Sets the maximum allowed number of children in an element.
    pub fn with_max_child_count(mut self, count: usize) -> Self {
        self.max_child_count = Some(count);
        self
    }

    /// Sets the maximum allowed length of the text in an element.
    pub fn with_max_text_length(mut self, length: usize) -> Self {
        self.max_text_length = Some(length);
        self
    }

    /// Sets the maximum allowed length of the attribute value.
    pub fn with_max_attribute_length(mut self, length: usize) -> Self {
        self.max_attribute_length = Some(length);
        self
    }

    /// Sets the maximum allowed length of the comment in an element.
    pub fn with_max_comment_length(mut self, length: usize) -> Self {
        self.max_comment_length = Some(length);
        self
    }

    /// Builds the validation instance with the specified configuration.
    pub fn build(self) -> XmlValidator {
        XmlValidator {
            max_depth: self.max_depth,
            max_attribute_count: self.max_attribute_count,
            max_child_count: self.max_child_count,
            max_text_length: self.max_text_length,
            max_attribute_length: self.max_attribute_length,
            max_comment_length: self.max_comment_length,
        }
    }
}

/// XML validator with configurable constraints.
pub struct XmlValidator {
    /// Maximum allowed depth of the XML document.
    max_depth: Option<usize>,
    /// Maximum allowed number of attributes in an element.
    max_attribute_count: Option<usize>,
    /// Maximum allowed number of children in an element.
    max_child_count: Option<usize>,
    /// Maximum allowed length of the text in an element.
    max_text_length: Option<usize>,
    /// Maximum allowed length of the attribute value.
    max_attribute_length: Option<usize>,
    /// Maximum allowed length of the comment in an element.
    max_comment_length: Option<usize>,
}

struct ElementInfo {
    child_count: usize,
}

impl XmlValidator {
    /// Validates XML from a PDK body stream.
    pub async fn validate_stream(&self, stream: BodyStream<'_>) -> Result<(), XmlValidationError> {
        self.validate(BodyStreamAsyncReader::new(stream).compat())
            .await
    }

    /// Validates XML from an async byte source.
    pub async fn validate(&self, reader: impl AsyncRead + Unpin) -> Result<(), XmlValidationError> {
        let reader = BufReader::new(reader);
        let mut buf = Vec::with_capacity(1024);
        let mut xml_reader = NsReader::from_reader(reader);
        let mut depth = 0usize;
        let mut element_stack: Vec<ElementInfo> = Vec::new();
        // Decoded length of the current text run. quick-xml 0.41 splits
        // entity-bearing text into Text + GeneralRef + Text events, so we must
        // accumulate the reassembled run and check the total against
        // `max_text_length` — checking each fragment in isolation (and letting
        // GeneralRef count as 0) would let entity-heavy text slip past a limit
        // that quick-xml 0.37 rejected. Reset on any structural boundary.
        let mut text_run_len = 0usize;

        loop {
            let event = xml_reader.read_event_into_async(&mut buf).await;
            match event {
                Ok(ref event @ Event::Start(ref e)) | Ok(ref event @ Event::Empty(ref e)) => {
                    // New element boundary: text limits are per-element.
                    text_run_len = 0;
                    depth += 1;

                    if let Some(max_depth) = self.max_depth {
                        if depth > max_depth {
                            return Err(XmlValidationError::MaxDepthExceeded(max_depth));
                        }
                    }

                    if let Event::Empty(_) = event {
                        depth -= 1;
                    }

                    if let Some(max_attr_count) = self.max_attribute_count {
                        if check_attribute_count(e, max_attr_count)
                            .map_err(|err| XmlValidationError::AttributeError(Box::new(err)))?
                        {
                            return Err(XmlValidationError::MaxAttributeCountExceeded(
                                max_attr_count,
                            ));
                        }
                    }

                    if let Some(max_attr_len) = self.max_attribute_length {
                        if check_attribute_length(e, max_attr_len)
                            .map_err(|err| XmlValidationError::AttributeError(Box::new(err)))?
                        {
                            return Err(XmlValidationError::MaxAttributeLengthExceeded(
                                max_attr_len,
                            ));
                        }
                    }

                    if !element_stack.is_empty() {
                        if let Some(parent) = element_stack.last_mut() {
                            parent.child_count += 1;
                            if let Some(max_child_count) = self.max_child_count {
                                if parent.child_count > max_child_count {
                                    return Err(XmlValidationError::MaxChildCountExceeded(
                                        max_child_count,
                                    ));
                                }
                            }
                        }
                    }

                    if let Event::Start(_) = event {
                        element_stack.push(ElementInfo { child_count: 0 });
                    }
                }
                Ok(Event::End(_)) => {
                    text_run_len = 0;
                    element_stack.pop();
                    depth -= 1;
                }
                Ok(Event::Eof) => {
                    if !element_stack.is_empty() {
                        return Err(XmlValidationError::InvalidStructure(
                            "XML document structures must start and end within the same entity."
                                .to_string(),
                        ));
                    }
                    return Ok(());
                }
                Ok(Event::Text(t)) => {
                    let is_only_whitespace = t.as_ref().iter().all(|&b| b.is_ascii_whitespace());
                    if is_only_whitespace {
                        continue;
                    }

                    if element_stack.is_empty() {
                        return Err(XmlValidationError::InvalidStructure(
                            "Text found outside of an element".to_string(),
                        ));
                    }

                    if let Some(max_text_len) = self.max_text_length {
                        text_run_len = text_run_len.saturating_add(unescape_event(&t)?.len());
                        if text_run_len > max_text_len {
                            return Err(XmlValidationError::MaxTextLengthExceeded(max_text_len));
                        }
                    }
                }
                Ok(Event::GeneralRef(r)) => {
                    // quick-xml 0.41 emits entity references (`&lt;`, `&#60;`, …)
                    // as their own event, splitting a text run. Count the
                    // expanded length toward the current run so `max_text_length`
                    // sees the same reassembled text quick-xml 0.37 did.
                    if let Some(max_text_len) = self.max_text_length {
                        text_run_len = text_run_len.saturating_add(resolve_general_ref(&r)?.len());
                        if text_run_len > max_text_len {
                            return Err(XmlValidationError::MaxTextLengthExceeded(max_text_len));
                        }
                    }
                }
                Ok(Event::Comment(t)) => {
                    if let Some(max_comment_len) = self.max_comment_length {
                        let comment = unescape_event(&t)?;
                        if comment.len() > max_comment_len {
                            return Err(XmlValidationError::MaxCommentLengthExceeded(
                                max_comment_len,
                            ));
                        }
                    }
                }
                Err(e) => {
                    return Err(XmlValidationError::MalformedDocument(Box::new(e)));
                }
                Ok(_) => {}
            }
            buf.clear();
        }
    }
}

/// Decodes an event's raw bytes and resolves XML entity escapes, mapping any
/// failure to a [`XmlValidationError::MalformedDocument`].
fn unescape_event(t: &quick_xml::events::BytesText<'_>) -> Result<String, XmlValidationError> {
    let decoded = t
        .decode()
        .map_err(|err| XmlValidationError::MalformedDocument(Box::new(err)))?;
    let unescaped = quick_xml::escape::unescape(&decoded)
        .map_err(|err| XmlValidationError::MalformedDocument(Box::new(err)))?;
    Ok(unescaped.into_owned())
}

/// Resolves an entity reference event (the `name` between `&` and `;`) to its
/// replacement text, so its expanded length can be counted toward the text-run
/// limit. Reconstructs the `&name;` form and reuses the same unescape path as
/// [`unescape_event`]; an unknown entity surfaces as
/// [`XmlValidationError::MalformedDocument`], matching quick-xml 0.37 behavior.
fn resolve_general_ref(r: &quick_xml::events::BytesRef<'_>) -> Result<String, XmlValidationError> {
    let name = r
        .decode()
        .map_err(|err| XmlValidationError::MalformedDocument(Box::new(err)))?;
    let reconstructed = format!("&{name};");
    let unescaped = quick_xml::escape::unescape(&reconstructed)
        .map_err(|err| XmlValidationError::MalformedDocument(Box::new(err)))?;
    Ok(unescaped.into_owned())
}

fn check_attribute_count(
    e: &BytesStart<'_>,
    max_attr_count: usize,
) -> Result<bool, quick_xml::Error> {
    let mut count = 0usize;
    for attr in e.attributes() {
        let attr_name = String::from_utf8_lossy(attr?.key.0);
        if !attr_name.starts_with("xmlns") {
            count += 1;
        }
        if count > max_attr_count {
            return Ok(true);
        }
    }
    Ok(false)
}

fn check_attribute_length(e: &BytesStart<'_>, max_len: usize) -> Result<bool, quick_xml::Error> {
    for attr in e.attributes() {
        let attr = attr?;
        // Decode bytes to UTF-8 string
        let decoded = std::str::from_utf8(attr.value.as_ref())
            .map_err(quick_xml::encoding::EncodingError::Utf8)?;
        // Unescape XML entities only (no normalization)
        let unescaped = quick_xml::escape::unescape(decoded)?;
        if unescaped.len() > max_len {
            return Ok(true);
        }
    }
    Ok(false)
}

#[cfg(test)]
mod tests {
    use super::*;
    use quick_xml::events::Event;
    use quick_xml::reader::Reader;

    fn build_validator() -> XmlValidator {
        XmlValidatorBuilder::new()
            .with_max_depth(3)
            .with_max_attribute_count(2)
            .with_max_child_count(2)
            .with_max_text_length(10)
            .with_max_attribute_length(5)
            .with_max_comment_length(10)
            .build()
    }

    #[tokio::test]
    async fn validate_accepts_valid_xml_stream() {
        let validator = build_validator();

        let xml = r#"<root><child attr="ok">text</child><child2/></root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn validate_accepts_vec_buffer_as_slice() {
        let validator = XmlValidatorBuilder::new().build();
        let bytes = b"<root><n/></root>".to_vec();
        let result = validator.validate(bytes.as_slice()).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn validate_vec_single_byte_slice_is_rejected() {
        let validator = XmlValidatorBuilder::new().build();
        let bytes = vec![0u8];
        let result = validator.validate(bytes.as_slice()).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn validate_rejects_depth_over_limit() {
        let validator = XmlValidatorBuilder::new().with_max_depth(1).build();

        let xml = r#"<root><child><grandchild/></child></root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::MaxDepthExceeded(1))
        ));
    }

    #[tokio::test]
    async fn validate_rejects_attribute_count_over_limit() {
        let validator = build_validator();
        let xml = r#"<root a="1" b="2" c="3"/>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::MaxAttributeCountExceeded(2))
        ));
    }

    #[tokio::test]
    async fn validate_rejects_text_outside_element() {
        let validator = build_validator();
        let xml = r#"text<root/>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::InvalidStructure(_))
        ));
    }

    #[tokio::test]
    async fn validate_rejects_attribute_length_over_limit() {
        let validator = build_validator();
        let xml = r#"<root attr="123456"/>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::MaxAttributeLengthExceeded(5))
        ));
    }

    #[tokio::test]
    async fn validate_rejects_child_count_over_limit() {
        let validator = XmlValidatorBuilder::new().with_max_child_count(1).build();
        let xml = r#"<root><c1/><c2/></root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::MaxChildCountExceeded(1))
        ));
    }

    #[tokio::test]
    async fn validate_rejects_text_length_over_limit() {
        let validator = XmlValidatorBuilder::new().with_max_text_length(3).build();
        let xml = r#"<root>abcd</root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::MaxTextLengthExceeded(3))
        ));
    }

    #[tokio::test]
    async fn validate_rejects_entity_text_length_over_limit() {
        // quick-xml 0.41 splits this into Text + GeneralRef + Text events. The
        // reassembled decoded text (`<<<<<<`, 6 chars) must be measured against
        // the limit as a whole, not per-fragment (each `&lt;` expands to a
        // single `<`, and a naive per-fragment/GeneralRef=0 check would let it
        // through).
        let validator = XmlValidatorBuilder::new().with_max_text_length(3).build();
        let xml = r#"<root>&lt;&lt;&lt;&lt;&lt;&lt;</root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::MaxTextLengthExceeded(3))
        ));
    }

    #[tokio::test]
    async fn validate_accepts_entity_text_within_limit() {
        // Same splitting, but the reassembled `<<<` (3 chars) is within the
        // limit, so it must be accepted.
        let validator = XmlValidatorBuilder::new().with_max_text_length(3).build();
        let xml = r#"<root>&lt;&lt;&lt;</root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(result.is_ok(), "unexpected: {:?}", result);
    }

    #[tokio::test]
    async fn validate_resets_text_run_between_elements() {
        // Text limits are per-element: two sibling elements each under the limit
        // must pass even though their combined text exceeds it.
        let validator = XmlValidatorBuilder::new().with_max_text_length(3).build();
        let xml = r#"<root><a>ab</a><b>cd</b></root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(result.is_ok(), "unexpected: {:?}", result);
    }

    #[tokio::test]
    async fn validate_rejects_comment_length_over_limit() {
        let validator = XmlValidatorBuilder::new()
            .with_max_comment_length(3)
            .build();
        let xml = r#"<root><!--abcd--></root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::MaxCommentLengthExceeded(3))
        ));
    }

    #[tokio::test]
    async fn validate_excludes_xmlns_from_attribute_count() {
        let validator = XmlValidatorBuilder::new()
            .with_max_attribute_count(1)
            .build();
        let xml = r#"<root xmlns="urn:test" xmlns:x="urn:test2" attr="1"/>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn validate_accepts_whitespace_outside_root() {
        let validator = XmlValidatorBuilder::new().build();
        let xml = "  \n\t<root/>";
        let result = validator.validate(xml.as_bytes()).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn validate_rejects_malformed_document() {
        let validator = XmlValidatorBuilder::new().build();
        let xml = r#"<root><child></root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::MalformedDocument(_))
        ));
    }

    #[tokio::test]
    async fn validate_rejects_unclosed_structure_at_eof() {
        let validator = XmlValidatorBuilder::new().build();
        let xml = r#"<root>"#;
        let result = validator.validate(xml.as_bytes()).await;
        assert!(matches!(
            result,
            Err(XmlValidationError::InvalidStructure(_))
        ));
    }

    #[test]
    fn check_attribute_helpers_cover_non_exceeded_paths() {
        let mut count_reader = Reader::from_str(r#"<root xmlns="urn:test" attr="123"/>"#);
        let mut count_buf = Vec::new();
        let count_event = count_reader
            .read_event_into(&mut count_buf)
            .expect("must parse start");
        let count_start = match count_event {
            Event::Start(e) | Event::Empty(e) => e,
            _ => panic!("expected start or empty event"),
        };

        let mut length_reader = Reader::from_str(r#"<root attr="123"/>"#);
        let mut length_buf = Vec::new();
        let length_event = length_reader
            .read_event_into(&mut length_buf)
            .expect("must parse start");
        let length_start = match length_event {
            Event::Start(e) | Event::Empty(e) => e,
            _ => panic!("expected start or empty event"),
        };

        let count_exceeded =
            check_attribute_count(&count_start, 1).expect("count check must succeed");
        let len_exceeded =
            check_attribute_length(&length_start, 3).expect("length check must succeed");
        assert!(!count_exceeded);
        assert!(!len_exceeded);
    }
}