Merge pull request ClickHouse#90501 from ClickHouse/backport/25.8/90031

Backport ClickHouse#90031 to 25.8: Do size checks when deserializing data from aggregation states and other sources

Author: clickhouse-gh[bot]

src/AggregateFunctions/AggregateFunctionGroupArray.cpp

@@ -1,6 +1,7 @@
 #include <AggregateFunctions/AggregateFunctionFactory.h>
 #include <AggregateFunctions/Helpers.h>
 #include <AggregateFunctions/FactoryHelpers.h>
+#include <AggregateFunctions/KeyHolderHelpers.h>
 #include <Interpreters/Context.h>
 #include <Core/ServerSettings.h>
 
@@ -464,7 +465,10 @@ struct GroupArrayNodeGeneral : public GroupArrayNodeBase<GroupArrayNodeGeneral>
         return node;
     }
 
-    void insertInto(IColumn & column) { std::ignore = column.deserializeAndInsertAggregationStateValueFromArena(data()); }
+    void insertInto(IColumn & column)
+    {
+        deserializeAndInsert<false>({data(), size}, column);
+    }
 };
 
 template <typename Node, bool has_sampler>

src/AggregateFunctions/AggregateFunctionGroupArrayIntersect.cpp

@@ -21,6 +21,7 @@
 #include <AggregateFunctions/FactoryHelpers.h>
 #include <AggregateFunctions/Helpers.h>
 #include <AggregateFunctions/IAggregateFunction.h>
+#include <AggregateFunctions/KeyHolderHelpers.h>
 
 #include <memory>
 
@@ -341,10 +342,7 @@ class AggregateFunctionGroupArrayIntersectGeneric final
 
         for (auto & elem : set)
         {
-            if constexpr (is_plain_column)
-                data_to.insertData(elem.getValue().data, elem.getValue().size);
-            else
-                std::ignore = data_to.deserializeAndInsertAggregationStateValueFromArena(elem.getValue().data);
+            deserializeAndInsert<is_plain_column>(elem.getValue(), data_to);
         }
     }
 };

src/AggregateFunctions/Combinators/AggregateFunctionDistinct.h

@@ -138,9 +138,10 @@ struct AggregateFunctionDistinctMultipleGenericData : public AggregateFunctionDi
                 Set::LookupResult it;
                 bool inserted;
                 history.emplace(ArenaKeyHolder{value, *arena}, it, inserted);
-                const char * pos = it->getValue().data;
+                ReadBufferFromString in({it->getValue().data, it->getValue().size});
+                /// Multiple columns are serialized one by one
                 for (auto & column : argument_columns)
-                    pos = column->deserializeAndInsertAggregationStateValueFromArena(pos);
+                    column->deserializeAndInsertAggregationStateValueFromArena(in);
             }
         }
     }

src/AggregateFunctions/KeyHolderHelpers.h

@@ -1,12 +1,18 @@
 #pragma once
 
-#include <Common/HashTable/HashTableKeyHolder.h>
 #include <Columns/IColumn.h>
+#include <Common/HashTable/HashTableKeyHolder.h>
+#include <IO/ReadBufferFromString.h>
 
 namespace DB
 {
 struct Settings;
 
+namespace ErrorCodes
+{
+    extern const int INCORRECT_DATA;
+}
+
 template <bool is_plain_column = false>
 static auto getKeyHolder(const IColumn & column, size_t row_num, Arena & arena)
 {
@@ -29,7 +35,14 @@ static void deserializeAndInsert(StringRef str, IColumn & data_to)
     if constexpr (is_plain_column)
         data_to.insertData(str.data, str.size);
     else
-        std::ignore = data_to.deserializeAndInsertAggregationStateValueFromArena(str.data);
+    {
+        ReadBufferFromString in({str.data, str.size});
+        data_to.deserializeAndInsertAggregationStateValueFromArena(in);
+        if (!in.eof())
+        {
+            throw Exception(ErrorCodes::INCORRECT_DATA, "Extra bytes ({}) found after deserializing aggregation state", in.available());
+        }
+    }
 }
 
 }

src/Columns/ColumnAggregateFunction.cpp

@@ -609,7 +609,7 @@ StringRef ColumnAggregateFunction::serializeValueIntoArena(size_t n, Arena & are
     return out.complete();
 }
 
-const char * ColumnAggregateFunction::deserializeAndInsertFromArena(const char * src_arena)
+void ColumnAggregateFunction::deserializeAndInsertFromArena(ReadBuffer & in)
 {
     ensureOwnership();
 
@@ -618,21 +618,10 @@ const char * ColumnAggregateFunction::deserializeAndInsertFromArena(const char *
       */
     Arena & dst_arena = createOrGetArena();
     pushBackAndCreateState(data, dst_arena, func.get());
-
-    /** We will read from src_arena.
-      * There is no limit for reading - it is assumed, that we can read all that we need after src_arena pointer.
-      * Buf ReadBufferFromMemory requires some bound. We will use arbitrary big enough number, that will not overflow pointer.
-      * NOTE Technically, this is not compatible with C++ standard,
-      *  as we cannot legally compare pointers after last element + 1 of some valid memory region.
-      *  Probably this will not work under UBSan.
-      */
-    ReadBufferFromMemory read_buffer(src_arena, std::numeric_limits<char *>::max() - src_arena - 1);
-    func->deserialize(data.back(), read_buffer, version, &dst_arena);
-
-    return read_buffer.position();
+    func->deserialize(data.back(), in, version, &dst_arena);
 }
 
-const char * ColumnAggregateFunction::skipSerializedInArena(const char *) const
+void ColumnAggregateFunction::skipSerializedInArena(ReadBuffer &) const
 {
     throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Method skipSerializedInArena is not supported for {}", getName());
 }

src/Columns/ColumnAggregateFunction.h

@@ -173,9 +173,9 @@ class ColumnAggregateFunction final : public COWHelper<IColumnHelper<ColumnAggre
 
     StringRef serializeValueIntoArena(size_t n, Arena & arena, char const *& begin) const override;
 
-    const char * deserializeAndInsertFromArena(const char * src_arena) override;
+    void deserializeAndInsertFromArena(ReadBuffer & in) override;
 
-    const char * skipSerializedInArena(const char *) const override;
+    void skipSerializedInArena(ReadBuffer & in) const override;
 
     void updateHashWithValue(size_t n, SipHash & hash) const override;
 

src/Columns/ColumnArray.cpp

@@ -17,6 +17,7 @@
 #include <Common/WeakHash.h>
 #include <Common/HashTable/Hash.h>
 #include <cstring> // memcpy
+#include <IO/ReadHelpers.h>
 
 
 namespace DB
@@ -296,39 +297,35 @@ std::optional<size_t> ColumnArray::getSerializedValueSize(size_t n) const
 }
 
 
-const char * ColumnArray::deserializeAndInsertFromArena(const char * pos)
+void ColumnArray::deserializeAndInsertFromArena(ReadBuffer & in)
 {
-    size_t array_size = unalignedLoad<size_t>(pos);
-    pos += sizeof(array_size);
+    size_t array_size;
+    readBinaryLittleEndian<size_t>(array_size, in);
 
     for (size_t i = 0; i < array_size; ++i)
-        pos = getData().deserializeAndInsertFromArena(pos);
+        getData().deserializeAndInsertFromArena(in);
 
     getOffsets().push_back(getOffsets().back() + array_size);
-    return pos;
 }
 
-const char * ColumnArray::deserializeAndInsertAggregationStateValueFromArena(const char * pos)
+void ColumnArray::deserializeAndInsertAggregationStateValueFromArena(ReadBuffer & in)
 {
-    size_t array_size = unalignedLoad<size_t>(pos);
-    pos += sizeof(array_size);
+    size_t array_size;
+    readBinaryLittleEndian<size_t>(array_size, in);
 
     for (size_t i = 0; i < array_size; ++i)
-        pos = getData().deserializeAndInsertAggregationStateValueFromArena(pos);
+        getData().deserializeAndInsertAggregationStateValueFromArena(in);
 
     getOffsets().push_back(getOffsets().back() + array_size);
-    return pos;
 }
 
-const char * ColumnArray::skipSerializedInArena(const char * pos) const
+void ColumnArray::skipSerializedInArena(ReadBuffer & in) const
 {
-    size_t array_size = unalignedLoad<size_t>(pos);
-    pos += sizeof(array_size);
+    size_t array_size;
+    readBinaryLittleEndian<size_t>(array_size, in);
 
     for (size_t i = 0; i < array_size; ++i)
-        pos = getData().skipSerializedInArena(pos);
-
-    return pos;
+        getData().skipSerializedInArena(in);
 }
 
 void ColumnArray::updateHashWithValue(size_t n, SipHash & hash) const

src/Columns/ColumnArray.h

@@ -82,9 +82,9 @@ class ColumnArray final : public COWHelper<IColumnHelper<ColumnArray>, ColumnArr
     StringRef serializeAggregationStateValueIntoArena(size_t n, Arena & arena, char const *& begin) const override;
     char * serializeValueIntoMemory(size_t, char * memory) const override;
     std::optional<size_t> getSerializedValueSize(size_t n) const override;
-    const char * deserializeAndInsertFromArena(const char * pos) override;
-    const char * deserializeAndInsertAggregationStateValueFromArena(const char * pos) override;
-    const char * skipSerializedInArena(const char * pos) const override;
+    void deserializeAndInsertFromArena(ReadBuffer & in) override;
+    void deserializeAndInsertAggregationStateValueFromArena(ReadBuffer & in) override;
+    void skipSerializedInArena(ReadBuffer & in) const override;
     void updateHashWithValue(size_t n, SipHash & hash) const override;
     WeakHash32 getWeakHash32() const override;
     void updateHashFast(SipHash & hash) const override;

src/Columns/ColumnBLOB.h

@@ -163,8 +163,8 @@ class ColumnBLOB : public COWHelper<IColumnHelper<ColumnBLOB>, ColumnBLOB>
     void popBack(size_t) override { throwInapplicable(); }
     StringRef serializeValueIntoArena(size_t, Arena &, char const *&) const override { throwInapplicable(); }
     char * serializeValueIntoMemory(size_t, char *) const override { throwInapplicable(); }
-    const char * deserializeAndInsertFromArena(const char *) override { throwInapplicable(); }
-    const char * skipSerializedInArena(const char *) const override { throwInapplicable(); }
+    void deserializeAndInsertFromArena(ReadBuffer &) override { throwInapplicable(); }
+    void skipSerializedInArena(ReadBuffer &) const override { throwInapplicable(); }
     void updateHashWithValue(size_t, SipHash &) const override { throwInapplicable(); }
     WeakHash32 getWeakHash32() const override { throwInapplicable(); }
     void updateHashFast(SipHash &) const override { throwInapplicable(); }

src/Columns/ColumnCompressed.h

@@ -99,8 +99,8 @@ class ColumnCompressed : public COWHelper<IColumnHelper<ColumnCompressed>, Colum
     void popBack(size_t) override { throwMustBeDecompressed(); }
     StringRef serializeValueIntoArena(size_t, Arena &, char const *&) const override { throwMustBeDecompressed(); }
     char * serializeValueIntoMemory(size_t, char *) const override { throwMustBeDecompressed(); }
-    const char * deserializeAndInsertFromArena(const char *) override { throwMustBeDecompressed(); }
-    const char * skipSerializedInArena(const char *) const override { throwMustBeDecompressed(); }
+    void deserializeAndInsertFromArena(ReadBuffer &) override { throwMustBeDecompressed(); }
+    void skipSerializedInArena(ReadBuffer &) const override { throwMustBeDecompressed(); }
     void updateHashWithValue(size_t, SipHash &) const override { throwMustBeDecompressed(); }
     WeakHash32 getWeakHash32() const override { throwMustBeDecompressed(); }
     void updateHashFast(SipHash &) const override { throwMustBeDecompressed(); }