diff --git a/README.md b/README.md index 90fce53..5849875 100644 --- a/README.md +++ b/README.md @@ -274,7 +274,10 @@ them, so every input within the size cap gets their result and error. The pre-sc MessagePack themselves; it returns a `MsgpackStructureError` whose `Display` is the bare reason, with no prefix. `msgpack_document_len(bytes, max_depth)` runs the same walk and returns the document's length, so a caller that must reject -trailing bytes compares it with `bytes.len()` instead of walking twice. See +trailing bytes compares it with `bytes.len()` instead of walking twice. The walk +allocates nothing for a document that nests at most 8 levels deep, and for a +deeper one allocates only for the levels past 8; a rejection allocates its error +message. See [`SECURITY.md`](SECURITY.md#envelope-decode-bounds). diff --git a/benches/perf_ir_baselines.json b/benches/perf_ir_baselines.json index d58557f..39190b2 100644 --- a/benches/perf_ir_baselines.json +++ b/benches/perf_ir_baselines.json @@ -15,9 +15,9 @@ "keyring_decrypt/1024": 34739, "keyring_decrypt/64": 31754, "keyring_decrypt/65536": 264387, - "prescan/1024": 985, - "prescan/64": 883, - "prescan/65536": 935, + "prescan/1024": 910, + "prescan/64": 815, + "prescan/65536": 864, "retrieve/1024": 7409, "retrieve/1048576": 6679353, "retrieve/40": 2637, diff --git a/src/msgpack_bounds.rs b/src/msgpack_bounds.rs index 9d1a98a..0f000fd 100644 --- a/src/msgpack_bounds.rs +++ b/src/msgpack_bounds.rs @@ -27,7 +27,10 @@ fn reject(reason: impl Into) -> MsgpackStructureError { /// A header-only walk: str/bin/ext payloads are skipped by offset, never read, /// and nothing is decoded. It keeps one `u64` per open non-empty collection, so /// at most `min(max_depth, bytes.len())` of them (8 KiB at a depth of 1024), and -/// allocates nothing proportional to any declared length. +/// allocates nothing proportional to any declared length. The first 8 live on +/// the call stack, so a document that nests no deeper than that (an envelope +/// `ByteStorage::store` writes nests 2 deep) is walked with no allocation. A +/// rejection allocates its error message. /// /// Bytes after the document are neither read nor rejected. To find where the /// document ends, call [`msgpack_document_len`], which runs this same walk. @@ -76,6 +79,57 @@ pub fn check_msgpack_structure( msgpack_document_len(bytes, max_depth).map(|_| ()) } +/// Open collections the walk keeps on the call stack before it spills to the heap. +const INLINE_LEVELS: usize = 8; + +/// Elements still owed per open collection, innermost last; `len` is the +/// depth. The outermost [`INLINE_LEVELS`] live inline and deeper ones in +/// `spill`, so a shallow document never allocates. `spill` holds exactly the +/// levels past the inline ones. +struct OpenCollections { + inline: [u64; INLINE_LEVELS], + len: usize, + spill: Vec, +} + +impl OpenCollections { + fn new() -> Self { + Self { + inline: [0; INLINE_LEVELS], + len: 0, + spill: Vec::new(), + } + } + + fn len(&self) -> usize { + self.len + } + + fn last_mut(&mut self) -> Option<&mut u64> { + match self.len { + 0 => None, + n @ 1..=INLINE_LEVELS => Some(&mut self.inline[n - 1]), + _ => self.spill.last_mut(), + } + } + + fn push(&mut self, owed: u64) { + if self.len < INLINE_LEVELS { + self.inline[self.len] = owed; + } else { + self.spill.push(owed); + } + self.len += 1; + } + + fn pop(&mut self) { + if self.len > INLINE_LEVELS { + self.spill.pop(); + } + self.len -= 1; + } +} + /// Runs the walk [`check_msgpack_structure`] describes and returns the length, /// in bytes, of the document at the start of `bytes`. /// @@ -119,9 +173,9 @@ pub fn msgpack_document_len( let mut pos = 0usize; let mut pending: u64 = 1; // elements owed across all open collections (the root is one) - let mut open: Vec = Vec::new(); // elements still owed per open collection = depth + let mut open = OpenCollections::new(); while pending > 0 { - while open.last() == Some(&0) { + while open.last_mut() == Some(&mut 0) { open.pop(); } let marker = *bytes @@ -319,6 +373,52 @@ mod tests { } } + #[test] + fn depth_bound_is_exact_where_the_stack_spills_to_the_heap() { + for max_depth in INLINE_LEVELS - 1..=INLINE_LEVELS + 1 { + let too_deep = Err(format!("nests deeper than {max_depth} levels")); + for (header, empty) in LEVELS { + let at = |levels, tail: &[u8]| { + check_msgpack_structure(&nested(header, levels, tail), max_depth) + .map_err(|e| e.to_string()) + }; + assert_eq!(at(max_depth, &[0xc0]), Ok(()), "max_depth {max_depth}"); + assert_eq!( + at(max_depth + 1, &[0xc0]), + too_deep, + "max_depth {max_depth}" + ); + assert_eq!(at(max_depth - 1, &[empty]), Ok(()), "max_depth {max_depth}"); + assert_eq!(at(max_depth, &[empty]), too_deep, "max_depth {max_depth}"); + } + } + } + + #[test] + fn a_sibling_after_a_spilled_subtree_is_measured_from_its_own_parent() { + // [a levels of [..[nil]..], b levels of the same]: depth is 1 + max(a, b). + // An `a` past the inline slots spills the stack to the heap and drains + // it again before `b` starts. + let max_depth = INLINE_LEVELS + 2; + for a in 0..=max_depth { + for b in 0..=max_depth { + let doc = [ + vec![0x92], + nested(&[0x91], a, &[0xc0]), + nested(&[0x91], b, &[0xc0]), + ] + .concat(); + let expected = if a.max(b) < max_depth { + Ok(doc.len()) + } else { + Err(format!("nests deeper than {max_depth} levels")) + }; + let got = msgpack_document_len(&doc, max_depth).map_err(|e| e.to_string()); + assert_eq!(got, expected, "a {a}, b {b}"); + } + } + } + #[test] fn map_pairs_cost_two_slots() { assert_eq!(check(&[0x81, 0xa0, 0xc0]), Ok(())); diff --git a/src/read_allocation_probe.rs b/src/read_allocation_probe.rs index ee91acd..7577430 100644 --- a/src/read_allocation_probe.rs +++ b/src/read_allocation_probe.rs @@ -25,7 +25,10 @@ //! //! The same count, [`bytes_requested`], also pins the write side: the //! `byte_storage` unit test `store_allocates_only_its_output_and_block` checks -//! that `ByteStorage::store` requests only its output and its LZ4 block. +//! that `ByteStorage::store` requests only its output and its LZ4 block, and +//! `accepted_envelope_prescan_allocates_nothing` below checks that the +//! structural pre-scan, which the fallback decode and SDK payload walks run, +//! requests nothing for an envelope a writer emits. //! //! This lives in the crate's unit tests, compiled only under `cfg(test)`, //! because the positive control needs a hook inside `extract_payload` @@ -151,6 +154,44 @@ mod tests { } } + /// The structural pre-scan keeps its nesting stack on the call stack for + /// an envelope, current or legacy, at every payload size, so walking one + /// requests nothing from the allocator. `retrieve` runs the walk only in + /// its fallback decode (`decode_envelope_full`); SDK bindings run it over + /// every payload they unpack. + #[test] + fn accepted_envelope_prescan_allocates_nothing() { + // The envelope decode's bound (`decode_envelope_full`). + const MAX_DEPTH: usize = 100; + let storage = ByteStorage::new(None); + for len in [0, 40, 64, 1024, 64 * 1024] { + let data: Vec = (0..len).map(|i| (i * 31 % 251) as u8).collect(); + let envelope = storage.store(&data, None).unwrap(); + let e: crate::StorageEnvelope = rmp_serde::from_slice(&envelope).unwrap(); + // Legacy writers encoded `compressed_data` as an array of ints. + let legacy = + rmp_serde::to_vec(&(&e.compressed_data, e.checksum, e.original_size, &e.format)) + .unwrap(); + for (shape, doc) in [("current", envelope), ("legacy", legacy)] { + assert!( + storage.retrieve(&doc).is_ok(), + "[{shape} {len}] not accepted" + ); + let (result, requested) = + bytes_requested(|| crate::check_msgpack_structure(&doc, MAX_DEPTH)); + assert_eq!(result, Ok(()), "[{shape} {len}]"); + assert_eq!(requested, 0, "[{shape} {len}] pre-scan allocated"); + } + } + // Positive control: past the inline levels the stack spills to the + // heap, and this probe sees it. + let deep = [vec![0x91; 9], vec![0xc0]].concat(); + let (result, requested) = + bytes_requested(|| crate::check_msgpack_structure(&deep, MAX_DEPTH)); + assert_eq!(result, Ok(())); + assert!(requested > 0, "probe missed the spilled nesting stack"); + } + /// On every `msgpack bin` vector, `retrieve_data` allocates what the LZ4 /// decoder does on its own and nothing else, since the one-pass decode /// takes every such vector without allocating, and `retrieve` adds only