From 22574d8ab3436bc0fd754c7f81b32ff0279c4b82 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 18 Sep 2026 02:12:03 +0300 Subject: [PATCH] Bound CSV recovery offsets across chunks --- arrow-csv/src/reader/records.rs | 68 +++++++++++++++++++++++++-------- 1 file changed, 52 insertions(+), 16 deletions(-) diff --git a/arrow-csv/src/reader/records.rs b/arrow-csv/src/reader/records.rs index a55aaf341829..f93e365ca6f9 100644 --- a/arrow-csv/src/reader/records.rs +++ b/arrow-csv/src/reader/records.rs @@ -230,13 +230,11 @@ impl RecordDecoder { return Ok((0, 0)); } - self.offsets - .resize(self.offsets_len + to_read * self.num_columns, 0); - let max_offsets_len = self - .offsets - .len() - .checked_add(MAX_RECOVERABLE_EXCESS_FIELDS) + let required_offsets = to_read + .checked_mul(self.num_columns) + .and_then(|additional| self.offsets_len.checked_add(additional)) .ok_or_else(|| ArrowError::CsvError("CSV field offset capacity overflowed".into()))?; + self.offsets.resize(required_offsets, 0); let mut input_offset = 0; let mut read = 0; @@ -247,12 +245,20 @@ impl RecordDecoder { self.data.resize(self.data_len + estimated_data, 0); loop { + let max_offsets_len = self + .record_offsets_start + .checked_add(self.num_columns) + .and_then(|expected| expected.checked_add(MAX_RECOVERABLE_EXCESS_FIELDS)) + .ok_or_else(|| { + ArrowError::CsvError("CSV field offset capacity overflowed".into()) + })?; + let offsets_end = self.offsets.len().min(max_offsets_len); let record_input_start = input_offset; let (result, bytes_read, bytes_written, end_positions) = self.delimiter.read_record( &input[input_offset..], &mut self.data[self.data_len..], - &mut self.offsets[self.offsets_len..], + &mut self.offsets[self.offsets_len..offsets_end], ); self.current_field += end_positions; @@ -272,20 +278,16 @@ impl RecordDecoder { } ReadRecordResult::OutputFull => break, ReadRecordResult::OutputEndsFull => { + if self.offsets_len >= max_offsets_len { + return Err(recovery_limit_error(self.line_number)); + } let new_len = self - .offsets - .len() + .offsets_len .checked_add(self.num_columns.max(1)) .ok_or_else(|| { ArrowError::CsvError("CSV field offset capacity overflowed".into()) })? .min(max_offsets_len); - if new_len == self.offsets.len() { - return Err(ArrowError::CsvError(format!( - "malformed CSV record on line {} exceeds the recovery limit of {} excess fields", - self.line_number, MAX_RECOVERABLE_EXCESS_FIELDS - ))); - } self.offsets.resize(new_len, 0); } ReadRecordResult::Record => { @@ -420,6 +422,12 @@ impl RecordDecoder { } } +fn recovery_limit_error(line_number: usize) -> ArrowError { + ArrowError::CsvError(format!( + "malformed CSV record on line {line_number} exceeds the recovery limit of {MAX_RECOVERABLE_EXCESS_FIELDS} excess fields" + )) +} + /// A collection of parsed, UTF-8 CSV records #[derive(Debug)] pub struct StringRecords<'a> { @@ -482,13 +490,14 @@ impl std::fmt::Display for StringRecord<'_> { #[cfg(test)] mod tests { - use crate::reader::records::RecordDecoder; use crate::reader::{CsvRecordError, CsvRecordErrorHandler}; use arrow_schema::ArrowError; use csv_core::Reader; use std::io::{BufRead, BufReader, Cursor}; use std::sync::{Arc, Mutex}; + use super::{MAX_RECOVERABLE_EXCESS_FIELDS, RecordDecoder, recovery_limit_error}; + #[derive(Debug, Clone, Eq, PartialEq)] struct OwnedRecordError { line_number: usize, @@ -637,6 +646,33 @@ mod tests { ); } + #[test] + fn test_invalid_fields_handler_bounds_offsets_across_input_chunks() { + let mut csv = Vec::with_capacity(2 * (MAX_RECOVERABLE_EXCESS_FIELDS + 3)); + for _ in 0..MAX_RECOVERABLE_EXCESS_FIELDS + 3 { + csv.extend_from_slice(b"x,"); + } + csv.push(b'\n'); + + let handler = Arc::new(CollectRecordErrors::default()); + let mut decoder = + RecordDecoder::new(Reader::new(), 2, false).with_record_error_handler(Some(handler)); + let mut input_offset = 0_usize; + let error = loop { + let input_end = input_offset.saturating_add(17).min(csv.len()); + match decoder.decode(&csv[input_offset..input_end], 1) { + Ok((_, bytes_read)) => { + assert!(bytes_read > 0); + input_offset += bytes_read; + } + Err(error) => break error, + } + }; + + assert_eq!(error.to_string(), recovery_limit_error(1).to_string()); + assert_eq!(decoder.offsets.len(), 1 + 2 + MAX_RECOVERABLE_EXCESS_FIELDS); + } + #[test] fn test_skip_insufficient_rows() { let csv = "a\nv\n";