pdk_xml_validator_lib/
lib.rs1use 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#[non_exhaustive]
33#[derive(Debug, Error)]
34pub enum XmlValidationError {
35 #[error("Node depth count has been exceeded. Maximum allowed is: {0}")]
37 MaxDepthExceeded(usize),
38 #[error("Attribute count has been exceeded. Maximum allowed is: {0}")]
40 MaxAttributeCountExceeded(usize),
41 #[error("Text length has been exceeded. Maximum allowed is: {0}")]
43 MaxTextLengthExceeded(usize),
44 #[error("Comment length has been exceeded. Maximum allowed is: {0}")]
46 MaxCommentLengthExceeded(usize),
47 #[error("Attribute value length has been exceeded. Maximum allowed is: {0}")]
49 MaxAttributeLengthExceeded(usize),
50 #[error("Child count has been exceeded. Maximum allowed is: {0}")]
52 MaxChildCountExceeded(usize),
53 #[error("Can not read attributes: {0}")]
55 AttributeError(Box<dyn Error>),
56 #[error("Malformed document: {0}")]
58 MalformedDocument(Box<dyn Error>),
59 #[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 pub fn new() -> Self {
77 Self::default()
78 }
79
80 pub fn with_max_depth(mut self, depth: usize) -> Self {
82 self.max_depth = Some(depth);
83 self
84 }
85
86 pub fn with_max_attribute_count(mut self, count: usize) -> Self {
88 self.max_attribute_count = Some(count);
89 self
90 }
91
92 pub fn with_max_child_count(mut self, count: usize) -> Self {
94 self.max_child_count = Some(count);
95 self
96 }
97
98 pub fn with_max_text_length(mut self, length: usize) -> Self {
100 self.max_text_length = Some(length);
101 self
102 }
103
104 pub fn with_max_attribute_length(mut self, length: usize) -> Self {
106 self.max_attribute_length = Some(length);
107 self
108 }
109
110 pub fn with_max_comment_length(mut self, length: usize) -> Self {
112 self.max_comment_length = Some(length);
113 self
114 }
115
116 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
129pub struct XmlValidator {
131 max_depth: Option<usize>,
133 max_attribute_count: Option<usize>,
135 max_child_count: Option<usize>,
137 max_text_length: Option<usize>,
139 max_attribute_length: Option<usize>,
141 max_comment_length: Option<usize>,
143}
144
145struct ElementInfo {
146 child_count: usize,
147}
148
149impl XmlValidator {
150 pub async fn validate_stream(&self, stream: BodyStream<'_>) -> Result<(), XmlValidationError> {
152 self.validate(BodyStreamAsyncReader::new(stream).compat())
153 .await
154 }
155
156 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}