fix: align fast paths with deserialize_bitmap on corrupt input

The fast path silently diverged from deserialize_bitmap:

1. Truncation: declared N buckets but fewer present was silently accepted;
  TreemapReader::new now walks all `size` headers up front.
2. Trailing data: bytes beyond declared buckets were decoded as undeclared buckets;
  `buf` is now sliced to the consumed range.
3. Mixed Small/Large paths in bitmap_has_any/has_all skipped roaring-side validation when small was empty;
   the reader is now built once before the loop and reused via the new TreemapReader::contains method (replaces reader::bitmap_contains).

Adds regression tests for all three cases.

Author: harry-hao

src/common/io/src/bitmap.rs

@@ -868,7 +868,7 @@ pub fn bitmap_contains(buf: &[u8], value: u64) -> Result<bool> {
     }
     validate_serialized_bitmap_header_and_length(buf)?;
     if let Some(roaring_buf) = as_roaring(buf) {
-        Ok(reader::bitmap_contains(roaring_buf, value)?)
+        Ok(reader::TreemapReader::new(roaring_buf)?.contains(value)?)
     } else {
         Ok(SmallReader::new(buf)?.contains(value))
     }
@@ -943,8 +943,9 @@ pub fn bitmap_has_all(lhs: &[u8], rhs: &[u8]) -> Result<bool> {
     // lhs is HybridLarge/Legacy, rhs is HybridSmall: probe
     let roaring_buf = as_roaring(lhs).unwrap();
     let rhs_small = SmallReader::new(rhs)?;
+    let tree = reader::TreemapReader::new(roaring_buf)?;
     for i in 0..rhs_small.len() {
-        if !reader::bitmap_contains(roaring_buf, rhs_small.values[i])? {
+        if !tree.contains(rhs_small.values[i])? {
             return Ok(false);
         }
     }
@@ -977,8 +978,9 @@ pub fn bitmap_has_any(lhs: &[u8], rhs: &[u8]) -> Result<bool> {
     } else {
         (as_roaring(rhs).unwrap(), SmallReader::new(lhs)?)
     };
+    let tree = reader::TreemapReader::new(roaring_buf)?;
     for i in 0..small.len() {
-        if reader::bitmap_contains(roaring_buf, small.values[i])? {
+        if tree.contains(small.values[i])? {
             return Ok(true);
         }
     }
@@ -1677,6 +1679,40 @@ mod tests {
         assert_eq!(decoded.into_iter().collect::<Vec<_>>(), vec![1, 5, 42]);
     }
 
+    // A corrupt large buffer paired with an empty HybridSmall must surface
+    // the decoding error, not silently return `false`/`true`. Without eager
+    // validation of the roaring side, the mixed-path loop runs zero times
+    // (small side is empty) and the corrupt large side slips through.
+    #[test]
+    fn has_any_corrupt_large_with_empty_small_is_rejected() {
+        let corrupt_large = 1u64.to_le_bytes(); // declares 1 bucket, no data
+        let empty_small = [
+            HYBRID_MAGIC[0],
+            HYBRID_MAGIC[1],
+            HYBRID_VERSION,
+            HYBRID_KIND_SMALL,
+            0, // 0 values
+        ];
+
+        assert!(bitmap_has_any(&corrupt_large, &empty_small).is_err());
+        assert!(bitmap_has_any(&empty_small, &corrupt_large).is_err());
+    }
+
+    #[test]
+    fn has_all_corrupt_large_with_empty_small_is_rejected() {
+        let corrupt_large = 1u64.to_le_bytes();
+        let empty_small = [
+            HYBRID_MAGIC[0],
+            HYBRID_MAGIC[1],
+            HYBRID_VERSION,
+            HYBRID_KIND_SMALL,
+            0,
+        ];
+
+        assert!(bitmap_has_all(&corrupt_large, &empty_small).is_err());
+        assert!(bitmap_has_all(&empty_small, &corrupt_large).is_err());
+    }
+
     // Tests for minimum deserialize bitmap functions.
     //
     // Cover public functions by comparing against RoaringTreemap:

src/common/io/src/bitmap/reader.rs

@@ -47,6 +47,17 @@ impl<'a> TreemapReader<'a> {
             .read_u64::<LittleEndian>()
             .map_err(|_| Error::other("fail to read size"))?;
 
+        // Validate that `size` buckets fit in `buf` and bound `buf` to the
+        // declared range, so the fast path matches `deserialize_bitmap` on
+        // corrupt input (truncation rejected, trailing data ignored).
+        let mut probe = buf;
+        for _ in 0..size {
+            let header = BitmapReader::decode(probe)?;
+            probe = &probe[header.buf.len()..];
+        }
+        let consumed = buf.len() - probe.len();
+        let buf = &buf[..consumed];
+
         Ok(Self { buf, _size: size })
     }
 
@@ -56,6 +67,26 @@ impl<'a> TreemapReader<'a> {
             offset: 0,
         }
     }
+
+    pub fn contains(&self, value: u64) -> io::Result<bool> {
+        let prefix = (value >> 32) as u32;
+        let container_key = ((value >> 16) & 0xFFFF) as u16;
+        let low16 = (value & 0xFFFF) as u16;
+
+        for bitmap_result in self.iter() {
+            let bitmap = bitmap_result?;
+            if bitmap.prefix() < prefix {
+                continue;
+            }
+            if bitmap.prefix() > prefix {
+                return Ok(false);
+            }
+            return Ok(bitmap
+                .find_container(container_key)?
+                .is_some_and(|c| c.contains(low16)));
+        }
+        Ok(false)
+    }
 }
 
 pub struct TreeMapIter<'a> {
@@ -853,27 +884,6 @@ pub(crate) fn bitmap_len_above(buf: &[u8], threshold: usize) -> io::Result<bool>
     Ok(false)
 }
 
-pub(crate) fn bitmap_contains(buf: &[u8], value: u64) -> io::Result<bool> {
-    let prefix = (value >> 32) as u32;
-    let container_key = ((value >> 16) & 0xFFFF) as u16;
-    let low16 = (value & 0xFFFF) as u16;
-
-    let tree = TreemapReader::new(buf)?;
-    for bitmap_result in tree.iter() {
-        let bitmap = bitmap_result?;
-        if bitmap.prefix() < prefix {
-            continue;
-        }
-        if bitmap.prefix() > prefix {
-            return Ok(false);
-        }
-        return Ok(bitmap
-            .find_container(container_key)?
-            .is_some_and(|c| c.contains(low16)));
-    }
-    Ok(false)
-}
-
 pub(crate) fn bitmap_min(buf: &[u8]) -> io::Result<Option<u64>> {
     let tree = TreemapReader::new(buf)?;
     for bitmap_result in tree.iter() {
@@ -992,6 +1002,51 @@ mod tests {
         Ok(())
     }
 
+    // A declared bucket count with no bucket data must be rejected up front.
+    #[test]
+    fn test_treemap_reader_rejects_truncated_buf() -> io::Result<()> {
+        let corrupt = 1u64.to_le_bytes();
+        assert!(TreemapReader::new(&corrupt).is_err());
+        Ok(())
+    }
+
+    // Partial truncation: the first bucket is intact and non-empty, but the
+    // declared count claims more buckets than exist. Eager validation in
+    // `TreemapReader::new` catches this before any caller can early-return.
+    #[test]
+    fn test_treemap_reader_rejects_partial_truncated_buf() -> io::Result<()> {
+        let mut tree = RoaringTreemap::new();
+        tree.insert(7);
+        let mut buf = Vec::new();
+        tree.serialize_into(&mut buf)?;
+
+        // Overwrite the leading u64 to claim 2 buckets while only 1 exists.
+        buf[..8].copy_from_slice(&2u64.to_le_bytes());
+
+        assert!(TreemapReader::new(&buf).is_err());
+        Ok(())
+    }
+
+    // Trailing data after the declared buckets must be ignored, matching
+    // `deserialize_bitmap`. A buffer that declares 0 buckets but carries a
+    // valid bucket trailer must decode as an empty treemap, not as one with
+    // an undeclared bucket.
+    #[test]
+    fn test_treemap_reader_ignores_trailing_data() -> io::Result<()> {
+        let mut tree = RoaringTreemap::new();
+        tree.insert(7);
+        let mut buf = Vec::new();
+        tree.serialize_into(&mut buf)?;
+
+        // Overwrite the leading u64 to claim 0 buckets; the rest is trailing.
+        buf[..8].copy_from_slice(&0u64.to_le_bytes());
+
+        let reader = TreemapReader::new(&buf)?;
+        let count = reader.iter().filter(|r| r.is_ok()).count();
+        assert_eq!(count, 0);
+        Ok(())
+    }
+
     #[test]
     fn test_intersection() -> io::Result<()> {
         let v1 = create_bitmap(123);
@@ -1283,10 +1338,11 @@ mod tests {
     #[test]
     fn test_bitmap_contains() -> io::Result<()> {
         for_each_fixture(|name, buf, tree, miss_value| {
+            let reader = TreemapReader::new(buf).unwrap();
             // Test a known-present value (min if exists)
             if let Some(hit) = tree.min() {
                 let expected = tree.contains(hit);
-                let actual = bitmap_contains(buf, hit).unwrap();
+                let actual = reader.contains(hit).unwrap();
                 assert_eq!(
                     actual, expected,
                     "bitmap_contains hit: fixture={name}, val={hit}"
@@ -1295,7 +1351,7 @@ mod tests {
             // Test a known-absent value
             let miss = miss_value;
             let expected = tree.contains(miss);
-            let actual = bitmap_contains(buf, miss).unwrap();
+            let actual = reader.contains(miss).unwrap();
             assert_eq!(
                 actual, expected,
                 "bitmap_contains miss: fixture={name}, val={miss}"