diff --git a/crates/format/src/unpack.rs b/crates/format/src/unpack.rs index 3199bbc5..d38f1776 100644 --- a/crates/format/src/unpack.rs +++ b/crates/format/src/unpack.rs @@ -117,8 +117,10 @@ impl<'a> Unpacked<'a> { /// The entries `[offset, length]` of the data after the JSON → each one's bytes, `None` for /// one that isn't there (the data ends before it). Errors: the compressed data is damaged, or - /// the entries are larger than [`MAX_BLOB`] each or [`MAX_BLOBS`] together. Entries are read - /// in the order of their offsets, each straight into its own buffer. + /// the entries of a compressed file are larger than [`MAX_BLOB`] each or [`MAX_BLOBS`] + /// together, or those of an uncompressed file larger together than both [`MAX_BLOBS`] and the + /// data (entries sharing bytes, only in hand-made files, would copy them once each). Entries + /// are read in the order of their offsets, each straight into its own buffer. pub(crate) fn entries(self, wanted: &[[u64; 2]]) -> Result>>, FormatError> { let mut out: Vec>> = vec![None; wanted.len()]; let in_memory = |tail: &[u8], out: &mut Vec>>| { @@ -127,8 +129,12 @@ impl<'a> Unpacked<'a> { } }; match self.tail { - Tail::Slice(tail) => in_memory(tail, &mut out), + Tail::Slice(tail) => { + within(wanted, u64::MAX, MAX_BLOBS.max(tail.len() as u64))?; + in_memory(tail, &mut out); + } Tail::Whole(at) => { + within(wanted, MAX_BLOB, MAX_BLOBS)?; let buf = match &self.json { Json::Whole { buf, .. } => buf.as_slice(), _ => &[], @@ -137,16 +143,7 @@ impl<'a> Unpacked<'a> { } Tail::None => in_memory(&[], &mut out), Tail::Stream { file, at, read, rest } => { - let mut total = 0u64; - for &[_, len] in wanted { - if len > MAX_BLOB { - return Err(FormatError::TooBig); - } - total = total.saturating_add(len); - } - if total > MAX_BLOBS { - return Err(FormatError::TooBig); - } + within(wanted, MAX_BLOB, MAX_BLOBS)?; drop(self.json); let mut order: Vec = (0..wanted.len()).collect(); order.sort_by_key(|&i| wanted.get(i).map_or(0, |w| w[0])); @@ -165,6 +162,18 @@ impl<'a> Unpacked<'a> { } } +/// Are the entries at most `each` bytes long each and `all` together? Checked before any is read. +fn within(wanted: &[[u64; 2]], each: u64, all: u64) -> Result<(), FormatError> { + let mut total = 0u64; + for &[_, len] in wanted { + if len > each { + return Err(FormatError::TooBig); + } + total = total.saturating_add(len); + } + if total > all { Err(FormatError::TooBig) } else { Ok(()) } +} + /// `offset..offset + len` as indices, when they fit. fn range(offset: u64, len: u64) -> Option<(usize, usize)> { let a = usize::try_from(offset).ok()?; @@ -215,3 +224,33 @@ impl Stream<'_> { Ok((got as u64 == len).then_some(bytes)) } } + +#[cfg(test)] +mod tests { + use std::io::Write; + + use super::*; + + fn gzip(bytes: &[u8]) -> Vec { + let mut gz = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + gz.write_all(bytes).unwrap(); + gz.finish().unwrap() + } + + #[test] + fn entries_sharing_bytes_obey_the_limits_in_memory_too() { + // Each entry within MAX_BLOB, together past MAX_BLOBS: many references to the same bytes. + let many = vec![[0, MAX_BLOB]; 5]; + // Uncompressed. + assert!(matches!(Unpacked::new(b"{}").unwrap().entries(&many), Err(FormatError::TooBig))); + // Compressed, unpacked whole (no marker right after the JSON). + let whole = gzip(b"{}\0x"); + assert!(matches!(Unpacked::new(&whole).unwrap().entries(&many), Err(FormatError::TooBig))); + assert!(matches!(Unpacked::new(&whole).unwrap().entries(&[[0, MAX_BLOB + 1]]), Err(FormatError::TooBig))); + // Uncompressed files keep no per-entry limit: a length the data doesn't have is just missing. + assert_eq!(Unpacked::new(b"{}").unwrap().entries(&[[0, MAX_BLOB + 1]]).unwrap(), vec![None]); + // Within the limits, shared bytes are still read once per entry. + let file = b"{}\n\0VCBLOBS\0abcd"; + assert_eq!(Unpacked::new(file).unwrap().entries(&[[0, 4], [0, 4]]).unwrap(), vec![Some(b"abcd".to_vec()), Some(b"abcd".to_vec())]); + } +}