Skip to content

Commit b93fa52

Browse files
authored
fix: ipc decode panic with invalid data (#8931)
Will panic if input buffer is invalid. **CLOUDFLARE ALERT!!!**
1 parent 7b2143b commit b93fa52

1 file changed

Lines changed: 53 additions & 1 deletion

File tree

arrow-ipc/src/reader.rs

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -569,7 +569,10 @@ impl<'a> RecordBatchDecoder<'a> {
569569
}
570570

571571
fn next_buffer(&mut self) -> Result<Buffer, ArrowError> {
572-
read_buffer(self.buffers.next().unwrap(), self.data, self.compression)
572+
let buffer = self.buffers.next().ok_or_else(|| {
573+
ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
574+
})?;
575+
read_buffer(buffer, self.data, self.compression)
573576
}
574577

575578
fn skip_buffer(&mut self) {
@@ -2001,6 +2004,55 @@ mod tests {
20012004
);
20022005
}
20032006

2007+
#[test]
2008+
fn test_missing_buffer_metadata_error() {
2009+
use crate::r#gen::Message::*;
2010+
use flatbuffers::FlatBufferBuilder;
2011+
2012+
let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2013+
2014+
// create RecordBatch buffer metadata with invalid buffer count
2015+
// Int32Array needs 2 buffers (validity + data) but we provide only 1
2016+
let mut fbb = FlatBufferBuilder::new();
2017+
let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2018+
let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2019+
let batch_offset = RecordBatch::create(
2020+
&mut fbb,
2021+
&RecordBatchArgs {
2022+
length: 2,
2023+
nodes: Some(nodes),
2024+
buffers: Some(buffers),
2025+
compression: None,
2026+
variadicBufferCounts: None,
2027+
},
2028+
);
2029+
fbb.finish_minimal(batch_offset);
2030+
let batch_bytes = fbb.finished_data().to_vec();
2031+
let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2032+
2033+
let data_buffer = Buffer::from(vec![0u8; 8]);
2034+
let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2035+
let metadata = MetadataVersion::V5;
2036+
2037+
let decoder = RecordBatchDecoder::try_new(
2038+
&data_buffer,
2039+
batch,
2040+
schema.clone(),
2041+
&dictionaries,
2042+
&metadata,
2043+
)
2044+
.unwrap();
2045+
2046+
let result = decoder.read_record_batch();
2047+
2048+
match result {
2049+
Err(ArrowError::IpcError(msg)) => {
2050+
assert_eq!(msg, "Buffer count mismatched with metadata");
2051+
}
2052+
other => panic!("unexpected error: {other:?}"),
2053+
}
2054+
}
2055+
20042056
#[test]
20052057
fn test_projection_array_values() {
20062058
// define schema

0 commit comments

Comments
 (0)