From a4921f116f71f5e22461d73030654214b9f662b3 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Date: Mon, 7 Sep 2026 09:29:55 +0530 Subject: [PATCH] fix(ipc): validate block lengths before allocation --- arrow-ipc/src/reader.rs | 121 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 110 insertions(+), 11 deletions(-) diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs index 52a026863b87..505c4751d606 100644 --- a/arrow-ipc/src/reader.rs +++ b/arrow-ipc/src/reader.rs @@ -36,9 +36,7 @@ use std::io::{BufReader, Read, Seek, SeekFrom}; use std::sync::Arc; use arrow_array::*; -use arrow_buffer::{ - ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, ScalarBuffer, -}; +use arrow_buffer::{BooleanBuffer, Buffer, MutableBuffer, NullBuffer, ScalarBuffer}; use arrow_data::{ArrayData, ArrayDataBuilder, UnsafeFlag}; use arrow_schema::*; @@ -951,12 +949,41 @@ fn get_dictionary_values( } /// Read the data for a given block -fn read_block(mut reader: R, block: &Block) -> Result { - reader.seek(SeekFrom::Start(block.offset() as u64))?; - let body_len = block.bodyLength().to_usize().unwrap(); - let metadata_len = block.metaDataLength().to_usize().unwrap(); - let total_len = body_len.checked_add(metadata_len).unwrap(); +fn read_block( + mut reader: R, + block: &Block, + file_len: u64, +) -> Result { + let offset = u64::try_from(block.offset()) + .map_err(|_| ArrowError::ParseError(format!("Invalid block offset: {}", block.offset())))?; + let body_len = usize::try_from(block.bodyLength()).map_err(|_| { + ArrowError::ParseError(format!("Invalid block body length: {}", block.bodyLength())) + })?; + let metadata_len = usize::try_from(block.metaDataLength()).map_err(|_| { + ArrowError::ParseError(format!( + "Invalid block metadata length: {}", + block.metaDataLength() + )) + })?; + let total_len = body_len.checked_add(metadata_len).ok_or_else(|| { + ArrowError::ParseError("Block body length + metadata length overflowed usize".to_string()) + })?; + let total_len_u64 = u64::try_from(total_len) + .map_err(|_| ArrowError::ParseError("Block total length overflowed u64".to_string()))?; + let block_end = offset.checked_add(total_len_u64).ok_or_else(|| { + ArrowError::ParseError(format!( + "Block offset {offset} + total length {total_len_u64} overflowed u64" + )) + })?; + + if block_end > file_len { + return Err(ArrowError::ParseError(format!( + "Invalid block: offset {offset} + length {total_len_u64} exceeds file length {file_len}" + ))); + } + + reader.seek(SeekFrom::Start(offset))?; let mut buf = MutableBuffer::try_from_len_zeroed(total_len) .map_err(|e| ArrowError::MemoryError(e.to_string()))?; reader.read_exact(&mut buf)?; @@ -1278,7 +1305,10 @@ impl FileReaderBuilder { pub fn build(self, mut reader: R) -> Result, ArrowError> { // Space for ARROW_MAGIC (6 bytes) and length (4 bytes) let mut buffer = [0; 10]; - reader.seek(SeekFrom::End(-10))?; + let end_pos = reader.seek(SeekFrom::End(-10))?; + let file_len = end_pos + .checked_add(10) + .ok_or_else(|| ArrowError::ParseError("File length overflowed u64".to_string()))?; reader.read_exact(&mut buffer)?; let footer_len = read_footer_length(buffer)?; @@ -1339,7 +1369,7 @@ impl FileReaderBuilder { // Create an array of optional dictionary value arrays, one per field. if let Some(dictionaries) = footer.dictionaries() { for block in dictionaries { - let buf = read_block(&mut reader, block)?; + let buf = read_block(&mut reader, block, file_len)?; decoder.read_dictionary(block, &buf)?; } } @@ -1352,6 +1382,7 @@ impl FileReaderBuilder { decoder, schema: projected_schema, custom_metadata, + file_len, }) } } @@ -1423,6 +1454,9 @@ pub struct FileReader { /// User defined metadata custom_metadata: HashMap, + + /// Length of the file/reader in bytes + file_len: u64, } impl fmt::Debug for FileReader { @@ -1500,7 +1534,7 @@ impl FileReader { self.current_block += 1; // read length - let buffer = read_block(&mut self.reader, block)?; + let buffer = read_block(&mut self.reader, block, self.file_len)?; self.decoder.read_record_batch(block, &buffer) } @@ -2020,6 +2054,7 @@ mod tests { use crate::writer::{ DictionaryTracker, IpcDataGenerator, IpcWriteOptions, unslice_run_array, write_message, }; + use arrow_buffer::ArrowNativeType; use super::*; @@ -4264,4 +4299,68 @@ mod tests { "unexpected error: {err}" ); } + + fn file_with_declared_block(body_length: i64, meta_length: i32, offset: i64) -> Vec { + let mut fbb = flatbuffers::FlatBufferBuilder::new(); + + let schema = crate::SchemaBuilder::new(&mut fbb).finish(); + let block = crate::Block::new(offset, meta_length, body_length); + let blocks = fbb.create_vector(&[block]); + + let mut footer = crate::FooterBuilder::new(&mut fbb); + footer.add_version(crate::MetadataVersion::V5); + footer.add_schema(schema); + footer.add_recordBatches(blocks); + let root = footer.finish(); + fbb.finish(root, None); + + let footer_data = fbb.finished_data(); + let footer_len = footer_data.len() as i32; + + let mut file = Vec::new(); + file.extend_from_slice(&crate::ARROW_MAGIC); + file.extend_from_slice(&[0u8; 100]); + file.extend_from_slice(footer_data); + file.extend_from_slice(&footer_len.to_le_bytes()); + file.extend_from_slice(&crate::ARROW_MAGIC); + file + } + + /// A block's declared offset, metadata length, and body length are read from the footer, + /// so they cannot be trusted. Declaring implausible lengths or offsets that extend beyond + /// the file length must be rejected with an ArrowError before attempting any buffer allocation. + #[test] + fn test_file_reader_rejects_implausible_block_length() { + for (body_length, meta_length, offset) in [ + (i64::MAX, 10, 8), + (1 << 50, 10, 8), + (-1, 10, 8), + (10, -1, 8), + (10, 10, -1), + (1000, 100, 8), + (100, 100, i64::MAX - 50), + ] { + let file = file_with_declared_block(body_length, meta_length, offset); + let mut reader = match FileReaderBuilder::new().build(std::io::Cursor::new(file)) { + Ok(reader) => reader, + Err(err) => { + assert!( + matches!(err, ArrowError::ParseError(_)), + "expected ParseError for block ({body_length}, {meta_length}, {offset}): {err}" + ); + continue; + } + }; + let err = reader + .next() + .expect("expected batch attempt") + .expect_err(&format!( + "block with ({body_length}, {meta_length}, {offset}) must be rejected" + )); + assert!( + matches!(err, ArrowError::ParseError(_)), + "expected ParseError for block ({body_length}, {meta_length}, {offset}): {err}" + ); + } + } }