use crate::de::{Event, MappingStartEvent, Progress, ScalarEvent, SequenceStartEvent};
use crate::budget::{BudgetEvent, BudgetTracker, check_yaml_budget};
use crate::error::{self, ErrorImpl, Result, ScanError as PreScanError};
use crate::libyaml::error::Mark;
use crate::libyaml::parser::{Anchor, Event as YamlEvent, Parser};
use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::Arc;
fn anchor_to_string(anchor: &Anchor) -> String {
String::from_utf8_lossy(&anchor.0).into_owned()
}
pub(crate) struct Loader<'input> {
parser: Option<Parser<'input>>,
document_count: usize,
budget: Option<BudgetTracker<Anchor>>,
explicit_end_marker: bool,
}
#[derive(Clone)]
pub(crate) struct Document<'input> {
pub events: Vec<(Event<'input>, Mark)>,
pub error: Option<Arc<ErrorImpl>>,
pub aliases: Vec<usize>,
pub end_mark: Option<Mark>,
}
impl<'input> Loader<'input> {
pub fn new(
progress: Progress<'input>,
options: &crate::de::DeserializerOptions,
) -> Result<Self> {
let parser = match progress {
Progress::Str(s) => {
let bytes = s.as_bytes();
if let Some(b) = &options.budget {
match check_yaml_budget(s, b) {
Ok(rep) => {
if let Some(breach) = rep.breached {
return Err(error::new(ErrorImpl::BudgetExceeded(breach)));
}
}
Err(se) => {
let pse: PreScanError = (&se).into();
return Err(error::new(ErrorImpl::PreScan(pse)));
}
}
}
Parser::new(Cow::Borrowed(bytes))?
}
Progress::Slice(bytes) => {
if let Some(b) = &options.budget {
let s = match std::str::from_utf8(bytes) {
Ok(s) => s,
Err(utf8_err) => {
let valid = utf8_err.valid_up_to();
return Err(error::new(ErrorImpl::PreScan(PreScanError::new(
"budget pre-scan requires UTF-8 input".to_owned(),
valid,
1,
valid + 1,
))));
}
};
match check_yaml_budget(s, b) {
Ok(rep) => {
if let Some(breach) = rep.breached {
return Err(error::new(ErrorImpl::BudgetExceeded(breach)));
}
}
Err(se) => {
let pse: PreScanError = (&se).into();
return Err(error::new(ErrorImpl::PreScan(pse)));
}
}
}
Parser::new(Cow::Borrowed(bytes))?
}
Progress::Read(rdr) => Parser::from_reader(
rdr,
options
.budget
.as_ref()
.map(|budget| budget.max_total_scalar_bytes),
)?,
Progress::Iterable(_) | Progress::Document(_) => {
return Err(error::new(ErrorImpl::MoreThanOneDocument));
}
Progress::Fail(err) => return Err(error::shared(err)),
};
Ok(Loader {
parser: Some(parser),
document_count: 0,
budget: options.budget.as_ref().map(BudgetTracker::new),
explicit_end_marker: false,
})
}
pub fn next_document(&mut self) -> Option<Document<'input>> {
let Some(parser) = &mut self.parser else {
return None;
};
let first = self.document_count == 0;
self.document_count += 1;
let mut anchors = HashMap::new();
let mut document = Document {
events: Vec::new(),
error: None,
aliases: Vec::new(),
end_mark: None,
};
loop {
let (event, mark) = match parser.next() {
Ok((event, mark)) => {
if let Some(budget) = &mut self.budget {
let budget_event = match &event {
YamlEvent::StreamStart => BudgetEvent::StreamStart,
YamlEvent::StreamEnd => BudgetEvent::StreamEnd,
YamlEvent::DocumentStart => BudgetEvent::DocumentStart,
YamlEvent::DocumentEnd { .. } => BudgetEvent::DocumentEnd,
YamlEvent::Alias(alias) => BudgetEvent::Alias(alias.clone()),
YamlEvent::Scalar(scalar) => BudgetEvent::Scalar {
anchor: scalar.anchor.clone(),
scalar_bytes: scalar.value.len(),
},
YamlEvent::SequenceStart(sequence_start) => BudgetEvent::SequenceStart {
anchor: sequence_start.anchor.clone(),
},
YamlEvent::SequenceEnd => BudgetEvent::SequenceEnd,
YamlEvent::MappingStart(mapping_start) => BudgetEvent::MappingStart {
anchor: mapping_start.anchor.clone(),
},
YamlEvent::MappingEnd => BudgetEvent::MappingEnd,
YamlEvent::Void => BudgetEvent::Nothing,
};
if let Err(breach) = budget.observe(budget_event) {
self.parser = None;
document.error = Some(error::new(ErrorImpl::BudgetExceeded(breach)).shared());
if document.events.is_empty() {
document.events.push((Event::Void, mark));
}
return Some(document);
}
}
(event, mark)
}
Err(err) => {
let err = err.shared();
if !first
&& document.events.is_empty()
&& self.explicit_end_marker
&& !matches!(err.as_ref(), ErrorImpl::BudgetExceeded(_))
{
self.parser = None;
return None;
}
if document.events.is_empty() {
document.events.push((Event::Void, Mark::default()));
}
document.error = Some(err);
self.parser = None;
return Some(document);
}
};
let event = match event {
YamlEvent::StreamStart => continue,
YamlEvent::StreamEnd => {
self.parser = None;
return if first {
if document.events.is_empty() {
document.events.push((Event::Void, mark));
}
Some(document)
} else {
None
};
}
YamlEvent::DocumentStart => continue,
YamlEvent::DocumentEnd { implicit } => {
self.explicit_end_marker = !implicit;
document.end_mark = Some(mark);
return Some(document);
}
YamlEvent::Alias(alias) => match anchors.get(&alias) {
Some(id) => Event::Alias(*id),
None => {
document.error =
Some(error::new(ErrorImpl::UnknownAnchor(mark, alias)).shared());
return Some(document);
}
},
YamlEvent::Scalar(mut scalar) => {
let anchor_name = scalar.anchor.take().map(|a| {
let name = anchor_to_string(&a);
let id = document.aliases.len();
anchors.insert(a, id);
document.aliases.push(document.events.len());
name
});
let raw = String::from_utf8_lossy(&scalar.value).into_owned();
Event::Scalar(ScalarEvent {
anchor: anchor_name,
value: scalar,
raw,
})
}
YamlEvent::SequenceStart(mut sequence_start) => {
let anchor_name = sequence_start.anchor.take().map(|a| {
let name = anchor_to_string(&a);
let id = document.aliases.len();
anchors.insert(a, id);
document.aliases.push(document.events.len());
name
});
Event::SequenceStart(SequenceStartEvent {
anchor: anchor_name,
tag: sequence_start.tag,
})
}
YamlEvent::SequenceEnd => Event::SequenceEnd,
YamlEvent::MappingStart(mut mapping_start) => {
let anchor_name = mapping_start.anchor.take().map(|a| {
let name = anchor_to_string(&a);
let id = document.aliases.len();
anchors.insert(a, id);
document.aliases.push(document.events.len());
name
});
Event::MappingStart(MappingStartEvent {
anchor: anchor_name,
tag: mapping_start.tag,
})
}
YamlEvent::MappingEnd => Event::MappingEnd,
YamlEvent::Void => Event::Void,
};
document.events.push((event, mark));
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn anchored_scalar_event_keeps_anchor() {
let yaml = "a: &id 1\nb: *id\n";
let mut loader = Loader::new(
Progress::Str(yaml),
&crate::de::DeserializerOptions::default(),
)
.unwrap();
let document = loader.next_document().unwrap();
let mut found = false;
for (event, _) in &document.events {
if let Event::Scalar(scalar) = event {
if let Some(name) = &scalar.anchor {
assert_eq!(name, "id");
found = true;
}
}
}
assert!(found, "anchored scalar not found");
}
#[test]
fn anchored_sequence_event_keeps_anchor() {
let yaml = "a: &id [1, 2]\nb: *id\n";
let mut loader = Loader::new(
Progress::Str(yaml),
&crate::de::DeserializerOptions::default(),
)
.unwrap();
let document = loader.next_document().unwrap();
let mut found = false;
for (event, _) in &document.events {
if let Event::SequenceStart(sequence) = event {
if let Some(name) = &sequence.anchor {
assert_eq!(name, "id");
found = true;
}
}
}
assert!(found, "anchored sequence not found");
}
#[test]
fn anchored_mapping_event_keeps_anchor() {
let yaml = "a: &id {b: 1}\nc: *id\n";
let mut loader = Loader::new(
Progress::Str(yaml),
&crate::de::DeserializerOptions::default(),
)
.unwrap();
let document = loader.next_document().unwrap();
let mut found = false;
for (event, _) in &document.events {
if let Event::MappingStart(mapping) = event {
if let Some(name) = &mapping.anchor {
assert_eq!(name, "id");
found = true;
}
}
}
assert!(found, "anchored mapping not found");
}
#[test]
fn budgeted_slice_rejects_non_utf8_input() {
let mut opts = crate::de::DeserializerOptions::default();
opts.budget = Some(crate::budget::Budget::default());
let err = match Loader::new(Progress::Slice(&[0xFF]), &opts) {
Ok(_) => panic!("expected pre-scan error for non-UTF-8 input"),
Err(err) => err,
};
assert!(
err.to_string()
.contains("budget pre-scan requires UTF-8 input"),
"unexpected error: {err}"
);
}
#[test]
fn budgeted_reader_rejects_second_document_without_buffering() {
use std::io::{self, Read};
struct ChunkedReader {
chunks: Vec<&'static [u8]>,
next_chunk: usize,
}
impl Read for ChunkedReader {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if self.next_chunk >= self.chunks.len() {
return Ok(0);
}
let chunk = self.chunks[self.next_chunk];
let len = chunk.len().min(buf.len());
buf[..len].copy_from_slice(&chunk[..len]);
self.next_chunk += 1;
Ok(len)
}
}
let reader = ChunkedReader {
chunks: vec![b"---\na: 1\n", b"---\nb: [1, 2]\n"],
next_chunk: 0,
};
let mut opts = crate::de::DeserializerOptions::default();
let mut budget = crate::budget::Budget::default();
budget.max_nodes = 3;
opts.budget = Some(budget);
let mut loader = Loader::new(Progress::Read(Box::new(reader)), &opts).unwrap();
let first_document = loader.next_document().unwrap();
assert!(first_document.error.is_none());
let second_document = loader.next_document().unwrap();
assert!(matches!(
second_document.error,
Some(ref err) if matches!(err.as_ref(), ErrorImpl::BudgetExceeded(_))
));
}
}