Harden RLE decompression against untrusted indices and offsets

Cargo.toml

@@ -154,7 +154,7 @@ datafusion-sqllogictest = { version = "54" }
 divan = { package = "codspeed-divan-compat", version = "5.0.0" }
 enum-iterator = "2.0.0"
 env_logger = "0.11"
-fastlanes = "0.5.1"
+fastlanes = "0.6.0"
 flatbuffers = "25.2.10"
 fsst-rs = "0.5.11"
 futures = { version = "0.3.31", default-features = false }

encodings/fastlanes/src/rle/array/rle_compress.rs

@@ -81,11 +81,15 @@ where
             // returned from `T::encode` are relative to the chunk.
             values_idx_offsets.push(value_count_acc as u64);
 
-            let value_count = NativeValue::<T>::encode(
-                input,
-                unsafe { &mut *(rle_vals.as_mut_ptr() as *mut [_; FL_CHUNK_SIZE]) },
-                rle_idxs,
-            );
+            // SAFETY: all three arguments are `FL_CHUNK_SIZE`-element arrays, which is
+            // all `encode_unchecked` requires.
+            let value_count = unsafe {
+                NativeValue::<T>::encode_unchecked(
+                    input,
+                    &mut *(rle_vals.as_mut_ptr() as *mut [_; FL_CHUNK_SIZE]),
+                    rle_idxs,
+                )
+            };
 
             value_count_acc += value_count;
         };

encodings/fastlanes/src/rle/array/rle_decompress.rs

@@ -12,6 +12,7 @@ use vortex_array::match_each_unsigned_integer_ptype;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
 use vortex_error::vortex_panic;
 
 use crate::FL_CHUNK_SIZE;
@@ -62,9 +63,27 @@ where
     let values = array.values().clone().execute::<PrimitiveArray>(ctx)?;
     let values = values.as_slice::<V>();
 
+    // The offsets come from (possibly untrusted) storage. Validate them once here so
+    // that the per-chunk slicing and the unchecked decodes below stay within `values`:
+    // offsets must be non-decreasing and span at most `values.len()` values.
+    vortex_ensure!(
+        values_idx_offsets.is_sorted(),
+        "RLE values_idx_offsets must be non-decreasing"
+    );
+    if let (Some(&first), Some(&last)) = (values_idx_offsets.first(), values_idx_offsets.last()) {
+        vortex_ensure!(
+            last - first <= values.len() as u64,
+            "RLE values_idx_offsets span {} values but only {} are present",
+            last - first,
+            values.len()
+        );
+    }
+
     let indices = array.indices().clone().execute::<PrimitiveArray>(ctx)?;
     assert!(indices.len().is_multiple_of(FL_CHUNK_SIZE));
-    let has_invalid = !indices.all_valid(ctx)?;
+    let indices_validity = indices.validity()?.execute_mask(indices.len(), ctx)?;
+    // `None` means every position is valid.
+    let validity_bits = (!indices_validity.all_true()).then(|| indices_validity.to_bit_buffer());
     let (indices_sl, _) = indices.as_slice::<I>().as_chunks::<FL_CHUNK_SIZE>();
 
     let chunk_start_idx = array.offset() / FL_CHUNK_SIZE;
@@ -89,6 +108,10 @@ where
         };
         let num_chunk_values = u16::try_from(next_value_idx_offset - value_idx_offset)
             .vortex_expect("There can be at most 1024 values in RLE chunk");
+        vortex_ensure!(
+            num_chunk_values > 0,
+            "RLE chunk {chunk_idx} references no values"
+        );
 
         // SAFETY: `MaybeUninit<T>` and `T` have the same layout.
         let buffer_values: &mut [V; FL_CHUNK_SIZE] = unsafe { std::mem::transmute(chunk_out) };
@@ -98,20 +121,47 @@ where
             // access. The indices may contain values other than 0 when they
             // have been further compressed (e.g., as a masked constant).
             buffer_values.fill(chunk_values[0]);
-        } else if has_invalid {
-            // When the indices array has invalid (null) positions, those
-            // positions may contain arbitrary garbage values after further
-            // compression. Clamp all indices into [0, num_chunk_values) to
-            // prevent out-of-bounds access in the fastlanes decoder.
-            let mut sanitized: [u16; FL_CHUNK_SIZE] = [0; FL_CHUNK_SIZE];
-            for (idx_out, idx) in sanitized.iter_mut().zip(chunk_indices) {
-                let idx: u16 =
-                    NumCast::from(*idx).vortex_expect("RLE indices are always less than u16");
-                *idx_out = idx.min(num_chunk_values - 1);
-            }
-            V::decode(chunk_values, &sanitized, buffer_values);
         } else {
-            V::decode(chunk_values, chunk_indices, buffer_values);
+            let chunk_start = chunk_idx * FL_CHUNK_SIZE;
+            let chunk_valid_count = validity_bits.as_ref().map_or(FL_CHUNK_SIZE, |bits| {
+                bits.count_range(chunk_start, chunk_start + FL_CHUNK_SIZE)
+            });
+            if chunk_valid_count == FL_CHUNK_SIZE {
+                decode_chunk_checked(
+                    chunk_values,
+                    chunk_indices,
+                    num_chunk_values,
+                    chunk_idx,
+                    buffer_values,
+                )?;
+            } else if chunk_valid_count == 0 {
+                // Entirely-null chunk: every position is masked out, so the indices are
+                // meaningless (possibly garbage after further compression) — skip the
+                // gather and emit an arbitrary in-bounds value.
+                buffer_values.fill(chunk_values[0]);
+            } else {
+                // Null positions may contain arbitrary garbage indices after further
+                // compression, so zero them out (their decoded values are masked by the
+                // validity). Valid positions must be genuinely in bounds and are still
+                // bound-checked before the unchecked gather: a corrupt index at a valid
+                // position is an error, never silently remapped.
+                let bits = validity_bits
+                    .as_ref()
+                    .vortex_expect("mixed-validity chunk implies a materialized mask");
+                let mut sanitized: [u16; FL_CHUNK_SIZE] = [0; FL_CHUNK_SIZE];
+                bits.slice(chunk_start..chunk_start + FL_CHUNK_SIZE)
+                    .for_each_set_index(|i| {
+                        sanitized[i] = NumCast::from(chunk_indices[i])
+                            .vortex_expect("RLE indices are always less than u16");
+                    });
+                decode_chunk_checked(
+                    chunk_values,
+                    &sanitized,
+                    num_chunk_values,
+                    chunk_idx,
+                    buffer_values,
+                )?;
+            }
         }
     }
 
@@ -128,3 +178,201 @@ where
         array.validity()?,
     ))
 }
+
+/// Bound-checks every index in the chunk, then runs the unchecked fastlanes gather.
+///
+/// The indices come from (possibly untrusted) storage: a single max-reduction over the
+/// chunk vectorizes and keeps the bounds check out of the per-element decode loop.
+fn decode_chunk_checked<V, I>(
+    chunk_values: &[V],
+    chunk_indices: &[I; FL_CHUNK_SIZE],
+    num_chunk_values: u16,
+    chunk_idx: usize,
+    out: &mut [V; FL_CHUNK_SIZE],
+) -> VortexResult<()>
+where
+    V: RLE,
+    I: Copy + Into<usize>,
+{
+    let max_index: usize = chunk_indices
+        .iter()
+        .map(|idx| (*idx).into())
+        .max()
+        .unwrap_or_default();
+    vortex_ensure!(
+        max_index < num_chunk_values as usize,
+        "RLE index {max_index} out of bounds for chunk {chunk_idx} with {num_chunk_values} values"
+    );
+    // SAFETY: just checked that every index in the chunk is below `num_chunk_values`,
+    // which the caller's offset validation bounds by `chunk_values.len()`.
+    unsafe { V::decode_unchecked(chunk_values, chunk_indices, out) };
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::IntoArray;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::assert_arrays_eq;
+    use vortex_array::validity::Validity;
+    use vortex_buffer::Buffer;
+    use vortex_error::VortexResult;
+
+    use crate::FL_CHUNK_SIZE;
+    use crate::RLE;
+    use crate::rle::array::rle_decompress::rle_decompress;
+    use crate::test::SESSION;
+
+    fn indices_with_oob(oob: u16) -> vortex_array::ArrayRef {
+        let mut indices = [0u16, 1]
+            .iter()
+            .cycle()
+            .take(FL_CHUNK_SIZE)
+            .copied()
+            .collect::<Vec<_>>();
+        indices[100] = oob;
+        PrimitiveArray::from_iter(indices).into_array()
+    }
+
+    #[test]
+    fn test_decode_rejects_out_of_bounds_index() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let values = PrimitiveArray::from_iter([10u32, 20]).into_array();
+        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
+
+        // Index 999 points far beyond the 2 values of the only chunk.
+        let rle = RLE::try_new(
+            values,
+            indices_with_oob(999),
+            values_idx_offsets,
+            0,
+            FL_CHUNK_SIZE,
+        )?;
+        assert!(rle_decompress(&rle, &mut ctx).is_err());
+        Ok(())
+    }
+
+    #[test]
+    fn test_decode_rejects_index_into_next_chunk() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        // Two chunks with 2 values each: an index of 2 in chunk 0 is within `values`
+        // but out of bounds for the chunk, and must be rejected rather than silently
+        // reading chunk 1's values.
+        let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array();
+        let values_idx_offsets = PrimitiveArray::from_iter([0u64, 2]).into_array();
+        let mut indices = [0u16, 1].repeat(FL_CHUNK_SIZE);
+        indices[100] = 2;
+        let indices = PrimitiveArray::from_iter(indices).into_array();
+
+        let rle = RLE::try_new(values, indices, values_idx_offsets, 0, 2 * FL_CHUNK_SIZE)?;
+        assert!(rle_decompress(&rle, &mut ctx).is_err());
+        Ok(())
+    }
+
+    #[test]
+    fn test_decode_rejects_non_monotonic_offsets() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array();
+        let values_idx_offsets = PrimitiveArray::from_iter([2u64, 0]).into_array();
+        let indices = PrimitiveArray::from_iter([0u16, 1].repeat(FL_CHUNK_SIZE)).into_array();
+
+        let rle = RLE::try_new(values, indices, values_idx_offsets, 0, 2 * FL_CHUNK_SIZE)?;
+        assert!(rle_decompress(&rle, &mut ctx).is_err());
+        Ok(())
+    }
+
+    #[test]
+    fn test_decode_rejects_offsets_beyond_values() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array();
+        // The offset span (100) exceeds the number of values present (4).
+        let values_idx_offsets = PrimitiveArray::from_iter([0u64, 100]).into_array();
+        let indices = PrimitiveArray::from_iter([0u16, 1].repeat(FL_CHUNK_SIZE)).into_array();
+
+        let rle = RLE::try_new(values, indices, values_idx_offsets, 0, 2 * FL_CHUNK_SIZE)?;
+        assert!(rle_decompress(&rle, &mut ctx).is_err());
+        Ok(())
+    }
+
+    #[test]
+    fn test_decode_rejects_empty_chunk() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let values = PrimitiveArray::from_iter([10u32, 20]).into_array();
+        // Chunk 0 references no values: offsets [0, 0].
+        let values_idx_offsets = PrimitiveArray::from_iter([0u64, 0]).into_array();
+        let indices = PrimitiveArray::from_iter([0u16, 1].repeat(FL_CHUNK_SIZE)).into_array();
+
+        let rle = RLE::try_new(values, indices, values_idx_offsets, 0, 2 * FL_CHUNK_SIZE)?;
+        assert!(rle_decompress(&rle, &mut ctx).is_err());
+        Ok(())
+    }
+
+    #[test]
+    fn test_decode_rejects_out_of_bounds_index_at_valid_position() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let values = PrimitiveArray::from_iter([10u32, 20]).into_array();
+        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
+
+        // Position 100 is VALID but holds an out-of-bounds index: this is corruption
+        // and must fail, even though other positions are null.
+        let mut indices = [0u16, 1].repeat(FL_CHUNK_SIZE / 2);
+        indices[100] = 999;
+        let mut validity = vec![true; FL_CHUNK_SIZE];
+        validity[200] = false;
+        let indices = PrimitiveArray::new(
+            indices.into_iter().collect::<Buffer<u16>>(),
+            Validity::from_iter(validity),
+        )
+        .into_array();
+
+        let rle = RLE::try_new(values, indices, values_idx_offsets, 0, FL_CHUNK_SIZE)?;
+        assert!(rle_decompress(&rle, &mut ctx).is_err());
+        Ok(())
+    }
+
+    #[test]
+    fn test_decode_tolerates_garbage_index_at_null_position() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let values = PrimitiveArray::from_iter([10u32, 20]).into_array();
+        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
+
+        // Position 100 is NULL and holds a garbage out-of-bounds index (e.g. after
+        // further compression of the indices): its value is masked, so decoding
+        // must succeed and every valid position must decode correctly.
+        let mut indices = [0u16, 1].repeat(FL_CHUNK_SIZE / 2);
+        indices[100] = 999;
+        let mut validity = vec![true; FL_CHUNK_SIZE];
+        validity[100] = false;
+        let indices = PrimitiveArray::new(
+            indices.into_iter().collect::<Buffer<u16>>(),
+            Validity::from_iter(validity),
+        )
+        .into_array();
+
+        let rle = RLE::try_new(values, indices, values_idx_offsets, 0, FL_CHUNK_SIZE)?;
+        let decoded = rle_decompress(&rle, &mut ctx)?;
+        let expected = PrimitiveArray::from_option_iter(
+            (0..FL_CHUNK_SIZE).map(|i| (i != 100).then_some(if i % 2 == 0 { 10u32 } else { 20 })),
+        );
+        assert_arrays_eq!(decoded, expected, &mut ctx);
+        Ok(())
+    }
+
+    #[test]
+    fn test_decode_in_bounds_indices_roundtrip() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let values = PrimitiveArray::from_iter([10u32, 20]).into_array();
+        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
+        let indices =
+            PrimitiveArray::from_iter([0u16, 1].iter().cycle().take(FL_CHUNK_SIZE).copied())
+                .into_array();
+
+        let rle = RLE::try_new(values, indices, values_idx_offsets, 0, FL_CHUNK_SIZE)?;
+        let decoded = rle_decompress(&rle, &mut ctx)?;
+        let expected =
+            PrimitiveArray::from_iter([10u32, 20].iter().cycle().take(FL_CHUNK_SIZE).copied());
+        assert_arrays_eq!(decoded, expected, &mut ctx);
+        Ok(())
+    }
+}