Skip to main content

pdk_xml_validator_lib/
lib.rs

1// Copyright (c) 2026, Salesforce, Inc.,
2// All rights reserved.
3// For full license text, see the LICENSE.txt file
4
5//! # PDK XML Validator Library
6//!
7//! Provides XML validation functionality with configurable structural constraints,
8//! validating XML payloads incrementally from the request body stream.
9//!
10//! ## Highlights
11//!
12//! - Stream-first validation API for policy usage
13//! - Configurable limits for depth, attributes, children, text and comments
14//!
15//! ## Primary types
16//!
17//! - [`XmlValidatorBuilder`]: builder to configure validation constraints
18//! - [`XmlValidator`]: validator instance used to validate XML payload streams
19//! - [`XmlValidator::validate`]: validate from an async byte source
20//! - [`XmlValidator::validate_stream`]: validate a PDK body stream
21//! - [`XmlValidationError`]: error type for structural violations and malformed XML
22
23use classy::hl::{BodyStream, BodyStreamAsyncReader};
24use quick_xml::events::BytesStart;
25use quick_xml::{events::Event, reader::NsReader};
26use std::error::Error;
27use thiserror::Error;
28use tokio::io::{AsyncRead, BufReader};
29use tokio_util::compat::FuturesAsyncReadCompatExt;
30
31/// Errors that can occur during XML validation.
32#[non_exhaustive]
33#[derive(Debug, Error)]
34pub enum XmlValidationError {
35    /// The maximum allowed depth of the XML document has been exceeded.
36    #[error("Node depth count has been exceeded. Maximum allowed is: {0}")]
37    MaxDepthExceeded(usize),
38    /// The maximum allowed number of attributes in an element has been exceeded.
39    #[error("Attribute count has been exceeded. Maximum allowed is: {0}")]
40    MaxAttributeCountExceeded(usize),
41    /// The maximum allowed length of the text in an element has been exceeded.
42    #[error("Text length has been exceeded. Maximum allowed is: {0}")]
43    MaxTextLengthExceeded(usize),
44    /// The maximum allowed length of the comment in an element has been exceeded.
45    #[error("Comment length has been exceeded. Maximum allowed is: {0}")]
46    MaxCommentLengthExceeded(usize),
47    /// The maximum allowed length of the attribute value has been exceeded.
48    #[error("Attribute value length has been exceeded. Maximum allowed is: {0}")]
49    MaxAttributeLengthExceeded(usize),
50    /// The maximum allowed number of children in an element has been exceeded.
51    #[error("Child count has been exceeded. Maximum allowed is: {0}")]
52    MaxChildCountExceeded(usize),
53    /// The attribute cannot be read.
54    #[error("Can not read attributes: {0}")]
55    AttributeError(Box<dyn Error>),
56    /// The XML document is malformed.
57    #[error("Malformed document: {0}")]
58    MalformedDocument(Box<dyn Error>),
59    /// The XML document structure is invalid.
60    #[error("XML Validation Error: {0}")]
61    InvalidStructure(String),
62}
63
64#[derive(Debug, Default)]
65pub struct XmlValidatorBuilder {
66    max_depth: Option<usize>,
67    max_attribute_count: Option<usize>,
68    max_child_count: Option<usize>,
69    max_text_length: Option<usize>,
70    max_attribute_length: Option<usize>,
71    max_comment_length: Option<usize>,
72}
73
74impl XmlValidatorBuilder {
75    /// Creates a new builder without restrictions.
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    /// Sets the maximum allowed depth of the XML document.
81    pub fn with_max_depth(mut self, depth: usize) -> Self {
82        self.max_depth = Some(depth);
83        self
84    }
85
86    /// Sets the maximum allowed number of attributes in an element.
87    pub fn with_max_attribute_count(mut self, count: usize) -> Self {
88        self.max_attribute_count = Some(count);
89        self
90    }
91
92    /// Sets the maximum allowed number of children in an element.
93    pub fn with_max_child_count(mut self, count: usize) -> Self {
94        self.max_child_count = Some(count);
95        self
96    }
97
98    /// Sets the maximum allowed length of the text in an element.
99    pub fn with_max_text_length(mut self, length: usize) -> Self {
100        self.max_text_length = Some(length);
101        self
102    }
103
104    /// Sets the maximum allowed length of the attribute value.
105    pub fn with_max_attribute_length(mut self, length: usize) -> Self {
106        self.max_attribute_length = Some(length);
107        self
108    }
109
110    /// Sets the maximum allowed length of the comment in an element.
111    pub fn with_max_comment_length(mut self, length: usize) -> Self {
112        self.max_comment_length = Some(length);
113        self
114    }
115
116    /// Builds the validation instance with the specified configuration.
117    pub fn build(self) -> XmlValidator {
118        XmlValidator {
119            max_depth: self.max_depth,
120            max_attribute_count: self.max_attribute_count,
121            max_child_count: self.max_child_count,
122            max_text_length: self.max_text_length,
123            max_attribute_length: self.max_attribute_length,
124            max_comment_length: self.max_comment_length,
125        }
126    }
127}
128
129/// XML validator with configurable constraints.
130pub struct XmlValidator {
131    /// Maximum allowed depth of the XML document.
132    max_depth: Option<usize>,
133    /// Maximum allowed number of attributes in an element.
134    max_attribute_count: Option<usize>,
135    /// Maximum allowed number of children in an element.
136    max_child_count: Option<usize>,
137    /// Maximum allowed length of the text in an element.
138    max_text_length: Option<usize>,
139    /// Maximum allowed length of the attribute value.
140    max_attribute_length: Option<usize>,
141    /// Maximum allowed length of the comment in an element.
142    max_comment_length: Option<usize>,
143}
144
145struct ElementInfo {
146    child_count: usize,
147}
148
149impl XmlValidator {
150    /// Validates XML from a PDK body stream.
151    pub async fn validate_stream(&self, stream: BodyStream<'_>) -> Result<(), XmlValidationError> {
152        self.validate(BodyStreamAsyncReader::new(stream).compat())
153            .await
154    }
155
156    /// Validates XML from an async byte source.
157    pub async fn validate(&self, reader: impl AsyncRead + Unpin) -> Result<(), XmlValidationError> {
158        let reader = BufReader::new(reader);
159        let mut buf = Vec::with_capacity(1024);
160        let mut xml_reader = NsReader::from_reader(reader);
161        let mut depth = 0usize;
162        let mut element_stack: Vec<ElementInfo> = Vec::new();
163
164        loop {
165            let event = xml_reader.read_event_into_async(&mut buf).await;
166            match event {
167                Ok(ref event @ Event::Start(ref e)) | Ok(ref event @ Event::Empty(ref e)) => {
168                    depth += 1;
169
170                    if let Some(max_depth) = self.max_depth {
171                        if depth > max_depth {
172                            return Err(XmlValidationError::MaxDepthExceeded(max_depth));
173                        }
174                    }
175
176                    if let Event::Empty(_) = event {
177                        depth -= 1;
178                    }
179
180                    if let Some(max_attr_count) = self.max_attribute_count {
181                        if check_attribute_count(e, max_attr_count)
182                            .map_err(|err| XmlValidationError::AttributeError(Box::new(err)))?
183                        {
184                            return Err(XmlValidationError::MaxAttributeCountExceeded(
185                                max_attr_count,
186                            ));
187                        }
188                    }
189
190                    if let Some(max_attr_len) = self.max_attribute_length {
191                        if check_attribute_length(e, max_attr_len)
192                            .map_err(|err| XmlValidationError::AttributeError(Box::new(err)))?
193                        {
194                            return Err(XmlValidationError::MaxAttributeLengthExceeded(
195                                max_attr_len,
196                            ));
197                        }
198                    }
199
200                    if !element_stack.is_empty() {
201                        if let Some(parent) = element_stack.last_mut() {
202                            parent.child_count += 1;
203                            if let Some(max_child_count) = self.max_child_count {
204                                if parent.child_count > max_child_count {
205                                    return Err(XmlValidationError::MaxChildCountExceeded(
206                                        max_child_count,
207                                    ));
208                                }
209                            }
210                        }
211                    }
212
213                    if let Event::Start(_) = event {
214                        element_stack.push(ElementInfo { child_count: 0 });
215                    }
216                }
217                Ok(Event::End(_)) => {
218                    element_stack.pop();
219                    depth -= 1;
220                }
221                Ok(Event::Eof) => {
222                    if !element_stack.is_empty() {
223                        return Err(XmlValidationError::InvalidStructure(
224                            "XML document structures must start and end within the same entity."
225                                .to_string(),
226                        ));
227                    }
228                    return Ok(());
229                }
230                Ok(Event::Text(t)) => {
231                    let is_only_whitespace = t.as_ref().iter().all(|&b| b.is_ascii_whitespace());
232                    if is_only_whitespace {
233                        continue;
234                    }
235
236                    if element_stack.is_empty() {
237                        return Err(XmlValidationError::InvalidStructure(
238                            "Text found outside of an element".to_string(),
239                        ));
240                    }
241
242                    if let Some(max_text_len) = self.max_text_length {
243                        let text = t
244                            .unescape()
245                            .map_err(|err| XmlValidationError::MalformedDocument(Box::new(err)))?;
246                        if text.len() > max_text_len {
247                            return Err(XmlValidationError::MaxTextLengthExceeded(max_text_len));
248                        }
249                    }
250                }
251                Ok(Event::Comment(t)) => {
252                    if let Some(max_comment_len) = self.max_comment_length {
253                        let comment = t
254                            .unescape()
255                            .map_err(|err| XmlValidationError::MalformedDocument(Box::new(err)))?;
256                        if comment.len() > max_comment_len {
257                            return Err(XmlValidationError::MaxCommentLengthExceeded(
258                                max_comment_len,
259                            ));
260                        }
261                    }
262                }
263                Err(e) => {
264                    return Err(XmlValidationError::MalformedDocument(Box::new(e)));
265                }
266                Ok(_) => {}
267            }
268            buf.clear();
269        }
270    }
271}
272
273fn check_attribute_count(
274    e: &BytesStart<'_>,
275    max_attr_count: usize,
276) -> Result<bool, quick_xml::Error> {
277    let mut count = 0usize;
278    for attr in e.attributes() {
279        let attr_name = String::from_utf8_lossy(attr?.key.0);
280        if !attr_name.starts_with("xmlns") {
281            count += 1;
282        }
283        if count > max_attr_count {
284            return Ok(true);
285        }
286    }
287    Ok(false)
288}
289
290fn check_attribute_length(e: &BytesStart<'_>, max_len: usize) -> Result<bool, quick_xml::Error> {
291    for attr in e.attributes() {
292        let attr = attr?;
293        let attr_value = attr.unescape_value()?;
294        if attr_value.len() > max_len {
295            return Ok(true);
296        }
297    }
298    Ok(false)
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304    use quick_xml::events::Event;
305    use quick_xml::reader::Reader;
306
307    fn build_validator() -> XmlValidator {
308        XmlValidatorBuilder::new()
309            .with_max_depth(3)
310            .with_max_attribute_count(2)
311            .with_max_child_count(2)
312            .with_max_text_length(10)
313            .with_max_attribute_length(5)
314            .with_max_comment_length(10)
315            .build()
316    }
317
318    #[tokio::test]
319    async fn validate_accepts_valid_xml_stream() {
320        let validator = build_validator();
321
322        let xml = r#"<root><child attr="ok">text</child><child2/></root>"#;
323        let result = validator.validate(xml.as_bytes()).await;
324        assert!(result.is_ok());
325    }
326
327    #[tokio::test]
328    async fn validate_accepts_vec_buffer_as_slice() {
329        let validator = XmlValidatorBuilder::new().build();
330        let bytes = b"<root><n/></root>".to_vec();
331        let result = validator.validate(bytes.as_slice()).await;
332        assert!(result.is_ok());
333    }
334
335    #[tokio::test]
336    async fn validate_vec_single_byte_slice_is_rejected() {
337        let validator = XmlValidatorBuilder::new().build();
338        let bytes = vec![0u8];
339        let result = validator.validate(bytes.as_slice()).await;
340        assert!(result.is_err());
341    }
342
343    #[tokio::test]
344    async fn validate_rejects_depth_over_limit() {
345        let validator = XmlValidatorBuilder::new().with_max_depth(1).build();
346
347        let xml = r#"<root><child><grandchild/></child></root>"#;
348        let result = validator.validate(xml.as_bytes()).await;
349        assert!(matches!(
350            result,
351            Err(XmlValidationError::MaxDepthExceeded(1))
352        ));
353    }
354
355    #[tokio::test]
356    async fn validate_rejects_attribute_count_over_limit() {
357        let validator = build_validator();
358        let xml = r#"<root a="1" b="2" c="3"/>"#;
359        let result = validator.validate(xml.as_bytes()).await;
360        assert!(matches!(
361            result,
362            Err(XmlValidationError::MaxAttributeCountExceeded(2))
363        ));
364    }
365
366    #[tokio::test]
367    async fn validate_rejects_text_outside_element() {
368        let validator = build_validator();
369        let xml = r#"text<root/>"#;
370        let result = validator.validate(xml.as_bytes()).await;
371        assert!(matches!(
372            result,
373            Err(XmlValidationError::InvalidStructure(_))
374        ));
375    }
376
377    #[tokio::test]
378    async fn validate_rejects_attribute_length_over_limit() {
379        let validator = build_validator();
380        let xml = r#"<root attr="123456"/>"#;
381        let result = validator.validate(xml.as_bytes()).await;
382        assert!(matches!(
383            result,
384            Err(XmlValidationError::MaxAttributeLengthExceeded(5))
385        ));
386    }
387
388    #[tokio::test]
389    async fn validate_rejects_child_count_over_limit() {
390        let validator = XmlValidatorBuilder::new().with_max_child_count(1).build();
391        let xml = r#"<root><c1/><c2/></root>"#;
392        let result = validator.validate(xml.as_bytes()).await;
393        assert!(matches!(
394            result,
395            Err(XmlValidationError::MaxChildCountExceeded(1))
396        ));
397    }
398
399    #[tokio::test]
400    async fn validate_rejects_text_length_over_limit() {
401        let validator = XmlValidatorBuilder::new().with_max_text_length(3).build();
402        let xml = r#"<root>abcd</root>"#;
403        let result = validator.validate(xml.as_bytes()).await;
404        assert!(matches!(
405            result,
406            Err(XmlValidationError::MaxTextLengthExceeded(3))
407        ));
408    }
409
410    #[tokio::test]
411    async fn validate_rejects_comment_length_over_limit() {
412        let validator = XmlValidatorBuilder::new()
413            .with_max_comment_length(3)
414            .build();
415        let xml = r#"<root><!--abcd--></root>"#;
416        let result = validator.validate(xml.as_bytes()).await;
417        assert!(matches!(
418            result,
419            Err(XmlValidationError::MaxCommentLengthExceeded(3))
420        ));
421    }
422
423    #[tokio::test]
424    async fn validate_excludes_xmlns_from_attribute_count() {
425        let validator = XmlValidatorBuilder::new()
426            .with_max_attribute_count(1)
427            .build();
428        let xml = r#"<root xmlns="urn:test" xmlns:x="urn:test2" attr="1"/>"#;
429        let result = validator.validate(xml.as_bytes()).await;
430        assert!(result.is_ok());
431    }
432
433    #[tokio::test]
434    async fn validate_accepts_whitespace_outside_root() {
435        let validator = XmlValidatorBuilder::new().build();
436        let xml = "  \n\t<root/>";
437        let result = validator.validate(xml.as_bytes()).await;
438        assert!(result.is_ok());
439    }
440
441    #[tokio::test]
442    async fn validate_rejects_malformed_document() {
443        let validator = XmlValidatorBuilder::new().build();
444        let xml = r#"<root><child></root>"#;
445        let result = validator.validate(xml.as_bytes()).await;
446        assert!(matches!(
447            result,
448            Err(XmlValidationError::MalformedDocument(_))
449        ));
450    }
451
452    #[tokio::test]
453    async fn validate_rejects_unclosed_structure_at_eof() {
454        let validator = XmlValidatorBuilder::new().build();
455        let xml = r#"<root>"#;
456        let result = validator.validate(xml.as_bytes()).await;
457        assert!(matches!(
458            result,
459            Err(XmlValidationError::InvalidStructure(_))
460        ));
461    }
462
463    #[test]
464    fn check_attribute_helpers_cover_non_exceeded_paths() {
465        let mut count_reader = Reader::from_str(r#"<root xmlns="urn:test" attr="123"/>"#);
466        let mut count_buf = Vec::new();
467        let count_event = count_reader
468            .read_event_into(&mut count_buf)
469            .expect("must parse start");
470        let count_start = match count_event {
471            Event::Start(e) | Event::Empty(e) => e,
472            _ => panic!("expected start or empty event"),
473        };
474
475        let mut length_reader = Reader::from_str(r#"<root attr="123"/>"#);
476        let mut length_buf = Vec::new();
477        let length_event = length_reader
478            .read_event_into(&mut length_buf)
479            .expect("must parse start");
480        let length_start = match length_event {
481            Event::Start(e) | Event::Empty(e) => e,
482            _ => panic!("expected start or empty event"),
483        };
484
485        let count_exceeded =
486            check_attribute_count(&count_start, 1).expect("count check must succeed");
487        let len_exceeded =
488            check_attribute_length(&length_start, 3).expect("length check must succeed");
489        assert!(!count_exceeded);
490        assert!(!len_exceeded);
491    }
492}