diff --git a/rust/lance-encoding/src/data.rs b/rust/lance-encoding/src/data.rs index 2f3a686ca89..03d2415340c 100644 --- a/rust/lance-encoding/src/data.rs +++ b/rust/lance-encoding/src/data.rs @@ -588,20 +588,253 @@ pub struct VariableWidthBlock { pub block_info: BlockInfo, } +/// Proof that a [`VariableWidthBlock`] satisfies the Arrow layout contract for +/// its target data type (offsets buffer long enough, offsets monotonic and +/// within the data buffer, values valid UTF-8 where required). +/// +/// Only [`VariableWidthBlock::validate_layout`] can construct it, which ties the +/// unchecked Arrow build below to an actual validation pass instead of a +/// caller-controlled flag. +struct ValidVariableWidthLayout; + +fn corrupt_file_named(name: &str, message: impl Into) -> Error { + Error::corrupt_file(name.into(), message) +} + impl VariableWidthBlock { - fn into_arrow(self, data_type: DataType, validate: bool) -> Result { + // The offsets buffer comes straight from file bytes, so an unchecked build would + // let a corrupt file smuggle out-of-bounds offsets into an Arrow array whose + // consumers then read (or crash on) memory outside the data buffer. This + // boundary therefore always validates the layout, ignoring the optional + // `validate` flag. Lance validates the common layouts itself (a branchless + // scan, measurably cheaper than Arrow's element-wise checked build) and only + // falls back to Arrow's checked build for the cold cases. + fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + let Some(expected_bits_per_offset) = Self::expected_bits_per_offset(&data_type) else { + // Not an [offsets, bytes] layout we know how to prove; let Arrow + // check it. + return self.into_arrow_checked(data_type); + }; + if self.bits_per_offset != expected_bits_per_offset { + return Err(self.layout_error( + &data_type, + format!( + "expected {}-bit offsets but got {}-bit offsets", + expected_bits_per_offset, self.bits_per_offset + ), + )); + } + if self.num_values == 0 { + // Cold path; Arrow handles the empty-offsets special cases. + return self.into_arrow_checked(data_type); + } + let proof = self.validate_layout(&data_type)?; + Ok(self.into_arrow_unchecked(data_type, proof)) + } + + /// The offset width Arrow mandates for `data_type`, or `None` if the type + /// does not use the `[offsets, bytes]` layout this block represents. + fn expected_bits_per_offset(data_type: &DataType) -> Option { + match data_type { + DataType::Binary | DataType::Utf8 => Some(32), + DataType::LargeBinary | DataType::LargeUtf8 => Some(64), + _ => None, + } + } + + fn layout_error(&self, data_type: &DataType, detail: impl std::fmt::Display) -> Error { + Self::format_layout_error( + data_type, + detail, + self.num_values, + self.bits_per_offset, + self.offsets.len(), + self.data.len(), + ) + } + + fn format_layout_error( + data_type: &DataType, + detail: impl std::fmt::Display, + num_values: u64, + bits_per_offset: u8, + offsets_size: usize, + data_size: usize, + ) -> Error { + corrupt_file_named( + "variable width data block", + format!( + "invalid variable-width layout for {}: {} (num_values: {}, bits_per_offset: {}, \ + offsets buffer size: {} bytes, data buffer size: {} bytes)", + data_type, detail, num_values, bits_per_offset, offsets_size, data_size, + ), + ) + } + + fn validate_layout(&self, data_type: &DataType) -> Result { + let bytes_per_offset = (self.bits_per_offset / 8) as u64; + let required_bytes = self + .num_values + .checked_add(1) + .and_then(|num_offsets| num_offsets.checked_mul(bytes_per_offset)) + .ok_or_else(|| self.layout_error(data_type, "offsets buffer size overflows"))?; + if (self.offsets.len() as u64) < required_bytes { + return Err(self.layout_error( + data_type, + format!( + "offsets buffer must hold at least {} offsets ({} bytes)", + self.num_values + 1, + required_bytes + ), + )); + } + let validate_utf8 = matches!(data_type, DataType::Utf8 | DataType::LargeUtf8); + match self.bits_per_offset { + 32 => self.validate_offsets_and_values::(data_type, validate_utf8), + 64 => self.validate_offsets_and_values::(data_type, validate_utf8), + other => Err(self.layout_error( + data_type, + format!("unsupported offset width: {} bits", other), + )), + } + } + + fn validate_offsets_and_values( + &self, + data_type: &DataType, + validate_utf8: bool, + ) -> Result { + let num_offsets = self.num_values as usize + 1; + // Slice before borrowing: the buffer may carry padding that is not a + // multiple of the offset width. + let offsets = self + .offsets + .slice_with_length(0, num_offsets * std::mem::size_of::()); + let offsets = offsets.borrow_to_typed_slice::(); + let offsets: &[T] = offsets.as_ref(); + let data = self.data.as_ref(); + + // A monotonic sequence with a non-negative first offset and an + // in-bounds last offset is entirely within [0, data.len()], so the hot + // loop only proves monotonicity; everything else is O(1) at the ends. + // The `&=` accumulation keeps the loop branchless so it vectorizes. + let mut is_monotonic = true; + for window in offsets.windows(2) { + is_monotonic &= window[0] <= window[1]; + } + let first = offsets[0]; + let last = offsets[num_offsets - 1]; + let bounds_ok = + first >= T::usize_as(0) && last.to_usize().is_some_and(|last| last <= data.len()); + if !is_monotonic || !bounds_ok { + return Err(self.offset_violation_error::(data_type, offsets)); + } + + if validate_utf8 { + let (first, last) = (first.as_usize(), last.as_usize()); + let values = std::str::from_utf8(&data[first..last]) + .map_err(|utf8_err| self.layout_error(data_type, utf8_err))?; + let mut on_char_boundaries = true; + for &offset in offsets { + on_char_boundaries &= values.is_char_boundary(offset.as_usize() - first); + } + if !on_char_boundaries { + // Cold path: rescan to pinpoint the offending offset. + let position = offsets + .iter() + .position(|offset| !values.is_char_boundary(offset.as_usize() - first)) + .expect("the fast scan found a non-boundary offset"); + return Err(self.layout_error( + data_type, + format!("offset at position {position} splits a UTF-8 character"), + )); + } + } + + Ok(ValidVariableWidthLayout) + } + + /// Cold path: pinpoint the first offending offset for the error message. + fn offset_violation_error( + &self, + data_type: &DataType, + offsets: &[T], + ) -> Error { + let data_size = self.data.len(); + for (position, window) in offsets.windows(2).enumerate() { + if window[0] > window[1] { + return self.layout_error( + data_type, + format!( + "non-monotonic offset at position {}: {:?} > {:?}", + position, window[0], window[1] + ), + ); + } + } + for (position, offset) in offsets.iter().enumerate() { + match offset.to_usize() { + None => { + return self.layout_error( + data_type, + format!("negative offset at position {}: {:?}", position, offset), + ); + } + Some(offset) if offset > data_size => { + return self.layout_error( + data_type, + format!( + "offset at position {} out of bounds: {} > {}", + position, offset, data_size + ), + ); + } + Some(_) => {} + } + } + // The fast scan only fails when one of the loops above finds the + // culprit; reaching here would be a bug in the fast scan itself. + self.layout_error(data_type, "offsets failed validation") + } + + fn into_arrow_checked(self, data_type: DataType) -> Result { + let num_values = self.num_values; + let bits_per_offset = self.bits_per_offset; + let offsets_size = self.offsets.len(); + let data_size = self.data.len(); + let builder = self.into_arrow_builder(data_type.clone()); + builder.build().map_err(|arrow_err| { + Self::format_layout_error( + &data_type, + arrow_err, + num_values, + bits_per_offset, + offsets_size, + data_size, + ) + }) + } + + fn into_arrow_unchecked( + self, + data_type: DataType, + _proof: ValidVariableWidthLayout, + ) -> ArrayData { + let builder = self.into_arrow_builder(data_type); + // SAFETY: `_proof` witnesses that `validate_layout` proved this block + // satisfies the Arrow layout contract for `data_type`. + unsafe { builder.build_unchecked() } + } + + fn into_arrow_builder(self, data_type: DataType) -> ArrayDataBuilder { + let num_values = self.num_values; let data_buffer = self.data.into_buffer(); let offsets_buffer = self.offsets.into_buffer(); - let builder = ArrayDataBuilder::new(data_type) + ArrayDataBuilder::new(data_type) .add_buffer(offsets_buffer) .add_buffer(data_buffer) - .len(self.num_values as usize) - .null_count(0); - if validate { - Ok(builder.build()?) - } else { - Ok(unsafe { builder.build_unchecked() }) - } + .len(num_values as usize) + .null_count(0) } fn into_buffers(self) -> Vec { @@ -1634,19 +1867,24 @@ mod tests { use std::sync::Arc; use arrow_array::{ - ArrayRef, DictionaryArray, Int8Array, LargeBinaryArray, StringArray, UInt8Array, - UInt16Array, make_array, new_null_array, + ArrayRef, BinaryArray, DictionaryArray, Int8Array, LargeBinaryArray, LargeStringArray, + StringArray, UInt8Array, UInt16Array, make_array, new_null_array, types::{Int8Type, Int32Type}, }; use arrow_buffer::{BooleanBuffer, NullBuffer}; use arrow_schema::{DataType, Field, Fields}; + use lance_core::Error; use lance_datagen::{ArrayGeneratorExt, DEFAULT_SEED, RowCount, array}; use rand::SeedableRng; + use rstest::rstest; use crate::buffer::LanceBuffer; - use super::{AllNullDataBlock, DataBlock}; + use super::{ + AllNullDataBlock, BlockInfo, DataBlock, DictionaryDataBlock, FixedWidthDataBlock, + VariableWidthBlock, + }; use arrow_array::Array; @@ -2009,4 +2247,184 @@ mod tests { let total_nulls_size_in_bytes = concatenated_array.nulls().unwrap().len().div_ceil(8); assert!(block.data_size() == (total_buffer_size + total_nulls_size_in_bytes) as u64); } + + #[test] + fn variable_width_rejects_out_of_bounds_offsets_without_optional_validation() { + let block = VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets: LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + bits_per_offset: 32, + num_values: 3, + block_info: BlockInfo::new(), + }; + + let error = block + .into_arrow(DataType::Binary, false) + .expect_err("out-of-bounds offsets must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("100000") && message.contains("data buffer size: 14 bytes"), + "error must report the offending offset and the data buffer size: {message}" + ); + } + + #[rstest] + #[case::binary_i32_tail_out_of_bounds( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::utf8_i32_tail_out_of_bounds( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::large_binary_i64_tail_out_of_bounds( + DataType::LargeBinary, + LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 100_000]), + 64, + 3, + b"alphabetagamma".as_slice() + )] + #[case::large_utf8_i64_tail_out_of_bounds( + DataType::LargeUtf8, + LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 100_000]), + 64, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_negative_offset( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, -1, 9, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_non_monotonic_offsets( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 9, 5, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_interior_offset_out_of_bounds( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 100_000, 100_000, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_offsets_buffer_too_short( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::utf8_invalid_byte_sequence( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2, 3]), + 32, + 3, + &[b'a', 0xFF, b'b'] + )] + #[case::utf8_offset_splits_multibyte_char( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2]), + 32, + 2, + "é".as_bytes() + )] + #[case::large_utf8_invalid_byte_sequence( + DataType::LargeUtf8, + LanceBuffer::reinterpret_vec(vec![0_i64, 1, 2, 3]), + 64, + 3, + &[b'a', 0xFF, b'b'] + )] + fn variable_width_rejects_malformed_layout( + #[case] data_type: DataType, + #[case] offsets: LanceBuffer, + #[case] bits_per_offset: u8, + #[case] num_values: u64, + #[case] data: &[u8], + ) { + let block = VariableWidthBlock { + data: LanceBuffer::copy_slice(data), + offsets, + bits_per_offset, + num_values, + block_info: BlockInfo::new(), + }; + + // The malformed layout must be rejected regardless of the optional + // `validate` flag: the flag selects extra validation, not the memory + // safety proof required to construct an Arrow array. + for validate in [false, true] { + let error = DataBlock::VariableWidth(block.clone()) + .into_arrow(data_type.clone(), validate) + .expect_err("malformed variable-width layout must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile with validate={validate}, got: {error}" + ); + } + } + + #[test] + fn dictionary_rejects_malformed_variable_width_values_without_optional_validation() { + let values = VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets: LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + bits_per_offset: 32, + num_values: 3, + block_info: BlockInfo::new(), + }; + let dictionary = DataBlock::Dictionary(DictionaryDataBlock { + indices: FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2]), + bits_per_value: 32, + num_values: 3, + block_info: BlockInfo::new(), + }, + dictionary: Box::new(DataBlock::VariableWidth(values)), + }); + + let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)); + let error = dictionary + .into_arrow(data_type, false) + .expect_err("dictionary with out-of-bounds value offsets must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + } + + #[rstest] + #[case::binary(Arc::new(BinaryArray::from_vec(vec![b"alpha", b"", b"gamma"])) as ArrayRef)] + #[case::large_binary( + Arc::new(LargeBinaryArray::from_vec(vec![b"alpha", b"", b"gamma"])) as ArrayRef + )] + #[case::utf8(Arc::new(StringArray::from(vec!["héllo", "", "world"])) as ArrayRef)] + #[case::large_utf8(Arc::new(LargeStringArray::from(vec!["héllo", "", "world"])) as ArrayRef)] + fn variable_width_valid_data_survives_mandatory_validation(#[case] array: ArrayRef) { + let block = DataBlock::from_array(array.clone()); + for validate in [false, true] { + let round_tripped = make_array( + block + .clone() + .into_arrow(array.data_type().clone(), validate) + .unwrap(), + ); + assert_eq!(&round_tripped, &array); + } + } } diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index aa6e0744191..e3b4242bc0d 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -2295,7 +2295,7 @@ impl FullZipScheduler { num_rows, bits_per_offset, bits_per_offset, - ))) + )?)) } } } @@ -2695,6 +2695,10 @@ struct VariableFullZipDecoder { num_rows: u64, } +fn corrupt_file_named(name: &str, message: impl Into) -> Error { + Error::corrupt_file(name.into(), message) +} + impl VariableFullZipDecoder { fn new( details: Arc, @@ -2702,7 +2706,7 @@ impl VariableFullZipDecoder { num_rows: u64, in_bits_per_length: u8, out_bits_per_offset: u8, - ) -> Self { + ) -> Result { let decompressor = match details.value_decompressor { PerValueDecompressor::Variable(ref d) => d.clone(), _ => unreachable!(), @@ -2747,9 +2751,9 @@ impl VariableFullZipDecoder { // - We could force each decode task to do a full unzip of all the data. Each decode task now // has to do more work but the work is all fused. // - We could just try doing this work on the decode thread and see if it is a problem. - decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows); + decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?; - decoder + Ok(decoder) } fn slice_batch_data_and_rebase_offsets_typed( @@ -2824,28 +2828,33 @@ impl VariableFullZipDecoder { } } - unsafe fn parse_length(data: &[u8], bits_per_offset: u8) -> u64 { - match bits_per_offset { - 8 => *data.get_unchecked(0) as u64, - 16 => u16::from_le_bytes([*data.get_unchecked(0), *data.get_unchecked(1)]) as u64, - 32 => u32::from_le_bytes([ - *data.get_unchecked(0), - *data.get_unchecked(1), - *data.get_unchecked(2), - *data.get_unchecked(3), - ]) as u64, - 64 => u64::from_le_bytes([ - *data.get_unchecked(0), - *data.get_unchecked(1), - *data.get_unchecked(2), - *data.get_unchecked(3), - *data.get_unchecked(4), - *data.get_unchecked(5), - *data.get_unchecked(6), - *data.get_unchecked(7), - ]), - _ => unreachable!(), + /// Reads a single length prefix from the front of `data`. + /// + /// The bytes come from the file. A page whose item walk ends with a partial + /// trailing item leaves fewer than `bits_per_offset / 8` bytes here, so this + /// is bounds checked and reports a corrupt file rather than reading past the + /// end of the buffer. + fn parse_length(data: &[u8], bits_per_offset: u8) -> Result { + let width = bits_per_offset as usize / 8; + if data.len() < width { + return Err(corrupt_file_named( + "variable_full_zip", + format!( + "truncated length prefix: {} byte(s) remain in the page buffer but a \ + {}-bit length prefix requires {}", + data.len(), + bits_per_offset, + width + ), + )); } + Ok(match bits_per_offset { + 8 => data[0] as u64, + 16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64, + 32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64, + 64 => u64::from_le_bytes(data[..8].try_into().unwrap()), + _ => unreachable!(), + }) } fn unzip( @@ -2854,7 +2863,7 @@ impl VariableFullZipDecoder { in_bits_per_length: u8, out_bits_per_offset: u8, num_rows: u64, - ) { + ) -> Result<()> { // This undercounts if there are lists but, at this point, we don't really know how many items we have let mut rep = Vec::with_capacity(num_rows as usize); let mut def = Vec::with_capacity(num_rows as usize); @@ -2905,9 +2914,7 @@ impl VariableFullZipDecoder { if ctrl_desc.is_visible { visible_item_count += 1; if ctrl_desc.is_valid_item { - // Safety: Data should have at least bytes_per_length bytes remaining - debug_assert!(databuf.len() >= bytes_per_length); - let length = unsafe { Self::parse_length(databuf, in_bits_per_length) }; + let length = Self::parse_length(databuf, in_bits_per_length)?; match out_bits_per_offset { 32 => offsets_data .extend_from_slice(&(current_offset as u32).to_le_bytes()), @@ -2943,6 +2950,7 @@ impl VariableFullZipDecoder { self.def = ScalarBuffer::from(def); self.data = LanceBuffer::from(unzipped_data); self.offsets = LanceBuffer::from(offsets_data); + Ok(()) } } @@ -7661,4 +7669,67 @@ mod tests { let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; } + + fn truncated_tail_details() -> std::sync::Arc { + use crate::compression::VariablePerValueDecompressor; + use crate::encodings::physical::binary::VariableDecoder; + use crate::repdef::{ControlWordParser, DefinitionInterpretation}; + use std::sync::Arc; + Arc::new(super::FullZipDecodeDetails { + value_decompressor: super::PerValueDecompressor::Variable(Arc::new( + VariableDecoder::default(), + ) + as Arc), + def_meaning: vec![DefinitionInterpretation::NullableItem].into(), + ctrl_word_parser: ControlWordParser::new(0, 0), + max_rep: 0, + max_visible_def: 0, + }) + } + + fn decode_variable_full_zip( + buf: Vec, + bits_per_offset: u8, + ) -> lance_core::Result { + use std::collections::VecDeque; + let mut data = VecDeque::new(); + data.push_back(crate::buffer::LanceBuffer::from(buf)); + super::VariableFullZipDecoder::new( + truncated_tail_details(), + data, + 1, + bits_per_offset, + bits_per_offset, + ) + } + + /// A page whose item walk ends with a partial length prefix must surface a + /// corrupt-file error rather than read past the end of the buffer. + /// + /// This asserts the error variant and message rather than merely expecting a + /// panic: before the length prefix was bounds checked, the read was + /// `get_unchecked` behind a `debug_assert!`, so a debug build panicked here + /// (which a `#[should_panic]` test would have accepted as a pass) while a + /// release build read up to 8 bytes out of a 4 byte allocation. + #[test] + fn variable_full_zip_truncated_length_prefix_is_corrupt_file() { + use lance_core::Error; + + for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] { + let err = decode_variable_full_zip(vec![0xAA; buf_len], bits) + .expect_err("a truncated length prefix must not decode"); + assert!( + matches!(err, Error::CorruptFile { .. }), + "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}", + bits, + buf_len, + err + ); + let msg = err.to_string(); + assert!( + msg.contains("truncated length prefix"), + "error should say what is wrong, got: {msg}" + ); + } + } } diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index d02cf2da693..ecb24ad6d0b 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -13,6 +13,10 @@ use arrow_array::OffsetSizeTrait; use byteorder::{ByteOrder, LittleEndian}; use core::panic; +fn corrupt_file_named(name: &str, message: impl Into) -> Error { + Error::corrupt_file(name.into(), message) +} + use crate::compression::{ BlockCompressor, BlockDecompressor, MiniBlockDecompressor, VariablePerValueDecompressor, }; @@ -288,52 +292,144 @@ impl BinaryMiniBlockDecompressor { } } +/// Cold path: pinpoint why the chunk-relative offsets of a binary mini-block +/// chunk failed validation. +fn chunk_offset_violation_error>(offsets: &[T], chunk_len: usize) -> Error { + let mut previous: u64 = offsets[0].into(); + for (position, &offset) in offsets.iter().enumerate().skip(1) { + let offset: u64 = offset.into(); + if offset < previous { + return corrupt_file_named( + "binary mini-block", + format!( + "value offset at position {position} decreases: {offset} < {previous} \ + (chunk is {chunk_len} bytes)" + ), + ); + } + previous = offset; + } + corrupt_file_named( + "binary mini-block", + format!("value offset {previous} is out of bounds for a chunk of {chunk_len} bytes"), + ) +} + impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { // decompress a MiniBlock of binary data, the num_values must be less than or equal // to the number of values this MiniBlock has, BinaryMiniBlock doesn't store `the number of values` // it has so assertion can not be done here and the caller of `decompress` must ensure // `num_values` <= number of values in the chunk. + // + // The chunk-relative value offsets at the front of the chunk come straight + // from the file and are used to slice the chunk buffer, so corrupt values + // must surface as a typed error instead of a panic or an out-of-bounds + // read. The monotonicity check rides along the existing rebase loop (the + // `&=` accumulation keeps it branchless) so validation adds no extra pass. fn decompress(&self, data: Vec, num_values: u64) -> Result { assert_eq!(data.len(), 1); let data = data.into_iter().next().unwrap(); - if self.bits_per_offset == 64 { - // offset and at least one value - assert!(data.len() >= 16); + let bytes_per_offset = self.bits_per_offset as usize / 8; + if !data.len().is_multiple_of(bytes_per_offset) { + return Err(corrupt_file_named( + "binary mini-block", + format!( + "chunk size {} is not a multiple of the {}-byte offset width", + data.len(), + bytes_per_offset + ), + )); + } + let num_offsets = (num_values as usize).checked_add(1).ok_or_else(|| { + corrupt_file_named( + "binary mini-block", + format!("cannot decode {num_values} values from a single chunk"), + ) + })?; + if data.len() / bytes_per_offset < num_offsets { + return Err(corrupt_file_named( + "binary mini-block", + format!( + "chunk of {} bytes holds {} offsets but decoding {} values requires {}", + data.len(), + data.len() / bytes_per_offset, + num_values, + num_offsets + ), + )); + } + + // The value region must start past the offsets being decoded, otherwise + // the offset table itself aliases into the value bytes. A lower bound + // (not equality) because a prefix read of the chunk legitimately leaves + // unrequested offsets between the requested prefix and the values. + let min_value_region_start = num_offsets * bytes_per_offset; + let value_region_overlap_error = |first: u64| { + corrupt_file_named( + "binary mini-block", + format!( + "value region starts at offset {first} which overlaps the {num_offsets} \ + requested offsets ({min_value_region_start} bytes)" + ), + ) + }; + if self.bits_per_offset == 64 { let offsets_buffer = data.borrow_to_typed_slice::(); - let offsets = offsets_buffer.as_ref(); + let offsets = &offsets_buffer.as_ref()[..num_offsets]; - let result_offsets = offsets[0..(num_values + 1) as usize] + let first = offsets[0]; + if first < min_value_region_start as u64 { + return Err(value_region_overlap_error(first)); + } + let mut previous = first; + let mut is_monotonic = true; + let result_offsets = offsets .iter() - .map(|offset| offset - offsets[0]) + .map(|&offset| { + is_monotonic &= previous <= offset; + previous = offset; + offset.wrapping_sub(first) + }) .collect::>(); + let last = offsets[num_offsets - 1]; + if !is_monotonic || last as usize > data.len() { + return Err(chunk_offset_violation_error(offsets, data.len())); + } Ok(DataBlock::VariableWidth(VariableWidthBlock { - data: LanceBuffer::from( - data[offsets[0] as usize..offsets[num_values as usize] as usize].to_vec(), - ), + data: LanceBuffer::from(data[first as usize..last as usize].to_vec()), offsets: LanceBuffer::reinterpret_vec(result_offsets), bits_per_offset: 64, num_values, block_info: BlockInfo::new(), })) } else { - // offset and at least one value - assert!(data.len() >= 8); - let offsets_buffer = data.borrow_to_typed_slice::(); - let offsets = offsets_buffer.as_ref(); + let offsets = &offsets_buffer.as_ref()[..num_offsets]; - let result_offsets = offsets[0..(num_values + 1) as usize] + let first = offsets[0]; + if (first as u64) < min_value_region_start as u64 { + return Err(value_region_overlap_error(first as u64)); + } + let mut previous = first; + let mut is_monotonic = true; + let result_offsets = offsets .iter() - .map(|offset| offset - offsets[0]) + .map(|&offset| { + is_monotonic &= previous <= offset; + previous = offset; + offset.wrapping_sub(first) + }) .collect::>(); + let last = offsets[num_offsets - 1]; + if !is_monotonic || last as usize > data.len() { + return Err(chunk_offset_violation_error(offsets, data.len())); + } Ok(DataBlock::VariableWidth(VariableWidthBlock { - data: LanceBuffer::from( - data[offsets[0] as usize..offsets[num_values as usize] as usize].to_vec(), - ), + data: LanceBuffer::from(data[first as usize..last as usize].to_vec()), offsets: LanceBuffer::reinterpret_vec(result_offsets), bits_per_offset: 32, num_values, @@ -462,18 +558,48 @@ impl BlockDecompressor for BinaryBlockDecompressor { // never be more than 255 and it's little endian so the last 3 bytes will always be 0. These will be the least // significant 3 bytes of the number of values in the old scheme. It's pretty unlikely these are all 0 (that would // mean there are at least 16M values in a single page) so we'll use this to determine if the old scheme is used. + // + // The header fields and the offsets themselves come straight from the file. + // The structural checks below (all O(1)) reject blocks whose regions do not + // line up; the offset *values* are validated later, by the mandatory layout + // validation in `VariableWidthBlock::into_arrow`, so they are not rescanned + // here. + if data.len() < 4 { + return Err(corrupt_file_named( + "variable-width block", + format!( + "block of {} bytes is too small to hold a header", + data.len() + ), + )); + } let is_old_scheme = data[1] != 0 || data[2] != 0 || data[3] != 0; + let ensure_header = |header_len: usize| { + if data.len() < header_len { + return Err(corrupt_file_named( + "variable-width block", + format!( + "block of {} bytes is too small for a {} byte header", + data.len(), + header_len + ), + )); + } + Ok(()) + }; let (bits_per_offset, bytes_start_offset, offset_start) = if is_old_scheme { // Old scheme let bits_per_offset = data[0]; match bits_per_offset { 32 => { + ensure_header(9)?; debug_assert_eq!(LittleEndian::read_u32(&data[1..5]), num_values as u32); let bytes_start_offset = LittleEndian::read_u32(&data[5..9]); - (bits_per_offset, bytes_start_offset as u64, 9) + (bits_per_offset, bytes_start_offset as u64, 9_u64) } 64 => { + ensure_header(17)?; debug_assert_eq!(LittleEndian::read_u64(&data[1..9]), num_values); let bytes_start_offset = LittleEndian::read_u64(&data[9..17]); (bits_per_offset, bytes_start_offset, 17) @@ -489,10 +615,12 @@ impl BlockDecompressor for BinaryBlockDecompressor { let bits_per_offset = LittleEndian::read_u32(&data[0..4]) as u8; match bits_per_offset { 32 => { + ensure_header(8)?; let bytes_start_offset = LittleEndian::read_u32(&data[4..8]); (bits_per_offset, bytes_start_offset as u64, 8) } 64 => { + ensure_header(16)?; let bytes_start_offset = LittleEndian::read_u64(&data[8..16]); (bits_per_offset, bytes_start_offset, 16) } @@ -504,9 +632,55 @@ impl BlockDecompressor for BinaryBlockDecompressor { } }; + // The offsets region sits between the header and `bytes_start_offset` + // and must hold exactly `num_values + 1` offsets starting at zero. + let expected_offsets_bytes = num_values + .checked_add(1) + .and_then(|num_offsets| num_offsets.checked_mul(bits_per_offset as u64 / 8)) + .ok_or_else(|| { + corrupt_file_named( + "variable-width block", + format!("offsets region size overflows for {num_values} values"), + ) + })?; + if bytes_start_offset < offset_start || bytes_start_offset > data.len() as u64 { + return Err(corrupt_file_named( + "variable-width block", + format!( + "bytes start offset {} is outside the block (header: {} bytes, block: {} bytes)", + bytes_start_offset, + offset_start, + data.len() + ), + )); + } + if bytes_start_offset - offset_start != expected_offsets_bytes { + return Err(corrupt_file_named( + "variable-width block", + format!( + "expected {} offset bytes for {} values but found {}", + expected_offsets_bytes, + num_values, + bytes_start_offset - offset_start + ), + )); + } + // the next `bytes_start_offset - offset_start` stores the offsets. - let offsets = - data.slice_with_length(offset_start, bytes_start_offset as usize - offset_start); + let offsets = data.slice_with_length( + offset_start as usize, + (bytes_start_offset - offset_start) as usize, + ); + let first_offset = match bits_per_offset { + 32 => LittleEndian::read_u32(&offsets[0..4]) as u64, + _ => LittleEndian::read_u64(&offsets[0..8]), + }; + if first_offset != 0 { + return Err(corrupt_file_named( + "variable-width block", + format!("first offset must be 0 but found {first_offset}"), + )); + } // the rest are the binary bytes. let data = data.slice_with_length( @@ -533,10 +707,12 @@ mod tests { use arrow_schema::{DataType, Field}; use crate::{ + buffer::LanceBuffer, constants::{ COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, }, + data::{BlockInfo, DataBlock, VariableWidthBlock}, testing::check_specific_random, }; use rstest::rstest; @@ -968,4 +1144,193 @@ mod tests { } } } + + #[test] + fn test_binary_miniblock_rejects_corrupt_offsets() { + use super::BinaryMiniBlockDecompressor; + use crate::compression::MiniBlockDecompressor; + use lance_core::Error; + + // Chunk layout mirrors the on-disk format for ["alpha", "beta", "gamma"]: + // LE u32 offsets [16, 21, 25, 30] followed by the value bytes, padded to + // a multiple of 8 bytes. + fn chunk_u32(offsets: &[u32], values: &[u8]) -> LanceBuffer { + let mut chunk = offsets + .iter() + .flat_map(|offset| offset.to_le_bytes()) + .collect::>(); + chunk.extend_from_slice(values); + chunk.resize(chunk.len().next_multiple_of(8), 0); + LanceBuffer::from(chunk) + } + + let decompressor = BinaryMiniBlockDecompressor::new(32); + + // The tail offset points past the end of the 32-byte chunk. + let err = decompressor + .decompress( + vec![chunk_u32(&[16, 21, 25, 100_000], b"alphabetagamma")], + 3, + ) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("out of bounds"), "{err}"); + + // Offsets go backwards, which would underflow the rebase subtraction. + let err = decompressor + .decompress(vec![chunk_u32(&[16, 25, 21, 30], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("decreases"), "{err}"); + + // The first offset points inside the offset table, which would alias + // the serialized offsets into the value bytes. + let err = decompressor + .decompress(vec![chunk_u32(&[0, 21, 25, 30], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("overlaps"), "{err}"); + + // The chunk stores fewer offsets than the requested value count needs. + let err = decompressor + .decompress(vec![chunk_u32(&[8, 8], &[])], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("requires 4"), "{err}"); + + // The chunk size is not a multiple of the offset width. + let err = decompressor + .decompress(vec![LanceBuffer::from(vec![0u8; 10])], 1) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("multiple"), "{err}"); + + // 64-bit offsets take the same validation path. + fn chunk_u64(offsets: &[u64], values: &[u8]) -> LanceBuffer { + let mut chunk = offsets + .iter() + .flat_map(|offset| offset.to_le_bytes()) + .collect::>(); + chunk.extend_from_slice(values); + chunk.resize(chunk.len().next_multiple_of(8), 0); + LanceBuffer::from(chunk) + } + let decompressor = BinaryMiniBlockDecompressor::new(64); + let err = decompressor + .decompress( + vec![chunk_u64(&[32, 37, 41, 100_000], b"alphabetagamma")], + 3, + ) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("out of bounds"), "{err}"); + let err = decompressor + .decompress(vec![chunk_u64(&[0, 37, 41, 46], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("overlaps"), "{err}"); + + // A valid chunk still decodes: offsets rebase to [0, 5, 9, 14]. + let decompressor = BinaryMiniBlockDecompressor::new(32); + let block = decompressor + .decompress(vec![chunk_u32(&[16, 21, 25, 30], b"alphabetagamma")], 3) + .unwrap(); + let DataBlock::VariableWidth(block) = block else { + panic!("expected a variable-width block"); + }; + assert_eq!(block.data.as_ref(), b"alphabetagamma"); + assert_eq!( + block.offsets, + LanceBuffer::reinterpret_vec(vec![0_u32, 5, 9, 14]) + ); + } + + fn encoded_binary_block(bits_per_offset: u8) -> Vec { + use crate::compression::BlockCompressor; + + let offsets = match bits_per_offset { + 32 => LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 14]), + 64 => LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 14]), + _ => unreachable!(), + }; + let block = DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets, + bits_per_offset, + num_values: 3, + block_info: BlockInfo::new(), + }); + BlockCompressor::compress(&super::VariableEncoder::default(), block) + .unwrap() + .as_ref() + .to_vec() + } + + /// The block decompressor only checks the block structure (all O(1)); bad + /// offset values inside a structurally-sound block are rejected by the + /// mandatory layout validation when the block is converted to Arrow. + #[rstest] + #[case::i32_tail_out_of_bounds(32, 3, 100_000, "out of bounds")] + #[case::i64_tail_out_of_bounds(64, 3, 15, "out of bounds")] + #[case::i32_non_monotonic(32, 2, 4, "non-monotonic")] + #[case::i64_non_monotonic(64, 2, 4, "non-monotonic")] + fn test_binary_block_bad_offsets_rejected_at_arrow_conversion( + #[case] bits_per_offset: u8, + #[case] mutated_offset_index: usize, + #[case] mutated_offset_value: u64, + #[case] expected_message: &str, + ) { + use crate::compression::BlockDecompressor; + use lance_core::Error; + + let mut encoded = encoded_binary_block(bits_per_offset); + let bytes_per_offset = (bits_per_offset / 8) as usize; + // The standard scheme header is two offset-width fields. + let mutated_offset_start = bytes_per_offset * (2 + mutated_offset_index); + encoded[mutated_offset_start..mutated_offset_start + bytes_per_offset] + .copy_from_slice(&mutated_offset_value.to_le_bytes()[..bytes_per_offset]); + + let block = super::BinaryBlockDecompressor::default() + .decompress(LanceBuffer::from(encoded), 3) + .unwrap(); + let data_type = match bits_per_offset { + 32 => DataType::Binary, + _ => DataType::LargeBinary, + }; + let err = block.into_arrow(data_type, false).unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains(expected_message), "{err}"); + } + + #[test] + fn test_binary_block_rejects_corrupt_structure() { + use crate::compression::BlockDecompressor; + use lance_core::Error; + + let decompressor = super::BinaryBlockDecompressor::default(); + + // The first offset must be zero. + let mut encoded = encoded_binary_block(32); + encoded[8..12].copy_from_slice(&5_u32.to_le_bytes()); + let err = decompressor + .decompress(LanceBuffer::from(encoded), 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("first offset"), "{err}"); + + // The offsets region must hold exactly num_values + 1 offsets. + let encoded = encoded_binary_block(32); + let err = decompressor + .decompress(LanceBuffer::from(encoded), 4) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("offset bytes"), "{err}"); + + // A block too small to hold its header is rejected, not a panic. + let err = decompressor + .decompress(LanceBuffer::from(vec![0_u8; 2]), 1) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("too small"), "{err}"); + } } diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index f631c0a7892..8c89a722428 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -1651,7 +1651,7 @@ mod tests { use std::{collections::BTreeMap, pin::Pin, sync::Arc}; use arrow_array::{ - RecordBatch, UInt32Array, + RecordBatch, RecordBatchIterator, UInt32Array, types::{Float64Type, Int32Type}, }; use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema}; @@ -1755,6 +1755,182 @@ mod tests { assert_eq!(remaining, 0); } + /// Writes `batch` to a fresh file, overwrites `patch` bytes at `patch_offset` + /// into the single occurrence of `pattern`, and reads the file back with the + /// default reader configuration. + async fn read_file_with_mutated_bytes( + version: LanceFileVersion, + batch: RecordBatch, + pattern: &[u8], + patch_offset: usize, + patch: &[u8], + ) -> lance_core::Result> { + let fs = FsFixture::default(); + let schema = batch.schema(); + write_lance_file( + RecordBatchIterator::new(vec![Ok(batch)], schema), + &fs, + FileWriterOptions { + format_version: Some(version), + ..Default::default() + }, + ) + .await; + + let mut bytes = fs + .object_store + .read_one_all(&fs.tmp_path) + .await + .unwrap() + .to_vec(); + let matches = bytes + .windows(pattern.len()) + .enumerate() + .filter_map(|(position, window)| (window == pattern).then_some(position)) + .collect::>(); + assert_eq!( + matches.len(), + 1, + "expected the byte pattern to appear exactly once in the file" + ); + let patch_start = matches[0] + patch_offset; + bytes[patch_start..patch_start + patch.len()].copy_from_slice(patch); + fs.object_store.put(&fs.tmp_path, &bytes).await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &test_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + file_reader + .read_stream( + lance_io::ReadBatchParams::RangeFull, + 1024, + 16, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await + } + + /// A corrupt file whose variable-width offsets point outside the value bytes + /// must fail with a typed error under the default reader configuration + /// (`validate_on_decode` disabled) instead of materializing values outside + /// the data buffer. + /// + /// Uses a dictionary-encoded string column because its values page stores + /// the offsets verbatim, so flipping the tail offset in the file reaches the + /// Arrow conversion boundary without being rejected by an intermediate + /// decompressor. + #[rstest] + #[tokio::test] + async fn test_default_reader_rejects_out_of_bounds_variable_width_offsets( + #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)] + version: LanceFileVersion, + ) { + use arrow_array::{Array, DictionaryArray, Int32Array, StringArray}; + + let values = StringArray::from(vec!["alpha", "beta", "gamma"]); + let indices = Int32Array::from((0..300).map(|i| i % 3).collect::>()); + let dictionary = DictionaryArray::new(indices, Arc::new(values)); + let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "category", + dictionary.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(dictionary)]).unwrap(); + + // The dictionary values page stores the value offsets as plain + // little-endian i32s ending with [5, 9, 14] (2.1 also stores the leading + // zero, 2.2+ omits it). If a future encoding change stops storing these + // offsets verbatim this lookup fails loudly and the test needs a new + // byte pattern. The patch rewrites the tail offset so it points far + // beyond the value bytes. + let offsets_tail_pattern = [5_i32, 9, 14] + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let error = read_file_with_mutated_bytes( + version, + batch, + &offsets_tail_pattern, + 8, + &100_000_i32.to_le_bytes(), + ) + .await + .expect_err("out-of-bounds offsets must fail the read"); + let error_message = error.to_string(); + assert!( + error_message.contains("corrupt file"), + "expected a corruption error, got: {error_message}" + ); + assert!( + error_message.contains("out of bounds"), + "unexpected message: {error_message}" + ); + } + + /// Same contract as the test above, but for a plain (non-dictionary) string + /// column: the mini-block chunk stores chunk-relative value offsets that are + /// used to slice the chunk, so a corrupt tail offset must surface as a typed + /// error from the chunk decompressor instead of a panic in the decode task. + #[rstest] + #[tokio::test] + async fn test_default_reader_rejects_out_of_bounds_miniblock_offsets( + #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)] + version: LanceFileVersion, + ) { + use arrow_array::StringArray; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "strings", + DataType::Utf8, + false, + )])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![Arc::new(StringArray::from(vec!["alpha", "beta", "gamma"]))], + ) + .unwrap(); + + // For ["alpha", "beta", "gamma"] the chunk stores LE i32 offsets + // [16, 21, 25, 30] (chunk-relative: a 16-byte offsets region precedes + // the value bytes). The patch rewrites the tail offset to point far + // past the chunk. + let chunk_offsets_pattern = [16_i32, 21, 25, 30] + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let error = read_file_with_mutated_bytes( + version, + batch, + &chunk_offsets_pattern, + 12, + &100_000_i32.to_le_bytes(), + ) + .await + .expect_err("an out-of-bounds chunk offset must fail the read"); + let error_message = error.to_string(); + assert!( + error_message.contains("corrupt file"), + "expected a corruption error, got: {error_message}" + ); + assert!( + error_message.contains("out of bounds"), + "unexpected message: {error_message}" + ); + } + #[tokio::test] async fn test_round_trip() { let fs = FsFixture::default();