Merge pull request #1783 from MarekSpodBiedry/issue-1698

Fix #1698: entries read from memory obey the size limits too, not only streamed ones (#1698)

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<Vec<Option<Vec<u8>>>, FormatError> {
         let mut out: Vec<Option<Vec<u8>>> = vec![None; wanted.len()];
         let in_memory = |tail: &[u8], out: &mut Vec<Option<Vec<u8>>>| {
@@ -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<usize> = (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<u8> {
+        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())]);
+    }
+}