diff --git a/clickhouse/columns/nullable.cpp b/clickhouse/columns/nullable.cpp index bb9f295e..cb2fa5ea 100644 --- a/clickhouse/columns/nullable.cpp +++ b/clickhouse/columns/nullable.cpp @@ -66,6 +66,9 @@ bool ColumnNullable::LoadBody(InputStream* input, size_t rows) { if (!nested_->LoadBody(input, rows)) { return false; } + if (nested_->Size() != rows) { + throw ProtocolError("Nullable nested column loaded " + std::to_string(nested_->Size()) + " rows, expected " + std::to_string(rows)); + } return true; } diff --git a/clickhouse/columns/string.cpp b/clickhouse/columns/string.cpp index e62a9b2b..2bad457a 100644 --- a/clickhouse/columns/string.cpp +++ b/clickhouse/columns/string.cpp @@ -78,7 +78,12 @@ void ColumnFixedString::Append(ColumnRef column) { } bool ColumnFixedString::LoadBody(InputStream * input, size_t rows) { - data_.resize(string_size_ * rows); + const size_t bytes = string_size_ * rows; + if (string_size_ != 0 && bytes / string_size_ != rows) { + throw ProtocolError("FixedString(" + std::to_string(string_size_) + ") column with " + std::to_string(rows) + + " rows exceeds the addressable size"); + } + data_.resize(bytes); if (!WireFormat::ReadBytes(*input, &data_[0], data_.size())) { return false; } diff --git a/clickhouse/columns/tuple.cpp b/clickhouse/columns/tuple.cpp index cde48f76..ff11cf44 100644 --- a/clickhouse/columns/tuple.cpp +++ b/clickhouse/columns/tuple.cpp @@ -120,6 +120,10 @@ bool ColumnTuple::LoadBody(InputStream* input, size_t rows) { if (!(*ci)->LoadBody(input, rows)) { return false; } + if ((*ci)->Size() != rows) { + throw ProtocolError("Tuple element " + std::to_string(ci - columns_.begin()) + " loaded " + std::to_string((*ci)->Size()) + + " rows, expected " + std::to_string(rows)); + } } return true; diff --git a/ut/columns_ut.cpp b/ut/columns_ut.cpp index 57650ea9..99d9acd0 100644 --- a/ut/columns_ut.cpp +++ b/ut/columns_ut.cpp @@ -339,6 +339,100 @@ TEST(ColumnsCase, FixedStringSwapDifferentSize) { EXPECT_EQ(column2->At(0), "bbbb"); } +// Width for which width * 4 wraps around size_t to 4. +static const size_t WRAPPING_FIXED_STRING_WIDTH = std::numeric_limits::max() / 4 + 2; + +TEST(ColumnsCase, FixedString_Load) { + const auto data = "abcdefgh"sv; + ArrayInput input(data.data(), data.size()); + ColumnFixedString col(4); + + ASSERT_TRUE(col.Load(&input, 2)); + ASSERT_EQ(col.Size(), 2u); + EXPECT_EQ(col.At(0), "abcd"); + EXPECT_EQ(col.At(1), "efgh"); +} + +TEST(ColumnsCase, FixedString_Load_SizeOverflow) { + const auto data = "ABCD"sv; + ArrayInput input(data.data(), data.size()); + ColumnFixedString col(WRAPPING_FIXED_STRING_WIDTH); + + EXPECT_THROW(col.Load(&input, 4), ProtocolError); +} + +TEST(ColumnsCase, Tuple_Load) { + const auto data = + "\x01\x02" + "abcdefgh"sv; + ArrayInput input(data.data(), data.size()); + ColumnTuple col({std::make_shared(), std::make_shared(4)}); + + ASSERT_TRUE(col.Load(&input, 2)); + ASSERT_EQ(col.Size(), 2u); + EXPECT_EQ(col[0]->As()->At(1), 2u); + EXPECT_EQ(col[1]->As()->At(0), "abcd"); + EXPECT_EQ(col[1]->As()->At(1), "efgh"); +} + +TEST(ColumnsCase, Tuple_Load_FixedStringSizeOverflow) { + const auto data = + "\x01\x02\x03\x04" + "ABCD"sv; + ArrayInput input(data.data(), data.size()); + ColumnTuple col({std::make_shared(), std::make_shared(WRAPPING_FIXED_STRING_WIDTH)}); + + EXPECT_THROW(col.Load(&input, 4), ProtocolError); +} + +TEST(ColumnsCase, Tuple_Load_ElementSizeMismatch) { + // Nothing::LoadBody adds to the rows it already holds while UInt8 replaces them, so a second Load leaves the elements uneven. + const auto data = + "\x01\x00" + "\x01\x02\x00\x00"sv; + ArrayInput input(data.data(), data.size()); + auto col = CreateColumnByType("Tuple(UInt8, Nothing)"); + + ASSERT_TRUE(col->Load(&input, 1)); + EXPECT_THROW(col->Load(&input, 2), ProtocolError); +} + +TEST(ColumnsCase, Nullable_Load) { + const auto data = + "\x00\x01" + "abcd\x00\x00\x00\x00"sv; + ArrayInput input(data.data(), data.size()); + ColumnNullable col(std::make_shared(4), std::make_shared()); + + ASSERT_TRUE(col.Load(&input, 2)); + ASSERT_EQ(col.Size(), 2u); + EXPECT_FALSE(col.IsNull(0)); + EXPECT_TRUE(col.IsNull(1)); + EXPECT_EQ(col.Nested()->As()->At(0), "abcd"); +} + +TEST(ColumnsCase, Nullable_Load_FixedStringSizeOverflow) { + const auto data = + "\x00\x00\x00\x00" + "ABCD"sv; + ArrayInput input(data.data(), data.size()); + ColumnNullable col(std::make_shared(WRAPPING_FIXED_STRING_WIDTH), std::make_shared()); + + EXPECT_THROW(col.Load(&input, 4), ProtocolError); +} + +TEST(ColumnsCase, Nullable_Load_NestedSizeMismatch) { + // Nothing::LoadBody adds to the rows it already holds while the null map replaces them, so a second Load leaves them uneven. + const auto data = + "\x00\x00" + "\x00\x00\x00\x00"sv; + ArrayInput input(data.data(), data.size()); + auto col = CreateColumnByType("Nullable(Nothing)"); + + ASSERT_TRUE(col->Load(&input, 1)); + EXPECT_THROW(col->Load(&input, 2), ProtocolError); +} + TEST(ColumnsCase, StringInit) { auto values = MakeStrings(); auto col = std::make_shared(values);