From 1cb622303ed9cc0d0826af9efcb5b5f24e58597a Mon Sep 17 00:00:00 2001 From: Ilia Alshanetsky Date: Sun, 18 Oct 2026 12:07:40 -0400 Subject: [PATCH 1/2] Grow column bodies and wire strings as the data arrives ColumnVector, ColumnEnum and ColumnFixedString resized to the declared row count before reading, ColumnString reserved the declared row count and allocated each declared string length up front, and WireFormat::ReadString resized to the declared length, so a packet of a few dozen bytes could make the client allocate gigabytes. I added WireFormat::ReadPODArray, which starts at 1 MiB, reserves eightfold ahead and zero-fills only what it is about to read, and use it for these bodies and for ReadString. ColumnString caps its initial row reservation and reads strings above 1 MiB into a block of their own that grows the same way. Default 64K-row blocks fit the first step and cost nothing extra; a single 80 MB column body copies 73 MiB while growing. LowCardinality's Load also ignored a short index column, which with incremental loading would deliver fewer rows than declared, so it now throws ProtocolError. --- clickhouse/base/wire_format.h | 44 ++++++- clickhouse/columns/enum.cpp | 3 +- clickhouse/columns/lowcardinality.cpp | 2 +- clickhouse/columns/numeric.cpp | 4 +- clickhouse/columns/string.cpp | 53 ++++++-- ut/BUILD.bazel | 1 + ut/CMakeLists.txt | 1 + ut/allocation_tracker.cpp | 51 ++++++++ ut/columns_ut.cpp | 170 ++++++++++++++++++++++++++ ut/stream_ut.cpp | 60 ++++++++- ut/utils.h | 17 +++ 11 files changed, 384 insertions(+), 22 deletions(-) create mode 100644 ut/allocation_tracker.cpp diff --git a/clickhouse/base/wire_format.h b/clickhouse/base/wire_format.h index be58d7a96..1648dc94c 100644 --- a/clickhouse/base/wire_format.h +++ b/clickhouse/base/wire_format.h @@ -1,7 +1,8 @@ #pragma once -#include +#include #include +#include namespace clickhouse { @@ -15,6 +16,20 @@ class WireFormat { static bool ReadString(InputStream& input, std::string* value); static bool SkipString(InputStream& input); static bool ReadBytes(InputStream& input, void* buf, size_t len); + + /// Replaces the contents of `container` with `count` elements read from `input`. + /// The container grows with the data instead of being sized to `count` up front, so + /// a count declared by the peer is never allocated before its bytes are read. + template + static bool ReadPODArray(InputStream& input, Container* container, size_t count); + + /// Capacity to grow to once `loaded` of `count` elements are held: eight times `loaded` + /// and at least `first`, capped at `count`, without overflowing `Size`. + template + static constexpr Size NextReadCapacity(Size loaded, Size count, Size first) { + return loaded > count / 8 ? count : std::min(count, std::max(first, static_cast(loaded * 8))); + } + static bool ReadUInt64(InputStream& input, uint64_t* value); static bool ReadVarint64(InputStream& output, uint64_t* value); @@ -43,8 +58,7 @@ inline bool WireFormat::ReadString(InputStream& input, std::string* value) { if (len > 0x00FFFFFFULL) { return false; } - value->resize((size_t)len); - return ReadAll(input, value->data(), (size_t)len); + return ReadPODArray(input, value, (size_t)len); } return false; @@ -54,6 +68,30 @@ inline bool WireFormat::ReadBytes(InputStream& input, void* buf, size_t len) { return ReadAll(input, buf, len); } +template +inline bool WireFormat::ReadPODArray(InputStream& input, Container* container, size_t count) { + using T = typename Container::value_type; + constexpr size_t first_chunk = std::max(1, (1 << 20) / sizeof(T)); + + container->clear(); + size_t loaded = 0; + while (loaded < count) { + // Reserving far ahead keeps reallocation copies small; resizing only by what is + // about to be read keeps the zero-filled part proportional to the data received. + if (container->capacity() <= loaded) { + container->reserve(NextReadCapacity(loaded, count, first_chunk)); + } + const size_t chunk = std::min({count - loaded, container->capacity() - loaded, std::max(loaded, first_chunk)}); + container->resize(loaded + chunk); + if (!ReadAll(input, container->data() + loaded, chunk * sizeof(T))) { + return false; + } + loaded += chunk; + } + + return true; +} + inline bool WireFormat::ReadUInt64(InputStream& input, uint64_t* value) { return ReadVarint64(input, value); } diff --git a/clickhouse/columns/enum.cpp b/clickhouse/columns/enum.cpp index 8243de7a8..86217b4f3 100644 --- a/clickhouse/columns/enum.cpp +++ b/clickhouse/columns/enum.cpp @@ -81,8 +81,7 @@ void ColumnEnum::Append(ColumnRef column) { template bool ColumnEnum::LoadBody(InputStream* input, size_t rows) { - data_.resize(rows); - return WireFormat::ReadBytes(*input, data_.data(), data_.size() * sizeof(T)); + return WireFormat::ReadPODArray(*input, &data_, rows); } template diff --git a/clickhouse/columns/lowcardinality.cpp b/clickhouse/columns/lowcardinality.cpp index 22b1995ba..031172444 100644 --- a/clickhouse/columns/lowcardinality.cpp +++ b/clickhouse/columns/lowcardinality.cpp @@ -452,7 +452,7 @@ auto Load(ColumnRef new_dictionary_column, InputStream& input, size_t rows) { if (number_of_rows != rows) throw AssertionError("LowCardinality column must be read in full."); - new_index_column->LoadBody(&input, number_of_rows); + if (!new_index_column->LoadBody(&input, number_of_rows)) throw ProtocolError("Failed to read index column."); if (auto nullable = new_dictionary_column->As()) { nullable->Append(true); diff --git a/clickhouse/columns/numeric.cpp b/clickhouse/columns/numeric.cpp index 2fdf393f1..8261eebf9 100644 --- a/clickhouse/columns/numeric.cpp +++ b/clickhouse/columns/numeric.cpp @@ -76,9 +76,7 @@ void ColumnVector::Append(ColumnRef column) { template bool ColumnVector::LoadBody(InputStream* input, size_t rows) { - data_.resize(rows); - - return WireFormat::ReadBytes(*input, data_.data(), data_.size() * sizeof(T)); + return WireFormat::ReadPODArray(*input, &data_, rows); } template diff --git a/clickhouse/columns/string.cpp b/clickhouse/columns/string.cpp index e62a9b2b5..87150ef06 100644 --- a/clickhouse/columns/string.cpp +++ b/clickhouse/columns/string.cpp @@ -8,6 +8,18 @@ namespace { constexpr size_t DEFAULT_BLOCK_SIZE = 4096; +// Upper bound on what LoadBody allocates ahead of the data, since row counts and +// string lengths come from the peer. +constexpr size_t MAX_PREALLOCATED_BYTES = 1 << 20; + +// Grows `items` eightfold rather than twofold so that loading many rows copies little. +void AppendItem(std::vector& items, size_t rows, std::string_view item) { + if (items.size() == items.capacity()) { + items.reserve(clickhouse::WireFormat::NextReadCapacity(items.size(), rows, size_t(1))); + } + items.emplace_back(item); +} + template size_t ComputeTotalSize(const Container & strings, size_t begin = 0, size_t len = -1) { size_t result = 0; @@ -78,12 +90,7 @@ void ColumnFixedString::Append(ColumnRef column) { } bool ColumnFixedString::LoadBody(InputStream * input, size_t rows) { - data_.resize(string_size_ * rows); - if (!WireFormat::ReadBytes(*input, &data_[0], data_.size())) { - return false; - } - - return true; + return WireFormat::ReadPODArray(*input, &data_, string_size_ * rows); } void ColumnFixedString::SaveBody(OutputStream* output) { @@ -134,6 +141,16 @@ struct ColumnString::Block data_(new CharT[capacity]) {} + // Moves the data, so it must not be called once a string_view points into the block. + void Grow(size_t new_capacity) { + std::unique_ptr new_data(new CharT[new_capacity]); + if (size > 0) { + memcpy(new_data.get(), data_.get(), size); + } + data_ = std::move(new_data); + capacity = new_capacity; + } + inline auto GetAvailable() const { return capacity - size; } @@ -163,7 +180,7 @@ struct ColumnString::Block } size_t size; - const size_t capacity; + size_t capacity; std::unique_ptr data_; }; @@ -274,7 +291,7 @@ bool ColumnString::LoadBody(InputStream* input, size_t rows) { decltype(items_) new_items; decltype(blocks_) new_blocks; - new_items.reserve(rows); + new_items.reserve(std::min(rows, MAX_PREALLOCATED_BYTES / sizeof(std::string_view))); // Suboptimzal if the first row string is >DEFAULT_BLOCK_SIZE, but that must be a very rare case. Block * block = &new_blocks.emplace_back(DEFAULT_BLOCK_SIZE); @@ -284,13 +301,31 @@ bool ColumnString::LoadBody(InputStream* input, size_t rows) { if (!WireFormat::ReadUInt64(*input, &len)) return false; + if (len > MAX_PREALLOCATED_BYTES) { + if (len > SIZE_MAX) return false; + + // A block of its own, grown as the bytes arrive. + block = &new_blocks.emplace_back(MAX_PREALLOCATED_BYTES); + while (block->size < len) { + if (block->GetAvailable() == 0) + block->Grow(WireFormat::NextReadCapacity(block->capacity, static_cast(len), size_t(0))); + + const size_t chunk = std::min(static_cast(len) - block->size, block->GetAvailable()); + if (!WireFormat::ReadBytes(*input, block->GetCurrentWritePos(), chunk)) return false; + block->size += chunk; + } + + AppendItem(new_items, rows, std::string_view(block->data_.get(), block->size)); + continue; + } + if (len > block->GetAvailable()) block = &new_blocks.emplace_back(std::max(DEFAULT_BLOCK_SIZE, len)); if (!WireFormat::ReadBytes(*input, block->GetCurrentWritePos(), len)) return false; - new_items.emplace_back(block->ConsumeTailAsStringViewUnsafe(len)); + AppendItem(new_items, rows, block->ConsumeTailAsStringViewUnsafe(len)); } items_.swap(new_items); diff --git a/ut/BUILD.bazel b/ut/BUILD.bazel index a41e3e545..c7db73558 100644 --- a/ut/BUILD.bazel +++ b/ut/BUILD.bazel @@ -51,6 +51,7 @@ cc_test( "types_ut.cpp", "utils_ut.cpp", # Test entry point and shared support code. + "allocation_tracker.cpp", "main.cpp", "tcp_server.cpp", "tcp_server.h", diff --git a/ut/CMakeLists.txt b/ut/CMakeLists.txt index b830ae175..2aa31a29e 100644 --- a/ut/CMakeLists.txt +++ b/ut/CMakeLists.txt @@ -31,6 +31,7 @@ SET ( clickhouse-cpp-ut-src bignum_samples.cpp utils.cpp + allocation_tracker.cpp value_generators.cpp low_cardinality_nullable_tests.cpp diff --git a/ut/allocation_tracker.cpp b/ut/allocation_tracker.cpp new file mode 100644 index 000000000..be01890e7 --- /dev/null +++ b/ut/allocation_tracker.cpp @@ -0,0 +1,51 @@ +#include + +#include +#include + +#include "utils.h" + +namespace { +thread_local bool track_allocations = false; +thread_local size_t largest_allocation = 0; +} // namespace + +// The default operator new[], operator delete[], and nothrow forms forward to these. +void* operator new(std::size_t size) { + if (track_allocations && size > largest_allocation) { + largest_allocation = size; + } + if (void* ptr = std::malloc(size ? size : 1)) { + return ptr; + } + throw std::bad_alloc(); +} + +void operator delete(void* ptr) noexcept { + std::free(ptr); +} + +void operator delete(void* ptr, std::size_t) noexcept { + std::free(ptr); +} + +LargestAllocationScope::LargestAllocationScope() { + largest_allocation = 0; + track_allocations = true; +} + +LargestAllocationScope::~LargestAllocationScope() { + track_allocations = false; +} + +size_t LargestAllocationScope::Get() const { + return largest_allocation; +} + +bool LargestAllocationScope::ObservesLibraryAllocations() { + constexpr size_t rows = 1 << 20; + LargestAllocationScope allocations; + clickhouse::ColumnUInt64 column; + column.Reserve(rows); + return allocations.Get() >= rows * sizeof(uint64_t); +} diff --git a/ut/columns_ut.cpp b/ut/columns_ut.cpp index 57650ea9e..547ace836 100644 --- a/ut/columns_ut.cpp +++ b/ut/columns_ut.cpp @@ -2349,3 +2349,173 @@ TEST(ColumnsCase, ColumnLowCardinalityT_Wrap_TypeMismatch) { EXPECT_THROW(TestLC::Wrap(bad_dict), ValidationError); EXPECT_THROW(TestLC::Wrap(not_lc), ValidationError); } + +namespace { + +// Truncated bodies below declare this much data but carry only a few bytes of it. +constexpr size_t DECLARED_BODY_SIZE = 128 * 1024 * 1024; +// Generous for growth-as-data-arrives, far below DECLARED_BODY_SIZE. +constexpr size_t MAX_ALLOCATION_FOR_TRUNCATED_BODY = 8 * 1024 * 1024; +// Above the size LoadBody allocates ahead of the data, so loading has to grow. +constexpr size_t LARGE_BODY_SIZE = 3 * 1024 * 1024 + 17; +// Above eight times that size, so loading grows more than once. +constexpr size_t HUGE_BODY_SIZE = 9 * 1024 * 1024 + 17; + +Buffer SaveColumn(Column& column) { + Buffer buffer; + BufferOutput output(&buffer); + column.Save(&output); + output.Flush(); + return buffer; +} + +} // namespace + +TEST(ColumnsCase, NumericLoad_TruncatedBodyDoesNotAllocateDeclaredRows) { + const uint64_t payload[] = {1, 2}; + ArrayInput input(payload, sizeof(payload)); + ColumnUInt64 column; + + SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED(); + LargestAllocationScope allocations; + EXPECT_FALSE(column.Load(&input, DECLARED_BODY_SIZE / sizeof(uint64_t))); + EXPECT_LT(allocations.Get(), MAX_ALLOCATION_FOR_TRUNCATED_BODY); +} + +TEST(ColumnsCase, ArrayLoad_TruncatedBodyDoesNotAllocateDeclaredNestedRows) { + const uint64_t payload[] = {DECLARED_BODY_SIZE / sizeof(uint64_t), 1, 2}; + ArrayInput input(payload, sizeof(payload)); + auto column = CreateColumnByType("Array(UInt64)"); + + SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED(); + LargestAllocationScope allocations; + EXPECT_FALSE(column->Load(&input, 1)); + EXPECT_LT(allocations.Get(), MAX_ALLOCATION_FOR_TRUNCATED_BODY); +} + +TEST(ColumnsCase, FixedStringLoad_TruncatedBodyDoesNotAllocateDeclaredSize) { + const char payload[] = "0123456789abcdef"; + ArrayInput input(payload, sizeof(payload) - 1); + ColumnFixedString column(DECLARED_BODY_SIZE); + + SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED(); + LargestAllocationScope allocations; + EXPECT_FALSE(column.Load(&input, 1)); + EXPECT_LT(allocations.Get(), MAX_ALLOCATION_FOR_TRUNCATED_BODY); +} + +TEST(ColumnsCase, StringLoad_TruncatedBodyDoesNotAllocateDeclaredLength) { + // Varint of DECLARED_BODY_SIZE, then three of its bytes. + const uint8_t payload[] = {0x80, 0x80, 0x80, 0x40, 'a', 'b', 'c'}; + ArrayInput input(payload, sizeof(payload)); + ColumnString column; + + SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED(); + LargestAllocationScope allocations; + EXPECT_FALSE(column.Load(&input, 1)); + EXPECT_LT(allocations.Get(), MAX_ALLOCATION_FOR_TRUNCATED_BODY); +} + +TEST(ColumnsCase, StringLoad_TruncatedBodyDoesNotAllocateDeclaredRows) { + const uint8_t payload[] = {3, 'a', 'b', 'c'}; + ArrayInput input(payload, sizeof(payload)); + ColumnString column; + + SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED(); + LargestAllocationScope allocations; + EXPECT_FALSE(column.Load(&input, DECLARED_BODY_SIZE / sizeof(std::string_view))); + EXPECT_LT(allocations.Get(), MAX_ALLOCATION_FOR_TRUNCATED_BODY); +} + +TEST(ColumnsCase, NumericLoad_LargeBodyRoundTrip) { + ColumnUInt64 source; + for (uint64_t i = 0; i < LARGE_BODY_SIZE / sizeof(uint64_t); ++i) { + source.Append(i * 2654435761u); + } + const Buffer buffer = SaveColumn(source); + + ArrayInput input(buffer.data(), buffer.size()); + ColumnUInt64 loaded; + ASSERT_TRUE(loaded.Load(&input, source.Size())); + EXPECT_TRUE(input.Exhausted()); + EXPECT_EQ(loaded.GetWritableData(), source.GetWritableData()); +} + +TEST(ColumnsCase, FixedStringLoad_LargeBodyRoundTrip) { + const size_t width = 1000; + ColumnFixedString source(width); + for (size_t i = 0; i < LARGE_BODY_SIZE / width; ++i) { + source.Append(std::string(width, static_cast('a' + i % 26))); + } + const Buffer buffer = SaveColumn(source); + + ArrayInput input(buffer.data(), buffer.size()); + ColumnFixedString loaded(width); + ASSERT_TRUE(loaded.Load(&input, source.Size())); + EXPECT_TRUE(input.Exhausted()); + ASSERT_EQ(loaded.Size(), source.Size()); + for (size_t i = 0; i < source.Size(); ++i) { + ASSERT_EQ(loaded.At(i), source.At(i)) << "at row " << i; + } +} + +TEST(ColumnsCase, StringLoad_LargeItemRoundTrip) { + std::string large(LARGE_BODY_SIZE, '\0'); + for (size_t i = 0; i < large.size(); ++i) { + large[i] = static_cast(i * 31); + } + const std::vector items = {"first", large, "", "after", std::string(5000, 'x'), large.substr(7)}; + + ColumnString source(items); + const Buffer buffer = SaveColumn(source); + + ArrayInput input(buffer.data(), buffer.size()); + ColumnString loaded; + ASSERT_TRUE(loaded.Load(&input, items.size())); + EXPECT_TRUE(input.Exhausted()); + ASSERT_EQ(loaded.Size(), items.size()); + for (size_t i = 0; i < items.size(); ++i) { + EXPECT_EQ(loaded.At(i), items[i]) << "at row " << i; + } +} + +TEST(ColumnsCase, StringLoad_HugeItemRoundTrip) { + std::string huge(HUGE_BODY_SIZE, '\0'); + for (size_t i = 0; i < huge.size(); ++i) { + huge[i] = static_cast(i * 131); + } + const std::vector items = {huge, "after"}; + + ColumnString source(items); + const Buffer buffer = SaveColumn(source); + + ArrayInput input(buffer.data(), buffer.size()); + ColumnString loaded; + ASSERT_TRUE(loaded.Load(&input, items.size())); + EXPECT_TRUE(input.Exhausted()); + ASSERT_EQ(loaded.Size(), items.size()); + EXPECT_EQ(loaded.At(0), huge); + EXPECT_EQ(loaded.At(1), "after"); +} + +TEST(ColumnsCase, NumericLoad_HugeBodyRoundTrip) { + ColumnUInt64 source; + for (uint64_t i = 0; i < HUGE_BODY_SIZE / sizeof(uint64_t); ++i) { + source.Append(i * 0x9E3779B97F4A7C15ULL); + } + const Buffer buffer = SaveColumn(source); + + ArrayInput input(buffer.data(), buffer.size()); + ColumnUInt64 loaded; + ASSERT_TRUE(loaded.Load(&input, source.Size())); + EXPECT_TRUE(input.Exhausted()); + EXPECT_EQ(loaded.GetWritableData(), source.GetWritableData()); +} + +TEST(ColumnsCase, ColumnLowCardinalityString_Load_TruncatedIndex) { + const auto& data = LOWCARDINALITY_STRING_FOOBAR_10_ITEMS_BINARY; + ArrayInput input(data.data(), data.size() - 1); + ColumnLowCardinalityT column; + + EXPECT_FALSE(column.Load(&input, 10)); +} diff --git a/ut/stream_ut.cpp b/ut/stream_ut.cpp index 0bcf669c8..fb75c8e69 100644 --- a/ut/stream_ut.cpp +++ b/ut/stream_ut.cpp @@ -1,10 +1,11 @@ -#include -#include -#include #include - +#include +#include +#include #include +#include "utils.h" + using namespace clickhouse; namespace { @@ -56,3 +57,54 @@ TEST(CompressedInputCase, RejectsBlocksShorterThanHeader) { << "compressed size: " << compressed_size; } } + +TEST(CodedStreamCase, ReadStringDoesNotAllocateDeclaredLength) { + Buffer buffer; + { + BufferOutput output(&buffer); + WireFormat::WriteUInt64(output, 0x00FFFFFFULL); + WireFormat::WriteBytes(output, "abc", 3); + output.Flush(); + } + ArrayInput input(buffer.data(), buffer.size()); + std::string value; + + SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED(); + LargestAllocationScope allocations; + EXPECT_FALSE(WireFormat::ReadString(input, &value)); + EXPECT_LT(allocations.Get(), 4u * 1024 * 1024); +} + +TEST(CodedStreamCase, ReadStringRoundTrip) { + const std::string large(3 * 1024 * 1024 + 17, 'q'); + Buffer buffer; + { + BufferOutput output(&buffer); + WireFormat::WriteString(output, ""); + WireFormat::WriteString(output, "name"); + WireFormat::WriteString(output, large); + output.Flush(); + } + ArrayInput input(buffer.data(), buffer.size()); + + std::string value = "stale"; + ASSERT_TRUE(WireFormat::ReadString(input, &value)); + EXPECT_EQ(value, ""); + ASSERT_TRUE(WireFormat::ReadString(input, &value)); + EXPECT_EQ(value, "name"); + ASSERT_TRUE(WireFormat::ReadString(input, &value)); + EXPECT_EQ(value, large); + EXPECT_TRUE(input.Exhausted()); +} + +TEST(CodedStreamCase, NextReadCapacityDoesNotOverflow) { + constexpr uint32_t first = 1 << 20; + EXPECT_EQ(WireFormat::NextReadCapacity(0, 3, first), 3u); + EXPECT_EQ(WireFormat::NextReadCapacity(0, 1u << 30, first), first); + EXPECT_EQ(WireFormat::NextReadCapacity(first, 1u << 30, first), 8 * first); + // loaded * 8 wraps in 32 bits from 2^29 on; the result must still cover the count. + constexpr uint32_t count = (1u << 29) + 4096; + EXPECT_EQ(WireFormat::NextReadCapacity(1u << 29, count, first), count); + EXPECT_EQ(WireFormat::NextReadCapacity(count / 8 + 1, count, first), count); + EXPECT_EQ(WireFormat::NextReadCapacity(uint64_t(1) << 61, UINT64_MAX, 1), UINT64_MAX); +} diff --git a/ut/utils.h b/ut/utils.h index 4d46c8d8a..354e98471 100644 --- a/ut/utils.h +++ b/ut/utils.h @@ -209,3 +209,20 @@ inline uint64_t versionNumber( uint64_t versionNumber(const clickhouse::ServerInfo & server_info); std::string ToString(const clickhouse::UUID& v); + +/// Records the largest single operator new request made by the current thread +/// while the object is alive. +class LargestAllocationScope { +public: + LargestAllocationScope(); + ~LargestAllocationScope(); + + size_t Get() const; + + /// Whether allocations made inside the library reach the replaced operator new; + /// they do not when the library brings its own allocator (e.g. a DLL build). + static bool ObservesLibraryAllocations(); +}; + +#define SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED() \ + if (!LargestAllocationScope::ObservesLibraryAllocations()) GTEST_SKIP() << "the library's allocations bypass the test's operator new" From bdf7de17b9c6c966387babbb64c4eb809a761a13 Mon Sep 17 00:00:00 2001 From: Ilia Alshanetsky Date: Sun, 18 Oct 2026 16:59:41 -0400 Subject: [PATCH 2/2] Reject block sizes that do not fit in size_t Client::Impl::ReadBlock() reads the block's column and row counts as uint64_t and passes the row count to Column::Load(), which takes a size_t. On 32-bit, a server that declares 2^41 rows has the count truncated to 0, and the client returns an empty block instead of an error. Counts that do not survive the conversion now throw ProtocolError. The test only runs where size_t is narrower than 64 bits and skips elsewhere. --- clickhouse/client.cpp | 4 ++++ ut/client_protocol_ut.cpp | 26 ++++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/clickhouse/client.cpp b/clickhouse/client.cpp index 531498c05..1d26a3e34 100644 --- a/clickhouse/client.cpp +++ b/clickhouse/client.cpp @@ -835,6 +835,10 @@ bool Client::Impl::ReadBlock(InputStream& input, Block* block) { if (!WireFormat::ReadUInt64(input, &num_rows)) { return false; } + if (static_cast(num_columns) != num_columns || static_cast(num_rows) != num_rows) { + throw ProtocolError("block declares " + std::to_string(num_columns) + " columns and " + std::to_string(num_rows) + + " rows, more than size_t can hold"); + } CreateColumnByTypeSettings create_column_settings; create_column_settings.low_cardinality_as_wrapped_column = options_.backward_compatibility_lowcardinality_as_wrapped_column; diff --git a/ut/client_protocol_ut.cpp b/ut/client_protocol_ut.cpp index 72cfdd5d6..e3fa1ec1e 100644 --- a/ut/client_protocol_ut.cpp +++ b/ut/client_protocol_ut.cpp @@ -238,3 +238,29 @@ TEST(ClientProtocol, WellFormedSelectResponseSucceeds) { })); EXPECT_EQ(blocks, 1u); } + +TEST(ClientProtocol, RowCountBeyondSizeTThrowsProtocolError) { + if (sizeof(size_t) >= sizeof(uint64_t)) { + GTEST_SKIP() << "every uint64_t row count fits in size_t here"; + } + + std::vector script = kServerHello; + const std::vector response = { + 0x01, // packet type: ServerCodes::Data (varint) + 0x01, // num_columns = 1 (varint) + 0x80, 0x80, 0x80, 0x80, 0x80, 0x40, // num_rows = 2^41 (varint), 0 once truncated to 32 bits + 0x01, 0x63, // column name: "c" + 0x05, 0x55, 0x49, 0x6E, 0x74, 0x38, // column type: "UInt8" + 0x05, // packet type: ServerCodes::EndOfStream (varint) + }; + script.insert(script.end(), response.begin(), response.end()); + + Client client(ScriptedClientOptions(), std::make_unique(script)); + + try { + client.Select("SELECT 1", [](const Block&) {}); + FAIL() << "expected ProtocolError"; + } catch (const ProtocolError& e) { + EXPECT_TRUE(std::string(e.what()).find("more than size_t can hold") != std::string::npos) << "unexpected message: " << e.what(); + } +}