From 5f44fda920882006e48899886711c83b4d679678 Mon Sep 17 00:00:00 2001 From: jskoiz <20649937+jskoiz@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:32:54 -1000 Subject: [PATCH] Aggregate byte-aware alias expansion accounting --- SECURITY.md | 5 +- docs/COMPATIBILITY.md | 2 +- docs/untrusted-input.md | 6 +- src/event_de/mod.rs | 24 +++--- src/event_de/serde_impl.rs | 25 +++--- src/event_de/source.rs | 168 ++++++++++++++++++++++++++----------- src/event_de/tests.rs | 25 +++++- src/parse.rs | 43 ++++++---- src/schema.rs | 62 ++++++++++++-- tests/diagnostics.rs | 49 +++++++++-- 10 files changed, 301 insertions(+), 108 deletions(-) diff --git a/SECURITY.md b/SECURITY.md index 600be7c..29311c6 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -23,8 +23,9 @@ The default loader posture is bounded for untrusted YAML inputs: - `LoadOptions` applies a default 64 MiB input byte ceiling. - Callers can tune `max_input_bytes()` or explicitly opt out with `without_input_limit()` only after bounding the source themselves. -- Alias expansion uses an input-derived budget by default and can be tuned with - `max_alias_expansion_nodes()`. +- Alias expansion uses an input-derived aggregate budget by default and can be + tuned with `max_alias_expansion_nodes()`. The budget accumulates clone work, + scalar bytes, and collection-item work across the complete input stream. - Recursive aliases are rejected. - Default nesting, scalar-size, and collection-item limits protect parser, loader, Serde, and lossless entrypoints from unbounded structural work: diff --git a/docs/COMPATIBILITY.md b/docs/COMPATIBILITY.md index 66555de..8000f94 100644 --- a/docs/COMPATIBILITY.md +++ b/docs/COMPATIBILITY.md @@ -139,7 +139,7 @@ The defended input is untrusted YAML at every load entrypoint. With default `LoadOptions`, the crate rejects: - input above **64 MiB** before parsing, -- alias-expansion bombs (input-derived budget) and recursive aliases, +- alias-expansion bombs (input-derived aggregate work budget) and recursive aliases, - nesting beyond **128**, scalars above **1 MiB**, and collections above **16,384** entries, diff --git a/docs/untrusted-input.md b/docs/untrusted-input.md index 6dbdeaa..22dabdb 100644 --- a/docs/untrusted-input.md +++ b/docs/untrusted-input.md @@ -18,7 +18,7 @@ out of the box: | Nesting depth | 128 | Deeply nested block/flow bombs | | Scalar size | 1 MiB | Single giant scalars | | Collection size | 16,384 entries | Wide sequence/mapping bombs | -| Alias expansion | input-derived budget | Billion-laughs alias bombs | +| Alias expansion | input-derived aggregate work budget | Billion-laughs alias bombs and oversized alias materialization | | Recursive aliases | — | always rejected | The defaults accept real-world config (Kubernetes CRDs, OpenAPI, Compose) while @@ -39,7 +39,9 @@ let cfg: Config = LoadOptions::new() ``` All knobs: `max_input_bytes`, `max_alias_expansion_nodes`, `max_nesting_depth`, -`max_scalar_bytes`, `max_collection_items`. +`max_scalar_bytes`, `max_collection_items`. The alias setting applies to the +aggregate stream work across clone operations, scalar bytes, and collection +items; it is not reset between documents. ## Relax — only when you've bounded the source yourself diff --git a/src/event_de/mod.rs b/src/event_de/mod.rs index 9061ea9..04a6567 100644 --- a/src/event_de/mod.rs +++ b/src/event_de/mod.rs @@ -9,13 +9,13 @@ use crate::{ Event, EventMeta, ScalarStyle, merge_policy_for_schema, parse_scalar_with_schema, schema_for_directives, }, - schema::{LoadOptions, Schema}, + schema::{AliasExpansionBudget, LoadOptions, Schema}, }; use serde::de::{ self, DeserializeOwned, DeserializeSeed, EnumAccess, IntoDeserializer, MapAccess, SeqAccess, VariantAccess, Visitor, }; -use std::{collections::HashMap, io::Read, marker::PhantomData}; +use std::{cell::Cell, collections::HashMap, io::Read, marker::PhantomData}; #[cfg_attr(not(test), allow(dead_code))] pub(crate) fn from_str_with_options<'de, T>(input: &'de str, options: LoadOptions) -> Result @@ -23,7 +23,7 @@ where T: serde::Deserialize<'de>, { let configured_schema = options.selected_schema(); - let replay_budget = options.alias_expansion_budget(input.len()); + let alias_budget = Cell::new(options.alias_expansion_budget(input.len())); let max_nesting_depth = options.selected_max_nesting_depth(); let events = crate::parse::EventStream::from_str_with_options(input, options)? .collect::>>()?; @@ -31,7 +31,7 @@ where input, events, configured_schema, - replay_budget, + &alias_budget, max_nesting_depth, ); source.enter_stream()?; @@ -69,13 +69,13 @@ where T: serde::Deserialize<'de>, { let configured_schema = options.selected_schema(); - let replay_budget = options.alias_expansion_budget(input.len()); + let alias_budget = Cell::new(options.alias_expansion_budget(input.len())); let max_nesting_depth = options.selected_max_nesting_depth(); Ok(EventDocumentIter { input, frames: EventDocumentFrames::from_str_with_options(input, options)?, configured_schema, - replay_budget, + alias_budget, max_nesting_depth, _marker: PhantomData, }) @@ -111,14 +111,14 @@ where ) })?; let configured_schema = options.selected_schema(); - let replay_budget = options.alias_expansion_budget(input.len()); + let alias_budget = Cell::new(options.alias_expansion_budget(input.len())); let max_nesting_depth = options.selected_max_nesting_depth(); let frames = EventDocumentFrames::from_str_with_options(&input, options)?; Ok(OwnedEventDocumentIter { input, frames, configured_schema, - replay_budget, + alias_budget, max_nesting_depth, _marker: PhantomData, }) @@ -128,7 +128,7 @@ pub(crate) struct EventDocumentIter<'de, T> { input: &'de str, frames: EventDocumentFrames, configured_schema: Schema, - replay_budget: usize, + alias_budget: Cell, max_nesting_depth: Option, _marker: PhantomData, } @@ -148,7 +148,7 @@ where self.input, events, self.configured_schema, - self.replay_budget, + &self.alias_budget, self.max_nesting_depth, ) }) @@ -161,7 +161,7 @@ pub(crate) struct OwnedEventDocumentIter { input: String, frames: EventDocumentFrames, configured_schema: Schema, - replay_budget: usize, + alias_budget: Cell, max_nesting_depth: Option, _marker: PhantomData, } @@ -181,7 +181,7 @@ where &self.input, events, self.configured_schema, - self.replay_budget, + &self.alias_budget, self.max_nesting_depth, ) }) diff --git a/src/event_de/serde_impl.rs b/src/event_de/serde_impl.rs index bb76439..c5bc203 100644 --- a/src/event_de/serde_impl.rs +++ b/src/event_de/serde_impl.rs @@ -2,11 +2,11 @@ use super::prepared::*; use super::source::EventSource; use super::*; -pub(super) struct EventNodeDeserializer<'a, 'de> { - pub(super) source: &'a mut EventSource<'de>, +pub(super) struct EventNodeDeserializer<'a, 'de, 'budget> { + pub(super) source: &'a mut EventSource<'de, 'budget>, } -impl<'de> EventNodeDeserializer<'_, 'de> { +impl<'de, 'budget> EventNodeDeserializer<'_, 'de, 'budget> { fn deserialize_prepared_current_node(self, visitor: V) -> Result where V: Visitor<'de>, @@ -38,7 +38,7 @@ impl<'de> EventNodeDeserializer<'_, 'de> { } } -impl<'de> de::Deserializer<'de> for EventNodeDeserializer<'_, 'de> { +impl<'de, 'budget> de::Deserializer<'de> for EventNodeDeserializer<'_, 'de, 'budget> { type Error = Error; fn deserialize_any(self, visitor: V) -> Result @@ -397,7 +397,7 @@ impl<'de> de::Deserializer<'de> for EventNodeDeserializer<'_, 'de> { } } -impl EventSource<'_> { +impl<'de, 'budget> EventSource<'de, 'budget> { fn peek_has_yaml_core_tag(&self, suffixes: &[&str]) -> bool { match self.peek() { Some(Event::SequenceStart { meta, .. }) | Some(Event::MappingStart { meta, .. }) => { @@ -424,12 +424,12 @@ impl EventSource<'_> { } } -struct EventSeqAccess<'a, 'de> { - source: &'a mut EventSource<'de>, +struct EventSeqAccess<'a, 'de, 'budget> { + source: &'a mut EventSource<'de, 'budget>, index: usize, } -impl<'de> SeqAccess<'de> for EventSeqAccess<'_, 'de> { +impl<'de, 'budget> SeqAccess<'de> for EventSeqAccess<'_, 'de, 'budget> { type Error = Error; fn next_element_seed(&mut self, seed: T) -> Result> @@ -450,12 +450,12 @@ impl<'de> SeqAccess<'de> for EventSeqAccess<'_, 'de> { } } -struct EventMapAccess<'a, 'de> { - source: &'a mut EventSource<'de>, +struct EventMapAccess<'a, 'de, 'budget> { + source: &'a mut EventSource<'de, 'budget>, value: Option, } -impl<'de> MapAccess<'de> for EventMapAccess<'_, 'de> { +impl<'de, 'budget> MapAccess<'de> for EventMapAccess<'_, 'de, 'budget> { type Error = Error; fn next_key_seed(&mut self, seed: K) -> Result> @@ -469,10 +469,9 @@ impl<'de> MapAccess<'de> for EventMapAccess<'_, 'de> { let depth = self.source.depth; let (events, pos) = self.source.current_events_and_pos(); let mut scan_anchors = self.source.anchors.clone(); - let mut replayed_events = 0usize; let segment = self .source - .mapping_key_at(events, pos, &mut scan_anchors, &mut replayed_events, depth)? + .mapping_key_at(events, pos, &mut scan_anchors, depth)? .map(|(node, _)| path_segment_for_node(&node)) .unwrap_or(ErrorPathSegment::ComplexKey); self.value = Some(segment.clone()); diff --git a/src/event_de/source.rs b/src/event_de/source.rs index 70c45d5..a622193 100644 --- a/src/event_de/source.rs +++ b/src/event_de/source.rs @@ -1,4 +1,5 @@ use super::*; +use crate::schema::AliasExpansionCost; pub(super) struct EventDocumentFrames { events: crate::parse::EventStream, @@ -98,7 +99,7 @@ pub(super) fn deserialize_document_frame<'de, T>( input: &'de str, events: Vec, configured_schema: Schema, - replay_budget: usize, + alias_budget: &Cell, max_nesting_depth: Option, ) -> Result where @@ -108,7 +109,7 @@ where input, events, configured_schema, - replay_budget, + alias_budget, max_nesting_depth, ); source.enter_stream()?; @@ -124,7 +125,7 @@ where } } -pub(super) struct EventSource<'de> { +pub(super) struct EventSource<'de, 'budget> { pub(super) input: &'de str, events: Vec, pos: usize, @@ -132,8 +133,7 @@ pub(super) struct EventSource<'de> { pub(super) schema: Schema, pub(super) anchors: HashMap>, inject: Vec, - replayed_events: usize, - replay_budget: usize, + alias_budget: &'budget Cell, max_nesting_depth: Option, pub(super) depth: usize, } @@ -144,12 +144,12 @@ struct InjectedEvents { pos: usize, } -impl<'de> EventSource<'de> { +impl<'de, 'budget> EventSource<'de, 'budget> { pub(super) fn new( input: &'de str, events: Vec, configured_schema: Schema, - replay_budget: usize, + alias_budget: &'budget Cell, max_nesting_depth: Option, ) -> Self { Self { @@ -160,13 +160,19 @@ impl<'de> EventSource<'de> { schema: configured_schema, anchors: HashMap::new(), inject: Vec::new(), - replayed_events: 0, - replay_budget, + alias_budget, max_nesting_depth, depth: 0, } } + fn charge_alias(&self, cost: AliasExpansionCost) -> bool { + let mut budget = self.alias_budget.get(); + let over_budget = budget.charge(cost); + self.alias_budget.set(budget); + over_budget + } + /// Records descent into a nested collection and enforces the configured /// nesting-depth ceiling. The event-backed path expands aliases lazily as /// it walks, so — unlike the tree-backed path's `AnchorTable::resolve` — the @@ -188,9 +194,9 @@ impl<'de> EventSource<'de> { self.depth = self.depth.saturating_sub(1); } - /// Same ceiling as [`enter_depth`], but for the read-only key/merge - /// materialization walk in [`node_at_for_key`], which threads an explicit - /// `depth` because it borrows `self` immutably. + /// Same ceiling as [`enter_depth`], but for the key/merge materialization + /// walk in [`node_at_for_key`], which threads an explicit `depth` through + /// its recursive calls. fn check_depth(&self, depth: usize, span: impl Into>) -> Result<()> { if self.max_nesting_depth.is_some_and(|max| depth > max) { return Err(Error::limit( @@ -279,15 +285,19 @@ impl<'de> EventSource<'de> { span, )); } - let events = self + let cost = self .anchors .get(&name) - .cloned() + .map(|events| event_expansion_cost(events)) .ok_or_else(|| Error::reference(format!("unknown anchor `{name}`"), span))?; - self.replayed_events = self.replayed_events.saturating_add(events.len()); - if self.replayed_events > self.replay_budget { + if self.charge_alias(cost) { return Err(Error::limit("alias event replay limit exceeded", span)); } + let events = self + .anchors + .get(&name) + .expect("anchor remains present while aliases are resolved") + .clone(); self.inject.push(InjectedEvents { anchor: name, events, @@ -308,7 +318,6 @@ impl<'de> EventSource<'de> { Event::DocumentStart { directives, .. } => { self.anchors.clear(); self.inject.clear(); - self.replayed_events = 0; self.depth = 0; self.schema = schema_for_directives(self.configured_schema, &directives); Ok(()) @@ -505,13 +514,11 @@ impl<'de> EventSource<'de> { pub(super) fn materialize_current_node_for_merge(&self) -> Result { let (events, pos) = self.current_events_and_pos(); let mut scan_anchors = self.anchors.clone(); - let mut replayed_events = 0usize; let (node, next) = self.node_at_for_key( events, pos, &mut scan_anchors, &mut Vec::new(), - &mut replayed_events, true, self.depth, )?; @@ -532,7 +539,6 @@ impl<'de> EventSource<'de> { }; let mut pos = start + 1; let mut scan_anchors = self.anchors.clone(); - let mut replayed_events = 0usize; while let Some(event) = events.get(pos) { if matches!(event, Event::MappingEnd { .. }) { return Ok(false); @@ -542,7 +548,6 @@ impl<'de> EventSource<'de> { pos, &mut scan_anchors, &mut Vec::new(), - &mut replayed_events, true, self.depth, )?; @@ -563,18 +568,13 @@ impl<'de> EventSource<'de> { let mut pos = start + 1; let mut seen = DuplicateKeyTracker::new(); let mut scan_anchors = self.anchors.clone(); - let mut replayed_events = 0usize; while let Some(event) = events.get(pos) { if matches!(event, Event::MappingEnd { .. }) { return Ok(()); } - if let Some((key, next_pos)) = self.mapping_key_at( - events, - pos, - &mut scan_anchors, - &mut replayed_events, - self.depth, - )? { + if let Some((key, next_pos)) = + self.mapping_key_at(events, pos, &mut scan_anchors, self.depth)? + { if node_is_merge_key(&key) { return Err(Error::data( "event-backed merge-key expansion is not implemented", @@ -605,7 +605,6 @@ impl<'de> EventSource<'de> { events: &[Event], pos: usize, scan_anchors: &mut HashMap>, - replayed_events: &mut usize, depth: usize, ) -> Result> { if let Some(name) = events.get(pos).and_then(event_anchor_name) { @@ -617,15 +616,7 @@ impl<'de> EventSource<'de> { | Some(Event::Alias { .. }) | Some(Event::SequenceStart { .. }) | Some(Event::MappingStart { .. }) => self - .node_at_for_key( - events, - pos, - scan_anchors, - &mut Vec::new(), - replayed_events, - false, - depth, - ) + .node_at_for_key(events, pos, scan_anchors, &mut Vec::new(), false, depth) .map(|(node, next)| Some((node, next))), Some(_) | None => Ok(None), } @@ -699,7 +690,6 @@ impl<'de> EventSource<'de> { pos: usize, scan_anchors: &mut HashMap>, active_aliases: &mut Vec, - replayed_events: &mut usize, allow_merge_key: bool, depth: usize, ) -> Result<(Node, usize)> { @@ -729,23 +719,28 @@ impl<'de> EventSource<'de> { anchor.span, )); } - let target = scan_anchors.get(name).cloned().ok_or_else(|| { - Error::reference(format!("unknown anchor `{name}`"), anchor.span) - })?; - *replayed_events = replayed_events.saturating_add(target.len()); - if *replayed_events > self.replay_budget { + let cost = scan_anchors + .get(name) + .map(|target| event_expansion_cost(target)) + .ok_or_else(|| { + Error::reference(format!("unknown anchor `{name}`"), anchor.span) + })?; + if self.charge_alias(cost) { return Err(Error::limit( "alias event replay limit exceeded", anchor.span, )); } + let target = scan_anchors + .get(name) + .expect("anchor remains present while aliases are resolved") + .clone(); active_aliases.push(name.clone()); let (mut node, end) = self.node_at_for_key( &target, 0, scan_anchors, active_aliases, - replayed_events, allow_merge_key, depth, )?; @@ -772,7 +767,6 @@ impl<'de> EventSource<'de> { next, scan_anchors, active_aliases, - replayed_events, allow_merge_key, depth + 1, )?; @@ -804,7 +798,6 @@ impl<'de> EventSource<'de> { next, scan_anchors, active_aliases, - replayed_events, allow_merge_key, depth + 1, )?; @@ -827,7 +820,6 @@ impl<'de> EventSource<'de> { after_key, scan_anchors, active_aliases, - replayed_events, allow_merge_key, depth + 1, )?; @@ -843,6 +835,84 @@ impl<'de> EventSource<'de> { } } +enum EventContainer { + Sequence, + Mapping { expect_key: bool }, +} + +fn event_expansion_cost(events: &[Event]) -> AliasExpansionCost { + let mut cost = AliasExpansionCost { + clone_work: events.len(), + ..AliasExpansionCost::default() + }; + let mut containers = Vec::new(); + + for event in events { + match event { + Event::Scalar { value, meta, .. } => { + count_container_item(&mut containers, &mut cost); + cost.scalar_bytes = cost.scalar_bytes.saturating_add(value.len()); + add_event_meta_bytes(&mut cost, meta); + } + Event::Alias { anchor } => { + count_container_item(&mut containers, &mut cost); + cost.scalar_bytes = cost.scalar_bytes.saturating_add(anchor.name.len()); + } + Event::SequenceStart { meta, .. } => { + count_container_item(&mut containers, &mut cost); + add_event_meta_bytes(&mut cost, meta); + containers.push(EventContainer::Sequence); + } + Event::MappingStart { meta, .. } => { + count_container_item(&mut containers, &mut cost); + add_event_meta_bytes(&mut cost, meta); + containers.push(EventContainer::Mapping { expect_key: true }); + } + Event::SequenceEnd { .. } => { + let _ = containers.pop(); + } + Event::MappingEnd { .. } => { + let _ = containers.pop(); + } + Event::StreamStart + | Event::StreamEnd + | Event::DocumentStart { .. } + | Event::DocumentEnd { .. } => {} + } + } + + cost +} + +fn count_container_item(containers: &mut Vec, cost: &mut AliasExpansionCost) { + let Some(container) = containers.last_mut() else { + return; + }; + match container { + EventContainer::Sequence => { + cost.container_items = cost.container_items.saturating_add(1); + } + EventContainer::Mapping { expect_key } => { + if *expect_key { + cost.container_items = cost.container_items.saturating_add(1); + } + *expect_key = !*expect_key; + } + } +} + +fn add_event_meta_bytes(cost: &mut AliasExpansionCost, meta: &EventMeta) { + if let Some(anchor) = &meta.anchor { + cost.scalar_bytes = cost.scalar_bytes.saturating_add(anchor.name.len()); + } + if let Some(tag) = &meta.tag { + cost.scalar_bytes = cost + .scalar_bytes + .saturating_add(tag.tag.handle.len()) + .saturating_add(tag.tag.suffix.len()); + } +} + pub(super) fn skip_node_in(events: &[Event], pos: usize) -> Result { match events .get(pos) diff --git a/src/event_de/tests.rs b/src/event_de/tests.rs index b7f260f..076edd7 100644 --- a/src/event_de/tests.rs +++ b/src/event_de/tests.rs @@ -688,6 +688,28 @@ fn event_deserializer_rejects_alias_replay_over_budget() { ); } +#[test] +fn event_document_iterator_aggregates_alias_replay_work_across_documents() { + let input = "---\nbase: &a []\ntarget: *a\n---\nbase: &a []\ntarget: *a\n"; + let options = LoadOptions::new().max_alias_expansion_nodes(3); + let mut documents = document_iter_str_with_options::(input, options) + .expect("event document iterator"); + + documents + .next() + .expect("first document") + .expect("first document is within the aggregate replay budget"); + let error = documents + .next() + .expect("second document") + .expect_err("event replay accounting must span documents"); + assert!( + error + .to_string() + .contains("alias event replay limit exceeded") + ); +} + #[test] fn event_deserializer_rejects_duplicate_keys_in_ignored_mappings() { let input = "base: &base {a: one, a: two}\ntarget: *base\n"; @@ -878,6 +900,7 @@ fn event_deserializer_document_errors_carry_document_index() { #[test] fn event_deserializer_skips_ignored_any_without_materializing_values() { let input = "root:\n - name: api\n ports: [80, 443]\n - nested:\n ok: true\n"; + let alias_budget = std::cell::Cell::new(LoadOptions::new().alias_expansion_budget(input.len())); IgnoredAny::deserialize(EventNodeDeserializer { source: &mut EventSource::new( input, @@ -886,7 +909,7 @@ fn event_deserializer_skips_ignored_any_without_materializing_values() { .collect::>>() .expect("events"), Schema::Yaml12, - LoadOptions::new().alias_expansion_budget(input.len()), + &alias_budget, LoadOptions::new().selected_max_nesting_depth(), ), }) diff --git a/src/parse.rs b/src/parse.rs index 9cbdbc8..43e47f5 100644 --- a/src/parse.rs +++ b/src/parse.rs @@ -16,7 +16,7 @@ use crate::{ de::read_to_end_with_options, error::utf8_error_span, key_identity::{DuplicateKeyTracker, check_duplicate_with_tracker_at_depth_limit}, - schema::{LoadOptions, Schema}, + schema::{AliasExpansionBudget, AliasExpansionCost, LoadOptions, Schema}, yaml11, }; use std::{ @@ -1142,17 +1142,15 @@ enum AnchorEntry { struct AnchorRegistry { entries: HashMap, generation: usize, - expanded_nodes: usize, - expansion_budget: usize, + expansion_budget: AliasExpansionBudget, options: LoadOptions, } impl AnchorRegistry { - fn new(expansion_budget: usize, options: LoadOptions) -> Self { + fn new(expansion_budget: AliasExpansionBudget, options: LoadOptions) -> Self { Self { entries: HashMap::new(), generation: 0, - expanded_nodes: 0, expansion_budget, options, } @@ -1160,7 +1158,6 @@ impl AnchorRegistry { fn reset_document(&mut self) { self.entries.clear(); - self.expanded_nodes = 0; } fn begin(&mut self, name: String, span: Span) -> usize { @@ -1203,9 +1200,7 @@ impl AnchorRegistry { } }; - let node_count = count_nodes(target); - self.expanded_nodes = self.expanded_nodes.saturating_add(node_count); - if self.expanded_nodes > self.expansion_budget { + if self.expansion_budget.charge(node_expansion_cost(target)) { return Err(Error::limit("alias expansion limit exceeded", span)); } if self @@ -4435,24 +4430,40 @@ fn core_scalar_tag_preserves_source(tag: &Tag) -> bool { .any(|suffix| tag.is_yaml_core(suffix)) } -fn count_nodes(node: &Node) -> usize { - let mut count = 0usize; +fn node_expansion_cost(node: &Node) -> AliasExpansionCost { + let mut cost = AliasExpansionCost::default(); let mut stack = vec![node]; while let Some(node) = stack.pop() { - count = count.saturating_add(1); + cost.clone_work = cost.clone_work.saturating_add(1); + if let Some(source) = node.scalar_source() { + cost.scalar_bytes = cost.scalar_bytes.saturating_add(source.raw().len()); + } match &node.value { - Value::Sequence(items) => stack.extend(items), + Value::Sequence(items) => { + cost.container_items = cost.container_items.saturating_add(items.len()); + stack.extend(items); + } Value::Mapping(entries) => { + cost.container_items = cost.container_items.saturating_add(entries.len()); for (key, value) in entries { stack.push(key); stack.push(value); } } - Value::Tagged(tagged) => stack.push(&tagged.value), - Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {} + Value::Tagged(tagged) => { + cost.scalar_bytes = cost + .scalar_bytes + .saturating_add(tagged.tag.handle.len()) + .saturating_add(tagged.tag.suffix.len()); + stack.push(&tagged.value); + } + Value::String(value) => { + cost.scalar_bytes = cost.scalar_bytes.saturating_add(value.len()); + } + Value::Null | Value::Bool(_) | Value::Number(_) => {} } } - count + cost } fn node_depth(node: &Node) -> usize { diff --git a/src/schema.rs b/src/schema.rs index 30b6e0a..52b679b 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -21,6 +21,54 @@ pub const DEFAULT_ALIAS_EXPANSION_FACTOR: usize = 64; /// Minimum alias expansion budget used by default loading options. pub const DEFAULT_MIN_ALIAS_EXPANSION_NODES: usize = 1024; +/// Aggregate accounting for one alias materialization. +/// +/// The dimensions are tracked independently because a small number of nodes +/// can still carry a large amount of scalar data, while a collection can do +/// substantial container work without carrying much scalar data. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) struct AliasExpansionCost { + pub(crate) clone_work: usize, + pub(crate) scalar_bytes: usize, + pub(crate) container_items: usize, +} + +/// Stream-level aggregate alias expansion budget. +/// +/// Every cost dimension is accumulated across all aliases and all documents in +/// the input stream. The configured limit is applied to each dimension so no +/// one kind of work can hide behind a different kind of cheap-looking node. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct AliasExpansionBudget { + limit: usize, + used: AliasExpansionCost, +} + +impl AliasExpansionBudget { + pub(crate) fn new(limit: usize) -> Self { + Self { + limit, + used: AliasExpansionCost::default(), + } + } + + /// Charges work before an alias target is materialized. + /// + /// Returning whether the aggregate is over budget lets callers construct + /// an error with the alias's source span without materializing the target. + pub(crate) fn charge(&mut self, cost: AliasExpansionCost) -> bool { + self.used.clone_work = self.used.clone_work.saturating_add(cost.clone_work); + self.used.scalar_bytes = self.used.scalar_bytes.saturating_add(cost.scalar_bytes); + self.used.container_items = self + .used + .container_items + .saturating_add(cost.container_items); + self.used.clone_work > self.limit + || self.used.scalar_bytes > self.limit + || self.used.container_items > self.limit + } +} + /// Default maximum constructed YAML nesting depth accepted by loading entrypoints. pub const DEFAULT_MAX_NESTING_DEPTH: usize = 128; @@ -170,16 +218,18 @@ impl LoadOptions { self.max_input_bytes } - /// Returns options with a maximum number of alias-expanded nodes. + /// Returns options with a maximum aggregate amount of alias expansion work. /// /// The default budget remains input-size derived. This option lets callers - /// loading untrusted configuration tighten that expansion work directly. + /// loading untrusted configuration tighten clone work, scalar bytes, and + /// collection-item work directly. The method name is retained for API + /// compatibility with the original node-count budget. pub const fn max_alias_expansion_nodes(mut self, max_alias_expansion_nodes: usize) -> Self { self.max_alias_expansion_nodes = Some(max_alias_expansion_nodes); self } - /// Returns the configured maximum number of alias-expanded nodes. + /// Returns the configured maximum aggregate alias expansion work. /// /// `None` means the default input-size-derived budget is selected. pub const fn selected_max_alias_expansion_nodes(self) -> Option { @@ -242,12 +292,12 @@ impl LoadOptions { self.max_collection_items } - pub(crate) fn alias_expansion_budget(self, input_len: usize) -> usize { - self.max_alias_expansion_nodes.unwrap_or_else(|| { + pub(crate) fn alias_expansion_budget(self, input_len: usize) -> AliasExpansionBudget { + AliasExpansionBudget::new(self.max_alias_expansion_nodes.unwrap_or_else(|| { input_len .saturating_mul(DEFAULT_ALIAS_EXPANSION_FACTOR) .max(DEFAULT_MIN_ALIAS_EXPANSION_NODES) - }) + })) } pub(crate) fn check_input_len(self, len: usize) -> Result<()> { diff --git a/tests/diagnostics.rs b/tests/diagnostics.rs index 70d1ec2..7d54b02 100644 --- a/tests/diagnostics.rs +++ b/tests/diagnostics.rs @@ -1435,18 +1435,18 @@ fn parser_depth_boundaries_cover_block_and_flow_shapes() { #[test] fn alias_expansion_boundary_keeps_raw_events_safe() { - let below = alias_expansion_chain(4); - parse_str(&below).expect("four-level alias expansion chain stays below budget"); + let below = alias_expansion_chain(3); + parse_str(&below).expect("three-level alias expansion chain stays below budget"); saneyaml::from_str::(&below).expect("serde reads below-budget alias chain"); assert!( saneyaml::parse_events(&below) .expect("raw events expose below-budget aliases") .iter() - .any(|event| matches!(event, saneyaml::Event::Alias { anchor } if anchor.name == "d")) + .any(|event| matches!(event, saneyaml::Event::Alias { anchor } if anchor.name == "c")) ); - let above = alias_expansion_chain(5); - let error = parse_str(&above).expect_err("five-level alias expansion chain crosses budget"); + let above = alias_expansion_chain(4); + let error = parse_str(&above).expect_err("four-level alias expansion chain crosses budget"); assert!(error.to_string().contains("alias expansion limit exceeded")); assert!(error.location().is_some()); let error = @@ -1457,10 +1457,47 @@ fn alias_expansion_boundary_keeps_raw_events_safe() { saneyaml::parse_events(&above) .expect("raw events remain safe because aliases are not expanded") .iter() - .any(|event| matches!(event, saneyaml::Event::Alias { anchor } if anchor.name == "e")) + .any(|event| matches!(event, saneyaml::Event::Alias { anchor } if anchor.name == "d")) ); } +#[test] +fn alias_expansion_budget_accounts_for_scalar_bytes_before_clone() { + let scalar = "value".repeat(40); + let input = format!("base: &base {scalar}\ntarget: *base\n"); + + let error = LoadOptions::new() + .max_alias_expansion_nodes(128) + .parse_str(&input) + .expect_err("scalar bytes must count against alias expansion work"); + assert!(error.to_string().contains("alias expansion limit exceeded")); + + let parsed: Value = LoadOptions::new() + .max_alias_expansion_nodes(512) + .from_str(&input) + .expect("a budget that covers the scalar clone preserves YAML semantics"); + assert_eq!(parsed["target"].as_str(), Some(scalar.as_str())); +} + +#[test] +fn alias_expansion_budget_aggregates_across_documents() { + let input = "---\nbase: &base {}\ntarget: *base\n---\nbase: &base {}\ntarget: *base\n"; + let options = LoadOptions::new().max_alias_expansion_nodes(1); + let mut documents = options + .stream_documents(input) + .expect("document stream construction"); + + documents + .next() + .expect("first document") + .expect("first document is within the aggregate budget"); + let error = documents + .next() + .expect("second document") + .expect_err("the stream must not reset alias accounting per document"); + assert!(error.to_string().contains("alias expansion limit exceeded")); +} + #[test] fn load_options_alias_expansion_limit_rejects_across_entrypoints() { let input = "base: &base {a: 1}\ntarget: *base\n";