Skip to content

Commit 25d8bed

Browse files
pitroupitrou
authored
GH-45497: [C++][CSV] Avoid buffer overflow when a line has too many columns (#45498)
### What changes are included in this PR? 1. Add guard against going past the buffer's end, while minimizing the performance overhead of the runtime check. 2. Add error propagation for buffer (re)allocation, instead of aborting. This is unrelated to the reported issue, but is desirable nevertheless. With these changes, a CSV line with an unexpectedly large number of columns will raise an error such as "CSV parse error: Expected 13 columns, got ...". ### Are these changes tested? Yes, by additional unit tests. ### Are there any user-facing changes? No. * GitHub Issue: #45497 Authored-by: Antoine Pitrou <[email protected]> Signed-off-by: Antoine Pitrou <[email protected]>
1 parent c2432af commit 25d8bed

2 files changed

Lines changed: 130 additions & 16 deletions

File tree

‎cpp/src/arrow/csv/parser.cc‎

Lines changed: 58 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -126,16 +126,30 @@ class ValueDescWriter {
126126
{static_cast<uint32_t>(parsed_writer->size()) & 0x7fffffffU, quoted_});
127127
}
128128

129-
void Finish(std::shared_ptr<Buffer>* out_values) {
130-
ARROW_CHECK_OK(values_buffer_->Resize(values_size_ * sizeof(*values_)));
131-
*out_values = values_buffer_;
129+
Result<std::shared_ptr<Buffer>> Finish() {
130+
RETURN_NOT_OK(values_buffer_->Resize(values_size_ * sizeof(*values_)));
131+
return std::move(values_buffer_);
132+
}
133+
134+
const Status& status() const { return status_; }
135+
136+
// Convenience error-checking factory. The arguments are forwarded to the
137+
// Derived class constructor.
138+
template <typename... Args>
139+
static Result<Derived> Make(Args&&... args) {
140+
auto self = Derived(std::forward<Args>(args)...);
141+
RETURN_NOT_OK(self.status());
142+
return self;
132143
}
133144

134145
protected:
135146
ValueDescWriter(MemoryPool* pool, int64_t values_capacity)
136-
: values_size_(0), values_capacity_(values_capacity) {
137-
values_buffer_ = *AllocateResizableBuffer(values_capacity_ * sizeof(*values_), pool);
138-
values_ = reinterpret_cast<ParsedValueDesc*>(values_buffer_->mutable_data());
147+
: values_size_(0), values_capacity_(values_capacity), status_(Status::OK()) {
148+
status_ &= AllocateResizableBuffer(values_capacity_ * sizeof(*values_), pool)
149+
.Value(&values_buffer_);
150+
if (status_.ok()) {
151+
values_ = reinterpret_cast<ParsedValueDesc*>(values_buffer_->mutable_data());
152+
}
139153
}
140154

141155
std::shared_ptr<ResizableBuffer> values_buffer_;
@@ -145,6 +159,7 @@ class ValueDescWriter {
145159
bool quoted_;
146160
// Checkpointing, for when an incomplete line is encountered at end of block
147161
int64_t saved_values_size_;
162+
Status status_;
148163
};
149164

150165
// A helper class handling a growable buffer for values offsets. This class is
@@ -157,11 +172,21 @@ class ResizableValueDescWriter : public ValueDescWriter<ResizableValueDescWriter
157172

158173
void PushValue(ParsedValueDesc v) {
159174
if (ARROW_PREDICT_FALSE(values_size_ == values_capacity_)) {
160-
values_capacity_ = values_capacity_ * 2;
161-
ARROW_CHECK_OK(values_buffer_->Resize(values_capacity_ * sizeof(*values_)));
162-
values_ = reinterpret_cast<ParsedValueDesc*>(values_buffer_->mutable_data());
175+
int64_t new_capacity = values_capacity_ * 2;
176+
auto resize_status = values_buffer_->Resize(new_capacity * sizeof(*values_));
177+
if (resize_status.ok()) {
178+
values_ = reinterpret_cast<ParsedValueDesc*>(values_buffer_->mutable_data());
179+
values_capacity_ = new_capacity;
180+
}
181+
status_ &= std::move(resize_status);
182+
}
183+
// The `values_` pointer may have become invalid if the `Resize` call above failed.
184+
// Note that ResizableValueDescWriter is less performance-critical than
185+
// PresizedValueDescWriter, as it should only be called on the first line(s)
186+
// of CSV data.
187+
if (ARROW_PREDICT_TRUE(status_.ok())) {
188+
values_[values_size_++] = v;
163189
}
164-
values_[values_size_++] = v;
165190
}
166191
};
167192

@@ -171,12 +196,26 @@ class ResizableValueDescWriter : public ValueDescWriter<ResizableValueDescWriter
171196
// faster CSV parsing code.
172197
class PresizedValueDescWriter : public ValueDescWriter<PresizedValueDescWriter> {
173198
public:
199+
// The number of offsets being written will be `1 + num_rows * num_cols`,
200+
// however we allow for one extraneous write in case of excessive columns,
201+
// hence `2 + num_rows * num_cols` (see explanation in PushValue below).
174202
PresizedValueDescWriter(MemoryPool* pool, int32_t num_rows, int32_t num_cols)
175-
: ValueDescWriter(pool, /*values_capacity=*/1 + num_rows * num_cols) {}
203+
: ValueDescWriter(pool, /*values_capacity=*/2 + num_rows * num_cols) {}
176204

177205
void PushValue(ParsedValueDesc v) {
178206
DCHECK_LT(values_size_, values_capacity_);
179-
values_[values_size_++] = v;
207+
values_[values_size_] = v;
208+
// We must take care not to write past the buffer's end if the line being
209+
// parsed has more than `num_cols` columns. The obvious solution of setting
210+
// an error status hurts too much on benchmarks, which is why we instead
211+
// cap `values_size_` to stay inside the buffer.
212+
//
213+
// Not setting an error immediately is not a problem since the `num_cols`
214+
// mismatch is detected later in ParseLine.
215+
//
216+
// Note that we want `values_size_` to reflect the number of written values
217+
// in the nominal case, which is why we choose a slightly larger `values_capacity_`.
218+
values_size_ += (values_size_ != values_capacity_ - 1);
180219
}
181220
};
182221

@@ -464,6 +503,7 @@ class BlockParserImpl {
464503
RETURN_NOT_OK((ParseLine<SpecializedOptions, true>(values_writer, parsed_writer,
465504
data, data_end, is_final,
466505
&line_end, bulk_filter)));
506+
RETURN_NOT_OK(values_writer->status());
467507
if (line_end == data) {
468508
// Cannot parse any further
469509
*finished_parsing = true;
@@ -477,6 +517,7 @@ class BlockParserImpl {
477517
RETURN_NOT_OK((ParseLine<SpecializedOptions, false>(values_writer, parsed_writer,
478518
data, data_end, is_final,
479519
&line_end, bulk_filter)));
520+
RETURN_NOT_OK(values_writer->status());
480521
if (line_end == data) {
481522
// Cannot parse any further
482523
*finished_parsing = true;
@@ -496,8 +537,7 @@ class BlockParserImpl {
496537
}
497538

498539
// Append new buffers and update size
499-
std::shared_ptr<Buffer> values_buffer;
500-
values_writer->Finish(&values_buffer);
540+
ARROW_ASSIGN_OR_RAISE(auto values_buffer, values_writer->Finish());
501541
if (values_buffer->size() > 0) {
502542
values_size_ +=
503543
static_cast<int32_t>(values_buffer->size() / sizeof(ParsedValueDesc) - 1);
@@ -535,7 +575,7 @@ class BlockParserImpl {
535575
// Can't presize values when the number of columns is not known, first parse
536576
// a single line
537577
const int32_t rows_in_chunk = 1;
538-
ResizableValueDescWriter values_writer(pool_);
578+
ARROW_ASSIGN_OR_RAISE(auto values_writer, ResizableValueDescWriter::Make(pool_));
539579
values_writer.Start(parsed_writer);
540580

541581
RETURN_NOT_OK(ParseChunk<SpecializedOptions>(
@@ -560,7 +600,9 @@ class BlockParserImpl {
560600
rows_in_chunk = std::min(kTargetChunkSize, max_num_rows_ - batch_.num_rows_);
561601
}
562602

563-
PresizedValueDescWriter values_writer(pool_, rows_in_chunk, batch_.num_cols_);
603+
ARROW_ASSIGN_OR_RAISE(
604+
auto values_writer,
605+
PresizedValueDescWriter::Make(pool_, rows_in_chunk, batch_.num_cols_));
564606
values_writer.Start(parsed_writer);
565607

566608
RETURN_NOT_OK(ParseChunk<SpecializedOptions>(

‎cpp/src/arrow/csv/parser_test.cc‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@
1717

1818
#include <algorithm>
1919
#include <cstdint>
20+
#include <sstream>
2021
#include <string>
22+
#include <string_view>
2123
#include <utility>
2224
#include <vector>
2325

@@ -586,6 +588,28 @@ TEST(BlockParser, QuotesSpecial) {
586588
}
587589
}
588590

591+
std::vector<std::string> MismatchingNumColumns(int32_t num_cols, int32_t mismatch,
592+
int64_t extra_lines = 0) {
593+
auto write_line = [](int32_t num_cols, std::string_view prefix, std::ostream* out) {
594+
for (int32_t i = 0; i < num_cols; ++i) {
595+
*out << prefix << i << ",";
596+
}
597+
out->seekp(-1, std::ios_base::cur);
598+
*out << "\n";
599+
};
600+
601+
std::stringstream csv_data;
602+
// Output first line with `num_cols` columns
603+
write_line(num_cols, "a", &csv_data);
604+
// Output second line with mismatching number of columns
605+
write_line(num_cols + mismatch, "b", &csv_data);
606+
// Output extra lines with `num_cols` columns each
607+
for (int64_t i = 0; i < extra_lines; ++i) {
608+
write_line(num_cols, "c", &csv_data);
609+
}
610+
return {csv_data.str()};
611+
}
612+
589613
TEST(BlockParser, MismatchingNumColumns) {
590614
uint32_t out_size;
591615
{
@@ -621,6 +645,25 @@ TEST(BlockParser, MismatchingNumColumns) {
621645
EXPECT_RAISES_WITH_MESSAGE_THAT(
622646
Invalid, testing::HasSubstr("CSV parse error: Expected 2 columns, got 1: a"), st);
623647
}
648+
// Vary the number of columns and mismatch, to catch buffer overflow issues
649+
for (int32_t num_cols : {1, 2, 5, 100}) {
650+
ARROW_SCOPED_TRACE("num_cols = ", num_cols);
651+
for (int32_t mismatch : {-5, -1, 1, 5, 10, 50, 1024, 32767}) {
652+
if (mismatch + num_cols <= 0) {
653+
continue;
654+
}
655+
ARROW_SCOPED_TRACE("mismatch = ", mismatch);
656+
// Try to parse CSV data
657+
auto csv_data = MismatchingNumColumns(num_cols, mismatch);
658+
BlockParser parser(ParseOptions::Defaults(), num_cols, /*first_row=*/1);
659+
Status st = Parse(parser, MakeCSVData(csv_data), &out_size);
660+
std::stringstream expected_error;
661+
expected_error << "Row #2: Expected " << num_cols << " columns, got "
662+
<< num_cols + mismatch << ":";
663+
EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, testing::HasSubstr(expected_error.str()),
664+
st);
665+
}
666+
}
624667
}
625668

626669
TEST(BlockParser, MismatchingNumColumnsHandler) {
@@ -724,6 +767,35 @@ TEST(BlockParser, MismatchingNumColumnsHandler) {
724767

725768
ASSERT_NO_FATAL_FAILURE(AssertLastRowEq(parser, {"j", "k"}, {false, false}));
726769
}
770+
771+
// Vary the number of columns and mismatch, to catch buffer overflow issues
772+
for (int32_t num_cols : {1, 2, 5, 100}) {
773+
ARROW_SCOPED_TRACE("num_cols = ", num_cols);
774+
for (int32_t mismatch : {-5, -1, 1, 5, 10, 50, 1024, 32767}) {
775+
if (mismatch + num_cols <= 0) {
776+
continue;
777+
}
778+
ARROW_SCOPED_TRACE("mismatch = ", mismatch);
779+
// Parse CSV data
780+
auto csv_data = MismatchingNumColumns(num_cols, mismatch, /*extra_lines=*/1);
781+
ParseOptions opts = ParseOptions::Defaults();
782+
CustomHandler handler;
783+
opts.invalid_row_handler = handler;
784+
BlockParser parser(opts, num_cols, /*first_row=*/1);
785+
ASSERT_NO_FATAL_FAILURE(AssertParseOk(parser, MakeCSVData(csv_data)));
786+
ASSERT_EQ(2, parser.num_rows());
787+
ASSERT_EQ(3, parser.total_num_rows());
788+
ASSERT_EQ(1, handler.rows.size());
789+
const auto& invalid_row = handler.rows[0];
790+
ASSERT_EQ(num_cols, invalid_row.first.expected_columns);
791+
ASSERT_EQ(num_cols + mismatch, invalid_row.first.actual_columns);
792+
ASSERT_EQ("b0", invalid_row.second.substr(0, 2));
793+
std::vector<std::string> last_row;
794+
GetLastRow(parser, &last_row);
795+
ASSERT_EQ(last_row.front(), "c0");
796+
ASSERT_EQ(last_row.back(), "c" + std::to_string(num_cols - 1));
797+
}
798+
}
727799
}
728800

729801
TEST(BlockParser, Escaping) {

0 commit comments

Comments
 (0)